# Nested-transport depth consistency

- ID: 1100
- Canonical URL: https://synthcore.org/idea/1100/nested-transport-depth-consistency
- API JSON: https://synthcore.org/api/idea/1100.json
- API Markdown: https://synthcore.org/api/idea/1100.md
- Verification status: unverified
- Source: [arXiv:2607.18121](https://arxiv.org/abs/2607.18121)
- Category: regularization
- Solves: accuracy, stability, speedup
- ML areas: training, regularization, rnn, ssm
- Math tags: stochastic-processes, optimal-transport, combinatorics, dynamical-systems
- Ratings: usefulness 6/10; difficulty 5/10; novelty 7/10

## 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.

## 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.

## Key formulas

- $$x_{t+1}=(1-\alpha)x_t+\alpha T(x_t)$$
- $$\mathbb{P}(\mathcal{P}^{k,j}_{m,n}\text{ path})=\alpha\,(\alpha^{2}\ell)^{k}\,(1-\alpha)^{n+m-1-j-2k}$$
- $$z_{i,j}=\begin{cases}\pi_i^n,&i=j\in\{0,\ldots,m\},\\[2pt]\pi_i^m-\pi_i^n,&i\in\{0,\ldots,m-1\},\ j=n,\\[2pt]\pi_j^n,&i=m,\ j\in\{m+1,\ldots,n-1\},\\[2pt]\pi_m^m-\displaystyle\sum_{j=m}^{n-1}\pi_j^n,&i=m,\ j=n,\\[2pt]0,&\text{otherwise.}\end{cases}$$
- $$\mathcal{L}_{\mathrm{NTDC}}=\sum_{i=0}^{m}\sum_{j=0}^{n}z_{i,j}\,\|h_i-h_j\|_2^2,\qquad \mathcal{L}=\mathcal{L}_{\mathrm{task}}+\lambda\mathcal{L}_{\mathrm{NTDC}}$$

## 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.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.