Improving numerical stability of degenerate eigenvector problems.

[proteneer]

I'm facing a rather annoying problem with derivatives of eigenvectors. I'm doing an eigen-decomposition of the 3x3 inertia tensor for two point clouds in order to align the principal axes. 99.999% of the time, the derivatives are perfectly stable since the eigenvalues are no where close to each other. However, that 0.001% of the time the eigenvalues become close due to symmetries that arise in a thermal bath, and my simulation blows up. Unfortunately we run for 500k+ steps routinely so these small problems will always occur. Is there a way to stabilize this derivative?

In the following example there are two distinct molecules from a public dataset,

The first molecule has eigenvalues of the inertia tensor:

[0.04587656 0.13561707 0.15451878] <-- fine

The second molecule has the following

[0.0276736  0.14507128 0.14508379] <-- blows up

The problem is that although we never backprop adjoints of eigenvalues, the backprop of eigenvector adjoints do involve a check in the eigenvalues. This is seen in the vjp of grad_eigh (taken from autograd):

F = off_diag / (w_repeated.T - w_repeated + np.eye(N))

Not sure if this is related to gimbal locks in classical mechanics, perhaps a quaternion treatment might be better?

Code sample:

from jax.config import config; config.update("jax_enable_x64", True)
import numpy as onp
import jax.numpy as np
import jax
import functools

x_a = onp.array([[ 2.73311633,  0.18177812, -1.2259971 ],
       [ 2.6979834 ,  0.04043467, -1.22612598],
       [ 2.63764162, -0.01410118, -1.35598765],
       [ 2.48079276, -0.0296    , -1.32917345],
       [ 2.47819522, -0.03812585, -1.17863794],
       [ 2.59164471,  0.0148576 , -1.12400423],
       [ 2.61155763,  0.02180472, -0.98558674],
       [ 2.50908202, -0.02726411, -0.9025647 ],
       [ 2.3882526 , -0.08075677, -0.95932773],
       [ 2.37435233, -0.0876205 , -1.10186291],
       [ 2.26638956, -0.14430827, -1.16753249],
       [ 2.14407869, -0.14717527, -1.11510023],
       [ 2.07100457, -0.27024709, -1.12313471],
       [ 1.94221273, -0.28502385, -1.07340893],
       [ 1.87233336, -0.39515079, -1.08716641],
       [ 1.88645201, -0.17213314, -1.01361103],
       [ 1.94467601, -0.04298015, -1.01393563],
       [ 1.88723829,  0.07905923, -0.9735346 ],
       [ 2.07476862, -0.03692507, -1.06932417],
       [ 2.76745785,  0.09095856, -0.91827045],
       [ 2.77289176,  0.07082749, -0.77003796],
       [ 2.89412935,  0.05528072, -0.98350214],
       [ 2.66757146,  0.08183684, -1.45844427],
       [ 1.82809764,  0.17520094, -0.9412489 ],
       [ 2.41275836,  0.08275558, -1.37509023],
       [ 2.70252212, -0.13417082, -1.38254163],
       [ 2.73618653,  0.25141057, -0.94660624],
       [ 2.8127551 ,  0.18381183, -1.28099194],
       [ 2.79083978, -0.01474471, -1.20346063],
       [ 2.44384181, -0.123096  , -1.37695732],
       [ 2.51941863, -0.01680053, -0.79634633],
       [ 2.30691973, -0.11270552, -0.89090591],
       [ 2.1178025 , -0.36358227, -1.15792079],
       [ 1.78259332, -0.19329293, -0.9702879 ],
       [ 2.12432485,  0.05913917, -1.06949944],
       [ 2.6413124 ,  0.25330167, -0.91536833],
       [ 2.74380936,  0.26297849, -1.0502067 ]])
x_b = onp.array([[ 2.57284285,  0.06760627, -1.39784694],
       [ 2.39290326,  0.04465644, -1.22753136],
       [ 2.51534151,  0.10008738, -1.25980957],
       [ 2.58913275,  0.18380994, -1.17319303],
       [ 2.52231642,  0.21967995, -1.05427739],
       [ 2.40285849,  0.15611561, -1.01952239],
       [ 2.34129871,  0.07101286, -1.1048996 ],
       [ 2.22058919,  0.02021698, -1.05675381],
       [ 2.2105963 , -0.11717335, -1.04298674],
       [ 2.31615665, -0.1977131 , -1.00610058],
       [ 2.3062712 , -0.33819435, -0.99548418],
       [ 2.18535067, -0.40277475, -1.03957463],
       [ 2.07685267, -0.32297073, -1.07003918],
       [ 1.95105107, -0.38156664, -1.10349816],
       [ 2.08699896, -0.18075882, -1.06093312],
       [ 2.75117595,  0.26237364, -1.2031492 ],
       [ 2.76015439,  0.3341335 , -1.33501717],
       [ 2.78137784,  0.34612478, -1.08680888],
       [ 2.88066072,  0.11882641, -1.20740952],
       [ 3.00556801,  0.16550629, -1.21802936],
       [ 2.8753517 ,  0.03191879, -1.09883737],
       [ 2.29467925, -0.0507217 , -1.34621636],
       [ 2.67418731, -0.03636819, -1.38731796],
       [ 1.86391525, -0.44821548, -1.13137375],
       [ 2.49833354,  0.03264615, -1.46676707],
       [ 2.62028234,  0.15530322, -1.44201163],
       [ 2.56450861,  0.28768358, -0.97544155],
       [ 2.36098737,  0.1771387 , -0.92419352],
       [ 2.41044556, -0.14261345, -0.98342749],
       [ 2.3841071 , -0.39311771, -0.94943743],
       [ 2.17129662, -0.51186501, -1.03277499],
       [ 2.0001237 , -0.11475051, -1.06882338],
       [ 2.86346044,  0.06202121, -1.29483951],
       [ 2.68249409, -0.09166717, -1.46721486]])

masses_a = onp.array([15.999, 12.011, 12.011, 12.011, 12.011, 12.011, 12.011, 12.011,
       12.011, 12.011, 15.999, 12.011, 12.011, 12.011, 18.998, 12.011,
       12.011, 12.011, 12.011, 32.067, 15.999, 15.999, 18.998, 14.007,
       18.998, 18.998, 14.007,  1.008,  1.008,  1.008,  1.008,  1.008,
        1.008,  1.008,  1.008,  1.008,  1.008])
masses_b = onp.array([12.011, 12.011, 12.011, 12.011, 12.011, 12.011, 12.011, 15.999,
       12.011, 12.011, 12.011, 12.011, 12.011, 12.011, 12.011, 32.067,
       15.999, 15.999, 12.011, 18.998, 18.998, 35.453, 15.999, 14.007,
        1.008,  1.008,  1.008,  1.008,  1.008,  1.008,  1.008,  1.008,
        1.008,  1.008])

def inertia_tensor(conf, masses):
    xs = conf[:, 0]
    ys = conf[:, 1]
    zs = conf[:, 2]
    xx = np.average(ys*ys + zs*zs, weights=masses)
    yy = np.average(xs*xs + zs*zs, weights=masses)
    zz = np.average(xs*xs + ys*ys, weights=masses)
    xy = np.average(-xs*ys, weights=masses)
    xz = np.average(-xs*zs, weights=masses)
    yz = np.average(-ys*zs, weights=masses)
    tensor = np.array([
        [xx, xy, xz],
        [xy, yy, yz],
        [xz, yz, zz]
    ])

    return tensor


def grad_inertia_tensor(conf, masses, tensor_grad):
    xs = conf[:, 0]
    ys = conf[:, 1]
    zs = conf[:, 2]
    xx = np.average(ys*ys + zs*zs, weights=masses)
    yy = np.average(xs*xs + zs*zs, weights=masses)
    zz = np.average(xs*xs + ys*ys, weights=masses)
    xy = np.average(-xs*ys, weights=masses)
    xz = np.average(-xs*zs, weights=masses)
    yz = np.average(-ys*zs, weights=masses)
    tensor = np.array([
        [xx, xy, xz],
        [xy, yy, yz],
        [xz, yz, zz]
    ])

    [[dxx, dxy, dxz], [dxy, dyy, dyz], [dxz, dyz, dzz]] = tensor_grad # 

    N = conf.shape[0]

    mass_sum = np.sum(masses)

    dxs = (dyy*2*xs + dzz*2*xs + -dxy*2*ys + -dxz*2*zs)*(masses/mass_sum)
    dys = (dzz*2*ys + dxx*2*ys + -dxy*2*xs + -dyz*2*zs)*(masses/mass_sum)
    dzs = (dxx*2*zs + dyy*2*zs + -dxz*2*xs + -dyz*2*ys)*(masses/mass_sum)

    dconf = np.stack([dxs, dys, dzs], axis=-1)

    return dconf

# taken from autograd but removed w_g
def grad_eigh(w, v, vg):
    vc = v # real 
    N = 3
    w_repeated = np.repeat(w[..., np.newaxis], N, axis=-1)
    if np.any(vg):
        off_diag = np.ones((N, N)) - np.eye(N)
        # this is the line that blows up the derivatives
        F = off_diag / (w_repeated.T - w_repeated + np.eye(N))
        # (this used to be += but we never do derivatives w.r.t. eigenvalues)
        vjp_temp = np.dot(np.dot(vc, F * np.dot(v.T, vg)), v.T)
    else:
        assert 0
        
    off_diag_mask = (onp.ones((3,3)) - onp.eye(3))/2

    final = vjp_temp*np.eye(vjp_temp.shape[-1]) + (vjp_temp + vjp_temp.T)*off_diag_mask

    return final


def potential(
    a_conf,
    b_conf,
    a_masses,
    b_masses,
    k):

    a_com_conf = a_conf - np.average(a_conf, axis=0, weights=a_masses)
    b_com_conf = b_conf - np.average(b_conf, axis=0, weights=b_masses)

    a_tensor = inertia_tensor(a_com_conf, a_masses)
    b_tensor = inertia_tensor(b_com_conf, b_masses)

    a_eval, a_evec = np.linalg.eigh(a_tensor)
    b_eval, b_evec = np.linalg.eigh(b_tensor)

    loss = []
    for a, b in zip(a_evec.T, b_evec.T):
        delta = 1 - np.abs(np.dot(a, b))
        loss.append(delta*delta*k)

    return np.sum(loss)

def analytical_grad(
    a_conf,
    b_conf,
    a_masses,
    b_masses,
    k):

    a_com_conf = a_conf - np.average(a_conf, axis=0, weights=a_masses)
    b_com_conf = b_conf - np.average(b_conf, axis=0, weights=b_masses)

    a_tensor = inertia_tensor(a_com_conf, a_masses)
    b_tensor = inertia_tensor(b_com_conf, b_masses)

    a_eval, a_evec = np.linalg.eigh(a_tensor)
    b_eval, b_evec = np.linalg.eigh(b_tensor)

    loss = []
    for a, b in zip(a_evec.T, b_evec.T):
        delta = 1 - np.abs(np.dot(a, b))
        loss.append(delta*delta*k)

    dl_daevec_T = []
    dl_dbevec_T = []
    for a, b in zip(a_evec.T, b_evec.T):
        delta = 1 - np.abs(np.dot(a, b))
        prefactor = -np.sign(np.dot(a, b))*2*delta*k
        dl_daevec_T.append(prefactor*b)
        dl_dbevec_T.append(prefactor*a)
        
    dl_daevec = np.transpose(np.array(dl_daevec_T))
    dl_dbevec = np.transpose(np.array(dl_dbevec_T))

    dl_datensor = grad_eigh(a_eval, a_evec, np.array(dl_daevec))
    dl_dbtensor = grad_eigh(b_eval, b_evec, np.array(dl_dbevec))
    
    dl_da_com_conf = grad_inertia_tensor(a_com_conf, a_masses, dl_datensor)
    dl_db_com_conf = grad_inertia_tensor(b_com_conf, b_masses, dl_dbtensor)

    # conservative forces are not affected by center of mass changes.
    # the vjp w.r.t. to the center of mass yields zeros since sum dx=0, dy=0 and dz=0
    return dl_da_com_conf, dl_db_com_conf



potential = functools.partial(potential, a_masses=masses_a, b_masses=masses_b, k=200.0)

grad_fn = jax.grad(potential, argnums=(0,1))
print(potential(x_a, x_b))
du_dx_a, du_dx_b = grad_fn(x_a, x_b) # large forces on x_b
print(du_dx_a)
print(du_dx_b)

# handwritten backprop, F matrix in grad_eigh is the problem
print(analytical_grad(x_a, x_b, masses_a, masses_b, k=200.0))