Fix transition covariance for non-normal drift matrices

[junipertcy]

When teaching myself how thermox simulates SPU trajectories, I noticed that thermox is exact only when the transformed drift matrix $A_y = L^{-1} A L$ (with $D = LL^\intercal$, as built in thermox.preprocess) is normal. This is because thermox uses the eigendecomposition of $A_y$'s symmetric part

$$A_s = \frac{1}{2}(A_y + A_y^\intercal) = V_s \Lambda_s V_s^\intercal$$

and evaluates, with $\lambda_s$ the diagonal of $\Lambda_s$,

$$\tilde{\Sigma}_t = V_s \mathrm{diag}\left(\frac{1 - e^{-2\lambda_s t}}{2\lambda_s}\right) V_s^\intercal = \int_0^t e^{-(A_y + A_y^\intercal)s} ds$$

Whereas in principle, we should evaluate

$$\Sigma_t = \int_0^t e^{-A_y s} e^{-A_y^\intercal s} , ds$$

The two equations are equal if and only if $A_y$ is normal, that is, $A_y$ and $A_y^\intercal$ commute. But normality is a property of the transformed drift, so a symmetric $A$ is fine with $D = \sigma^2 I$ (this is why thermox.linalg.solve/inv are not affected; expnegm of a non-normal matrix is), but not with a general $D$ (unless $A$ and $D$ commute); a non-symmetric $A$ is wrong even with $D = I$.

The mean is not affected: thermox computes it from the eigendecomposition of $A_y$ itself (expm_vp). The problem is the covariance, and three functions inherit it: sample, conditional.covariance and log_prob (and its gradients). On main, the relative error of $\Sigma_t$ is a few percent for the non-normal cases in the test suite and 17% for a symmetric $A$ with $D = \mathrm{diag}(1, 4, 9)$. The numbers and the script are in this comment, and the one before it explains why the old tests did not catch it.

What this PR does

In short: the covariance is exact for any stable drift; the cost is $O(d^3 + T d^2)$ on uniform grids and $O(U d^3 + T d^2)$ on any other grid, with $U$ a constant (53 for sample, a few hundred for log_prob); normal drifts take the old path and give the same results as before.

Below, $\Sigma(t)$ is the $\Sigma_t$ above (the code calls it cov(t)), and $E(t) = e^{-A_y t}$ is the propagator (E).

For a normal $A_y$, a single eigendecomposition serves every gap, and this is what thermox has always exploited. For a non-normal $A_y$ the matrices $\Sigma(t)$ share no fixed eigenbasis, so an exact sampler appears to require its own $O(d^3)$ factorization at every step. The way out rests on two observations. The first is a composition law, $\Sigma(a + b) = \Sigma(b) + E(b) \Sigma(a) E(b)^\intercal$, by which the operators of a few elementary gaps generate the operators of many. The second is smoothness: $\log\det\Sigma(t)$ and $\Sigma(t)^{-1}$ vary smoothly with the gap, so they can be interpolated across a grid instead of factored at each of its steps. The first observation gives the sampler, the second the density.

Both constructions need $\Sigma(t)$ at arbitrary $t$, evaluated stably, and this is the one new primitive. transition_expm_and_cov(A, dt) evaluates Van Loan's block exponential at the fine scale $h = dt / 2^{12}$ and applies the composition law twelve times, returning $E(dt)$ and $\Sigma(dt)$ with no eigendecomposition; the result is exact for any stable $A_y$, including $dt = 0$, and accurate for $|A_y| dt$ up to about $10^5$. conditional.covariance now calls it directly, and preprocess_drift_matrix gains one flag, is_normal. (The eigenbasis formula I posted first turned out to amplify rounding by $\mathrm{cond}(V)^2$, so it is gone.)

On top of this primitive, sample and log_prob choose one of three paths per call (a lax.switch on is_normal and on whether the grid is uniform):

1. Normal $A_y$, any grid: the existing path, untouched.
2. Non-normal $A_y$, uniform grid: the common case in this library, since thermox.linalg and the examples all use one step size. The operator is built once and applied to every step; uniform_dt tolerates floating-point rounding and lets the first gap differ, so the [0, burnin * dt, dt, ...] grids that thermox.linalg builds qualify, and the cost is $O(d^3 + T d^2)$.
3. Non-normal $A_y$, any other grid:
   • sample builds a dyadic ladder. Although a grid of $T$ different gaps appears to need $T$ operators, it does not: each gap is an integer multiple of a lattice unit $\delta$ (a power of two, chosen so that the gaps are reproduced to the rounding of the time stamps), and is therefore a sum of powers of two. One pair $E_j = E(\delta 2^j)$, $L_j L_j^\intercal = \Sigma(\delta 2^j)$ is built per bit, and the steps are swept once per bit, $e \leftarrow E_j e + L_j z_j$ on the steps whose gap has that bit set; by the composition law, every step ends with exactly the covariance of its gap. The mean goes through the eigenbasis as before, and the once-only count is $U = 53$ bits in float64, 24 in float32.
   • log_prob uses Chebyshev interpolation. The density needs $\log\det\Sigma(dt)$ and $\Sigma(dt)^{-1}$, which the ladder does not compose. Here the smoothness enters. The gaps are grouped into bands, each covering a factor of two in $dt$; within a band, $\Sigma$ is factored at 17 Chebyshev points, and every step reads the interpolant as a 17-term sum. Because interpolation can fail quietly, it is guarded: a run-time check requires the last two Chebyshev coefficients of both series to fall below the accuracy of the covariance routine ($10^{-10}$ in float64, $3 \times 10^{-4}$ in float32), and when the check fails (oscillatory drifts, very ill-conditioned eigenvectors, a zero gap) log_prob falls back to one eigh per step at $O(T d^3)$, never wrong, only slower. The once-only count is $U = 17$ per band that contains a gap, at most $17 \times 53$.

Two remarks concern gradients, which users take through log_prob for maximum-likelihood estimation (the use the README advertises and test_MLE exercises). First, cond and switch keep every branch's residuals for the backward pass, so the per-step branch used to be stored even when another path ran; it now sits under jax.checkpoint, and the peak memory of jax.grad(log_prob) at $d = 64$, $T = 10^4$ drops to 1.3 GB on the interpolation path (from ~10 GB across runs) and to 0.5 GB on the uniform path (from 9.9 GB). Second, gradients follow the branch taken: at an exactly normal $A_y$ they are those of the old formula, which sees $A_y$ only through $(A_y + A_y^\intercal)/2$; this set has measure zero and the behavior is the same on main (noted in the docstring of transition_cov_eigh).

Cost

	main	this PR
normal $A_y$, any grid	$O(d^3 + T d^2)$	unchanged
non-normal $A_y$, uniform grid	$O(T d^2)$, inexact	$O(d^3 + T d^2)$
non-normal $A_y$, other grid, sample	$O(T d^2)$, inexact	$O(U d^3 + T d^2)$, $U = 53$ (float64) or 24 (float32)
non-normal $A_y$, other grid, log_prob	$O(T d^2)$, inexact	$O(U d^3 + T d^2)$, $U \le 17 \times 53$; $O(T d^3)$ if the check fails

Tests

Every reference in tests/test_nonnormal.py (37 tests) is computed independently, by Van Loan's block exponential, by the Lyapunov equation, or from the stationary identity $\Sigma_\infty - e^{-At} \Sigma_\infty e^{-A^\intercal t}$, and the three are cross-checked against one another before use. Against them the tests verify the covariance, log_prob, and their gradients with respect to $A$ and to ts, and the law of the ladder's draws, whitened step by step over 20 000 irregular gaps. The safeguards are exercised too: each kind of grid reaches its intended path, the fallback is bitwise equal to the per-step path when the check declines a drift, and the number of operators built stays constant as $T$ grows. tests/test_conditional.py now compares against expm and Van Loan rather than against thermox's own samples. The full suite, 114 tests, passes on Python 3.12 with jax 0.11.1 and on Python 3.10 with jax 0.6.2, and pre-commit is clean.

Unchanged for current users

The interface does not change: no new arguments, and the grid is detected automatically. For normal drifts the outputs are unchanged up to rounding; I compared sample, log_prob and its gradient with main on uniform and irregular grids, and they agree to $10^{-14}$, bitwise for a symmetric $A$ with $D = I$. conditional.covariance, which now reaches every $A$ through the new routine, agrees to $4 \times 10^{-13}$. What users see of the change is one cost clause in the docstrings of sample and log_prob and one sentence in the README.

Limits

• The interpolation declines oscillatory drifts and very ill-conditioned eigenvectors; log_prob is then correct at $O(T d^3)$.
• transition_expm_and_cov returns NaN beyond $|A_y| dt \approx 3 \times 10^5$, not a wrong number.
• The mean on the per-step paths still uses the eigendecomposition of $A_y$ (expm_vp), as on main.

Measurements (Apple M4, CPU, float64, jit, best of 3)

Accuracy. The covariance the ladder composes is exact to rounding: against Van Loan on a 300-step jittered grid, the max relative error is $1.1 \times 10^{-12}$ at $d = 16$ and $2.2 \times 10^{-12}$ at $d = 64$, where the symmetric-part formula of main is off by up to $7.9 \times 10^{-2}$ on the same steps. A 40 000-step ladder trajectory at $d = 8$ reproduces the Lyapunov stationary covariance to $9.9 \times 10^{-3}$, at the Monte Carlo floor of about $10^{-2}$. The interpolation is as accurate as its nodes: the relative error of log_prob is at most $1.1 \times 10^{-13}$ against the independent reference on every family where the check passes (symmetric; $\mathrm{cond}(V) = 10$; stiff, $10^4$; triangular with $\mathrm{cond}(V) = 1.6 \times 10^5$; $M / \sqrt{d} + 1.1 I$ at $d = 64$); the check declines an oscillatory family and $\mathrm{cond}(V) = 1000$, which take the per-step path.

Sampling time. On a grid where every gap differs, the ladder sits an order of magnitude below the per-step path and about thirteen times above the normal-drift floor: at $d = 64$ and $T = 1000$ it takes 0.08 s, against 0.75 s for the per-step path on the same inputs and 0.006 s for upstream's path on a normal drift. At $d = 256$ the ladder takes 0.89 s, and at $d = 512$, 4.9 s.

Density time. The interpolation's advantage grows with $T$, because its once-only part barely moves: at $d = 64$ it takes 0.21 s against 0.85 s for the per-step path on the same inputs at $T = 1000$ (gradients 0.90 s), and 0.42 s against 8.6 s at $T = 10^4$ (gradient 1.5 s); at $d = 256$ and $T = 1000$, 8.0 s against 19.6 s. The band count behind $U$ grows slowly, from 10 at $T = 100$ to 23 at $T = 10^5$ on grids of sorted uniform random times ($U = 170$ and $391$).

Uniform grids. The operator-once path runs at or below upstream's normal-drift path itself: at $T = 10^4$ and $d = 64$, sample takes 0.012 s and log_prob 0.007 s (normal-drift path, 0.020 s and 0.075 s; per-step path, 9.4 s and 9.3 s); at $d = 1024$, 4.4 s and 3.9 s.

Memory. Peak memory of jax.grad(log_prob) at $d = 64$ and $T = 10^4$: 1.3 GB on the interpolation path, 0.5 GB on the uniform path, 7.3 GB on the per-step path.

[junipertcy]

I also looked at why the old tests did not catch this. test_log_prob_numeric in tests/test_log_prob.py is the only upstream test whose reference does not come from thermox itself (it integrates the covariance numerically), and its case is non-normal: A symmetric $A$ with a dense $D$, where the covariance is 3.5% off. It asserts rtol=1e-2, though, and the actual discrepancy on main is below $2 \times 10^{-3}$. The reason is that the samples come from thermox's own model; the expected log-likelihood is stationary at the data-generating distribution, so the first-order error averages out over the transitions and what remains (a KL divergence) is second order in the covariance error. With this PR it is about $10^{-6}$, so rtol=1e-4 would fail on main but pass here.

test_mean_and_cov in tests/test_conditional.py does use a non-symmetric $A$ (covariance 3 to 4% off), but it compares conditional.covariance with samples drawn by thermox from the same formula, so the two sides move together and the test cannot see the error. I changed it to compare against expm and Van Loan references instead. Lastly, examples/matrix_exponentials is asymmetric by intent, but it uses orthogonal matrices, which are normal, so it happened to be exact.

[junipertcy]

The script below reproduces the numbers above on main and on this branch.

Standalone check (jax + thermox only): relative Frobenius error of the transition covariance against references that do not go through thermox

"""Relative Frobenius error of thermox's transition covariance against references
that do not go through thermox. Run it on main and on this branch."""
import jax, jax.numpy as jnp, thermox
jax.config.update("jax_enable_x64", True)

def van_loan(A, D, t):  # Sigma_t = int_0^t e^{-As} D e^{-A^T s} ds (block exponential)
    d = A.shape[0]
    F = jax.scipy.linalg.expm(jnp.block([[-A, D], [jnp.zeros((d, d)), A.T]]) * t)
    return F[:d, d:] @ jax.scipy.linalg.expm(-A.T * t)

def lyapunov(A, D):  # stationary covariance: A S + S A^T = D
    d = A.shape[0]; I = jnp.eye(d)
    K = jnp.kron(I, A) + jnp.kron(A, I)
    return jnp.linalg.solve(K, D.reshape(-1, order="F")).reshape(d, d, order="F")

def relerr(X, ref):
    return float(jnp.linalg.norm(X - ref) / jnp.linalg.norm(ref))

def is_normal(A, D):  # of the transformed drift A_y = L^{-1} A L, D = L L^T
    L = jnp.linalg.cholesky(D); Ay = jnp.linalg.solve(L, A @ L)
    return float(jnp.linalg.norm(Ay @ Ay.T - Ay.T @ Ay) / jnp.linalg.norm(Ay) ** 2) < 1e-10

A_sym = jnp.array([[3.0, 2, 1], [2, 4, 2], [1, 2, 5]])
D_dense = jnp.array([[1.0, 0.3, -0.1], [0.3, 1, 0.2], [-0.1, 0.2, 1]])
cases = [
    ("symmetric A, D = I",                          A_sym, jnp.eye(3)),
    ("rotation-like A = [[1,2],[-2,1]], D = I",     jnp.array([[1.0, 2], [-2, 1]]), jnp.eye(2)),
    ("symmetric A, D = diag(1,4,9)",                A_sym, jnp.diag(jnp.array([1.0, 4, 9]))),
    ("symmetric A, dense D   (test_log_prob)",      A_sym, D_dense),
    ("A = [[3,2.5],[2,4]], D = 2I  (test_conditional)", jnp.array([[3.0, 2.5], [2, 4]]), 2 * jnp.eye(2)),
    ("triangular A, D = I",                         jnp.array([[2.0, 1.5, 0], [0, 3, 1.5], [0, 0, 4]]), jnp.eye(3)),
]
print(f"{'case':50s} {'A_y normal':>10s} {'Sigma_t(0.7)':>13s} {'Sigma_inf':>10s} {'sampled Sigma_inf':>18s}")
for name, A, D in cases:
    d = A.shape[0]; S_inf = lyapunov(A, D)
    e_t = relerr(thermox.conditional.covariance(0.7, A, D), van_loan(A, D, 0.7))
    e_inf = relerr(thermox.conditional.covariance(200.0, A, D), S_inf)
    ts = jnp.arange(0.0, 10000.0, 0.5)
    xs = thermox.sample(jax.random.PRNGKey(0), ts, jnp.zeros(d), A, jnp.zeros(d), D)
    e_mc = relerr(jnp.cov(xs[2000:].T), S_inf)
    print(f"{name:50s} {str(is_normal(A, D)):>10s} {e_t:13.1e} {e_inf:10.1e} {e_mc:18.1e}")

On main:

case                                               A_y normal  Sigma_t(0.7)  Sigma_inf  sampled Sigma_inf
symmetric A, D = I                                       True       2.2e-15    5.3e-16            1.6e-02
rotation-like A = [[1,2],[-2,1]], D = I                  True       1.5e-16    0.0e+00            7.2e-03
symmetric A, D = diag(1,4,9)                            False       1.7e-01    2.6e-01            1.2e-01
symmetric A, dense D   (test_log_prob)                  False       3.5e-02    4.7e-02            2.9e-02
A = [[3,2.5],[2,4]], D = 2I  (test_conditional)         False       3.1e-02    4.0e-02            1.4e-02
triangular A, D = I                                     False       6.3e-02    9.6e-02            3.3e-02

On this branch:

case                                               A_y normal  Sigma_t(0.7)  Sigma_inf  sampled Sigma_inf
symmetric A, D = I                                       True       2.8e-15    1.2e-15            1.6e-02
rotation-like A = [[1,2],[-2,1]], D = I                  True       2.9e-16    2.4e-16            7.2e-03
symmetric A, D = diag(1,4,9)                            False       2.4e-15    4.8e-16            1.2e-02
symmetric A, dense D   (test_log_prob)                  False       2.2e-15    1.2e-15            1.3e-02
A = [[3,2.5],[2,4]], D = 2I  (test_conditional)         False       6.6e-16    2.5e-16            2.6e-03
triangular A, D = I                                     False       7.3e-16    5.7e-16            1.7e-02

The last column is the empirical covariance of a thermox.sample trajectory (18k points), so it carries about 1e-2 of Monte Carlo noise on top of any bias; the first two rows show that noise floor.

[SamDuffield]

Hi!! Thanks for investigating and highlighting this! The $O(Td^3)$ cost is scaring me though, do you think there might be a way to support non-normal matrices with cost $O(d^3 + Td^2)$ ?

[junipertcy]

Hey Sam, thanks for the speedy reply. Yes! When the time grid is uniform, the transition operator can be built once, from the block matrix exponential, without using the eigenvectors, and applied to every step by a matmul. So a non-normal matrix costs $O(d^3 + T d^2)$, or $O(d^3 \log T + T d^2)$ with the associative scan. Whether we can improve it for arbitrary grids is open though. Is sampling on non-uniform grids something you expect to need?

This second commit is larger than the first, so here is the map when you review: sample and log_prob now first check whether the time grid is uniform after its first gap (like the burn-in grids in thermox.linalg). If it is, and the transformed drift is not normal, we build the transition operator once in transition_expm_and_cov and apply it to every step: for sample with a scan, either the sequential one or a tree, and for log_prob with a single matmul over all steps, since the density needs no scan. I left everything else as it was, and the normal path still agrees with main to rounding.

Let me know if you find anything.

[SamDuffield]

Yeah I'd say that supporting non-uniform grids is a requirement

[junipertcy]

Good challenge! I might have a idea (solving a smaller optimization problem at every time step). Let me dig around and update.

[junipertcy]

Update: My idea did not work. I wanted to work in the eigenbasis of $A_y$, but then realized that $\Sigma_t$ will span the the whole $S_d(\mathbb{R})$, so full rank, dim = $d(d+1)/2$. Propagating $\Sigma_t$ across time makes it rotate. That's two dense matrices multiplying with each other. Without assumptions on $A_y$, there is no "smaller optimization problem".

And then I learned that we can use established tricks on exact discretization of linear SDEs (because the transition operators compose along the semigroup) and approximation theory, respectively, to reach the bound that you are asking for. Specifically, we used "dyadic ladder" for sampling and "piecewise Chebyshev interpolation" for the probability.

This gives $O( U d^3 + T d^2 )$ for both. For sampling, $U$ is a true constant. ($U$ is the number of binary digits; float64 has $U=53$ and float32 has $U=24$). For probability density, $U \sim \log T$ because it depends on how spread out the time gaps are. I measured that $U$ can go from 170 at $T = 100$ to about 390 at $T = 10^5$. There's a hard cap of $U = 17 \times 53$ though because we will hit the number of significant binary digits (that's $53$) of a float64, and we always use $17$ nodes to fit each time piece for a similar reason.

This PR has been a wild ride for me! Please ask me questions. Not this weekend, but next week I will have time to iron out the missing details. As of now, the PR is backward compatible, and can take non-normal matrices at uniform or non-uniform grids.

[junipertcy]

@SamDuffield The PR is complete and ready for review. I rewrote the description so it matches the code as it is now (four commits). I have no further changes planned. Anything from here will be in response to your review. I am around this week and will answer within a day.