Distribution-Aware Contraction Scheduler
Implementation & benchmark of arXiv:2607.14291 — Wasserstein Stability of Contracting Flows: Effective Rates, Euler Self-Correction, and Noise Tightening
Source paper: Wasserstein Stability of Contracting Flows: Effective Rates, Euler Self-Correction, and Noise Tightening arXiv:2607.14291 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Estimate the local contraction rate along minibatch couplings of neural ODE or flow-matching trajectories instead of using one global Lipschitz lower bound. Use the resulting displacement-weighted rate to trigger adaptive solver tolerances, training-time regularization, or early stopping when the transported distributions have entered a strongly contracting region.
Formulas
Mathematical statement
For an ODE \(\dot{x}=v(x)\), let \(X_t\sim\mu_t\) and \(Y_t\sim\nu_t\) be coupled trajectories, with displacement \(Z_t=X_t-Y_t\). Define the pairwise local contraction rate \(r(x,y)=-\langle x-y,v(x)-v(y)\rangle/\|x-y\|^2\) when \(x\neq y\). A globally contracting field satisfies \(r(x,y)\geq\lambda>0\), yielding \(W_2(\mu_t,\nu_t)\leq e^{-\lambda t}W_2(\mu_0,\nu_0)\). The paper's transferable refinement replaces \(\lambda\) by the displacement-weighted average \(\bar r_t=\mathbb{E}[\|Z_t\|^2r(X_t,Y_t)]/\mathbb{E}[\|Z_t\|^2]\), producing the empirical decay law \(d\log W_2/dt\approx-\bar r_t\). For the scalar drift \(v(x)=-\lambda x-\beta x^3\), \(r(x,y)=\lambda+\beta(x^2+xy+y^2)\), so broad distributions should contract faster than the worst-case rate predicts.
Implementation notes
Integrate this into a flow-matching or diffusion-ODE trainer as a diagnostic and adaptive schedule. At each training step, draw paired initial states from two nearby minibatches, or use the same latent seeds with two perturbations, and integrate both through the current vector field for a short time window. Compute \(d_i=x_i-y_i\), \(r_i=-d_i^\top(v(x_i)-v(y_i))/(\|d_i\|^2+\epsilon)\), and the weighted estimate \(\hat r=\sum_i\|d_i\|^2r_i/(\sum_i\|d_i\|^2)\). Pseudocode: sample pairs; evaluate the vector field at both states; compute \(\hat r\); maintain an exponential moving average; if \(\hat r\) is substantially above the baseline global rate, increase the ODE solver step size or reduce contraction regularization, while if \(\hat r\) approaches zero or becomes negative, reduce the step size and add a Jacobian-symmetric-part penalty \(\max(0,\lambda-r_i)^2\). The paper supplies the coupling argument and weighted-rate formula; the neural implementation estimates expectations with minibatches and must validate whether the sampled coupling is representative. First experiment: train a 2D flow-matching MLP on a two-Gaussian-to-Gaussian transport task, comparing fixed-step Euler, global spectral contraction regularization, and this adaptive scheduler. Measure empirical \(W_2\) or Sinkhorn distance between perturbed trajectory ensembles and fit its log-slope. The quantitative prediction is that the measured decay slope tracks \(-\hat r_t\), and that on a cubic-like learned field with particles far from equilibrium, \(\hat r_t\) exceeds the global minimum rate by at least 20%; using the global rate should therefore overestimate the time needed for distributional convergence.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: No measurable effect
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a readable scalar Distribution-Aware Contraction Scheduler diagnostic for v(x)=-lambda*x-beta*x^3, including pairwise rates, displacement-weighted estimates, adaptive stepping, and a global-rate baseline. The algebra identity matched with maximum absolute error 4.12e-5; beta=0 gave r_bar=lambda across all scales, scale increased r_bar from 0.708 to 13.81 at beta=1, and the excess rate scaled linearly with beta (0.832, 1.664, 3.327, 6.654). Early log-displacement slopes followed -r_bar qualitatively, with agreement best at small scales; adaptive stepping used roughly half as many steps as fixed-step Euler. This is a real toy-mechanism signal, not evidence of a neural training win.
- Agent confidence
- 8/10
- Baseline
- Fixed Euler dt=0.01: 59, 150, and 150 steps for scales 0.25, 1.0, and 2.0 respectively; convergence times 0.59, 1.5, and 1.5 (horizon-limited).
- Idea
- Rate-adaptive Euler: 30, 75, and 77 steps; times 0.60, 1.5, and 1.5. Step reductions were 49.2%, 50.0%, and 48.7%. For beta=1, r_bar at scales 0.25/0.5/1/2 was 0.708/1.332/3.827/13.808 versus global lambda=0.5.
Limitations: Only a scalar analytic cubic ODE was tested; no 2D neural flow-matching model, Sinkhorn/W2 estimator, GPU execution, FLOP-normalized training comparison, or solver-accuracy study was performed. The slope test used finite early windows and Euler integration, so agreement degrades as the distribution contracts and the weighted rate changes rapidly.
How to run: python3 contraction_scheduler.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.14291)
No significant difference vs. the tuned baseline (+0.0%, p=1, wins 3 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0014
- Idea mean
- 0.0014
- Effect (Δ)
- 0 (+0.0%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 1 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 6 configs
Only the built-in dynamics pendulum track was tested; no flow-matching or neural ODE solver-tolerance experiment was run. The local implementation used a contraction penalty rather than adaptive ODE step-size control, and the mechanism signature measured a one-step learned-map displacement ratio rather than a continuous-time trajectory slope.
How to run: python3 stage2_bench.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Implemented the Distribution-Aware Contraction Scheduler as a training-time contraction penalty on the matched dynamics track using the shared rnn_small architecture, with 8 paired seeds and equal learning-rate search spaces. The idea was essentially tied with the tuned baseline but slightly worse (delta_mean=+3.84e-7, permutation p=1.0), so there is no neural-training win. The trained-model mechanism signature was not confirmed: predicted local log-rate averaged -0.0484 versus observed output/input displacement log-ratio -2.6054.
- Baseline
- Best cfg lr=0.01, weight_decay=0.0001; full 8-seed MSE mean=0.00141069785, std=0.00038184105.
- Delta mean
- 3.840395947918296e-07
- Idea
- Best cfg lr=0.01, strength=0.08, target=0.02; full 8-seed MSE mean=0.00141108189, std=0.00043187111.
- Idea wins
- 3
- Mechanism confirmed
- false
- Mechanism observed mean
- -2.605360746383667
- Mechanism predicted mean
- -0.04839520575478673
- N pairs
- 8
- P value
- 1
Benchmark result
No significant difference vs. the tuned baseline (+0.0%, p=1, wins 3 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0014
- Idea mean
- 0.0014
- Effect (Δ)
- 0 (+0.0%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 1 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 6 configs
Only the built-in dynamics pendulum track was tested; no flow-matching or neural ODE solver-tolerance experiment was run. The local implementation used a contraction penalty rather than adaptive ODE step-size control, and the mechanism signature measured a one-step learned-map displacement ratio rather than a continuous-time trajectory slope.
How to run: python3 stage2_bench.py
Artifacts
- bench_report.json 5.3 KB View Raw
- contraction_scheduler.py 4.2 KB View Raw
- report.md 1.7 KB View
- report_bench_2026-08-30T182758.md 4.6 KB View
- results.json 2.9 KB View Raw
- stage2_bench.py 5.7 KB View Raw