Reversible Low-Rank Neural ODE State

(1) Exact integration point: use a matrix-valued Neural ODE or recurrent residual block whose hidden state is F in R^(n by m), for example sequence positions by channels. Store the terminal factors X, S, V and occasional checkpoints instead of storing F at every time step. Define a vector field G_theta(t,F) = MLP_theta(F) + A F + F B, or replace this with a token-mixing and channel-mixing neural field. The low-rank state is always reconstructed as F = X S V^T before evaluating the vector field.

(2) Pseudocode:

factorize F0 into X, S, V using QR/SVD
for k = 0,...,N-1:
    K = X @ S
    K = implicit_midpoint(K, dK = G(t, K @ V.T) @ V, step = h)
    X, R = qr(K)
    S = R
    S = implicit_midpoint(S, dS = projected_core_rhs(X, S, V), step = h)
    L = V @ S.T
    L = implicit_midpoint(L, dL = G(t, X @ L.T).T @ X, step = h)
    V, R = qr(L)
    S = R.T
    save X, S, V every C steps
loss = task_loss(X @ S @ V.T)
reverse K, S, L substeps in reverse order with step = -h
accumulate parameter and adjoint gradients

The K, S, and L updates are the projector-splitting factor substeps. In practice, implicit midpoint for a nonlinear neural field can be solved with a few fixed-point or Newton iterations; a cheap explicit symmetric alternative should be tested separately, but its reversal defect must be measured.

(3) Computed from the paper's mathematics: orthogonal factorization, tangent-space projection, reversible substep ordering, and adjoint accumulation. Estimated empirically: the smallest useful rank r, checkpoint interval C, nonlinear solver tolerance, and whether the chosen neural vector field preserves low-rank structure. Measure the normalized reversal defect delta_k = ||F_k - Reverse(Forward(F_k))||_F / (||F_k||_F + epsilon), and compare reconstructed-state gradients with dense autograd.

(4) First experiment: train a rank-constrained Neural ODE or 12-layer recurrent residual model on Sequential MNIST, comparing dense activation storage, ordinary adjoint recomputation, and this reversible low-rank method at ranks 8, 16, and 32. Measure peak GPU memory, wall-clock time, final accuracy, gradient cosine similarity against dense autograd, and loss at equal memory. A successful result is at least 2x lower activation memory, less than 1 percentage point accuracy loss, and gradient cosine similarity above 0.99 at moderate rank.