Nested-transport depth consistency
Implementation & benchmark of arXiv:2607.18121 — Krasnosel'skii-Mann iterations beyond asymptotics: a combinatorial analysis
Source paper: Krasnosel'skii-Mann iterations beyond asymptotics: a combinatorial analysis arXiv:2607.18121 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Apply the paper's nested coupling between path distributions at two Krasnosel'skii–Mann depths to an iterative neural block. Penalize discrepancies between intermediate representations using the coupling mass, so that the short unroll learns to approximate the long unroll while preserving the block's actual computational-path geometry. At inference, use the resulting coupled discrepancy as an early-exit criterion.
Formulas
Mathematical statement
The paper considers an under-relaxed fixed-point iteration x_{t+1}=(1-alpha)x_t+alpha T(x_t), where alpha is in (0,1] and T is an ell-Lipschitz map. Its path stratification assigns probability alpha (alpha^2 ell)^k (1-alpha)^(n+m-1-j-2k) to every path in the class P^{k,j}_{m,n}, consisting of paths from lattice position (m,n) to winning position (-1,j) with k unsettled phases. For two depths m<n, let pi_i^m and pi_j^n be the nonnegative path-mass distributions over i=0,...,m and j=0,...,n, each summing to one. Equation (18) gives a nested optimal transport plan z for alpha >= 1/2: diagonal mass z_{i,i}=pi_i^n for i=j<=m, excess mass z_{i,n}=pi_i^m-pi_i^n for i<m, mass z_{m,j}=pi_j^n for m<j<n, and the remaining corner mass z_{m,n}=pi_m^m-sum_{j=m}^{n-1} pi_j^n. Equation (19) gives a simpler feasible coupling tilde z that keeps shared indices aligned and distributes remaining mass proportionally. The paper proves d_{m,n}<=c_{m,n}, where d_{m,n} is the optimal transport cost and c_{m,n} is the cost under tilde z. We use z or tilde z as weights in a representation-consistency loss.
Implementation notes
1. Integration point: replace a standard recurrent or equilibrium block by an under-relaxed unroll. Given input u, initialize h_0, compute h_{t+1}=(1-alpha)h_t+alpha F_theta(h_t,u), and retain hidden states h_0,...,h_n. During training, choose a short depth m and long depth n>m; the task head uses h_n, while the new regularizer couples all short and long states. Use alpha >= 0.5. 2. Pseudocode: compute h[0]=init(u); for t=0,...,n-1, set h[t+1]=(1-alpha)h[t]+alpha F_theta(h[t],u). Construct path masses pi^t either from the paper's lattice-path dynamic program using the displayed per-path weight alpha(alpha^2 ell)^k(1-alpha)^(n+m-1-j-2k), or initially use normalized empirical masses obtained by counting sampled paths. Build z from equation (18), then compute L_NTDC=sum_{i,j} z[i,j] ||stopgrad(h[i])-h[j]||^2. Backpropagate L_task+lambda L_NTDC. 3. Mathematics versus estimation: equation (18) and the transport loss are exact once pi^m and pi^n are known. Estimate ell cheaply with one-step Jacobian-vector power iteration, ell approximately ||J_F(h)v||_2/||v||_2, and renormalize the resulting path masses if the estimated contraction is above one. For a first MVP, bypass exact pi derivation and compare the paper's nested z against uniform all-pairs weights and diagonal-only matching. 4. First experiment: train a 4-layer iterative MLP or small GRU on sequential MNIST and CIFAR-10 classification, with m=4,n=8, against the same model without the regularizer and standard deep supervision. Measure long-depth accuracy, short-depth accuracy, hidden-state variance, and the number of iterations required for the prediction to agree with the long-depth prediction. Success is higher accuracy at depth m, reduced variance across stopping depths, or at least 25% fewer inference iterations at unchanged accuracy.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.