Source-linked AI summary
Avoiding Premature Collapse: Adaptive Annealing for Entropy-Regularized Structural Inference
Yizhi Liu
TL;DR
The paper studies why annealing entropy-regularized inference toward zero temperature can collapse before recovering the intended discrete structure. It analyzes Sinkhorn sensitivity and non-normal dynamics, then proposes EPH-ASC, which adaptively pauses cooling; the method reaches target accuracy in 47 epochs and addresses late-stage instability in mHC training.
Problem
As ϵ decreases, inference sensitivity grows as O(1/ϵ), making annealing fragile under distributional shifts and high-variance gradients.
Method
EPH-ASC uses the derived linear stability law to monitor distributional drift and adaptively pause cooling when drift exceeds solver capacity.
Results
EPH-ASC preserves entropy during drift spikes, resumes annealing after feature maturation, and reaches target accuracy in 47 epochs—a 1.60× speedup over Gumbel-Sinkhorn.
Takeaways & Limitations
EPH-ASC provides a lightweight adaptive schedule for preventing premature collapse and stabilizing Manifold-Constrained Hyper-Connections during large-scale language-model training.
Takeaways & Limitations
The theoretical sensitivity analysis relies on a localized non-degeneracy assumption fixing an active support set for small ϵ.
Abstract
from arXiv · showhide
Differentiable matching layers and residual connection paradigms, often implemented via entropy-regularized Optimal Transport (OT), serve as critical mechanisms in structural prediction and architectural scaling. However, recovering discrete permutations or maintaining identity mappings via annealing $ε\to 0$ is notoriously unstable. In this work, we identify a fundamental mechanism for this failure: \textbf{Premature Mode Collapse}. By analyzing the non-normal dynamics of the Sinkhorn fixed-point map, we reveal a theoretical thermodynamic speed limit: standard exponential cooling outpaces the contraction rate of the inference operator, which degrades as $O(1/ε)$. To address this, we propose \textbf{Efficient Piecewise Hybrid Adaptive Stability Control (EPH-ASC)}, an adaptive scheduling algorithm that monitors the stability of the inference process. We demonstrate that EPH-ASC is essential for stabilizing Manifold-Constrained Hyper-Connections (mHC) during large-scale training on the FineWeb-Edu dataset, effectively preventing late-stage gradient explosions by enforcing a linear stability law.
1 Introduction
Entropy-regularized OT supports hard assignment recovery and identity-preserving architectural mappings, but cooling toward zero regularization is fragile. EPH-ASC addresses this by monitoring distributional drift and pausing cooling when stability is violated.
- Entropy-regularized OT is used for combinatorial inference and residual mappings projected onto the Birkhoff polytope.
- As ϵ decreases, sensitivity of the optimal plan to cost perturbations grows as O(1/ϵ).
- Standard annealing breaches the stability threshold and locks inference into a spurious mode.
- EPH-ASC monitors primal drift and triggers a Thermodynamic Pause when distributional shift exceeds the solver’s restoring capacity.
2 The Mechanism of Inference Collapse
As entropy temperature approaches zero, soft beliefs sharpen toward permutation vertices while the inference basin contracts. Non-normal Jacobian dynamics and 1/ϵ sensitivity make early locking possible when posterior drift outpaces correction.
- As ε →0, the entropic map sharpens soft belief states into near-permutations organized around permutation-vertex attraction basins.
- Under localized non-degeneracy, Sinkhorn-map sensitivity to ε scales as 1/ε.
- The effective stability basin shrinks proportionally with ε, requiring permissible distributional drift τt ∝ε.
- Non-normal Jacobian structure can amplify transient inference errors beyond what eigenvalue contraction alone predicts.
- For sufficiently small ϵ, Proposition 2.1 gives a spectral-gap lower bound 1 −ρ(Jϵ) ≥γ · ϵ.
- The permissible drift τmax must scale linearly with ϵ as ϵ →0.
3 Theoretical Analysis: The Thermodynamic Speed Limit
The analysis models annealing as tracking a moving Sinkhorn fixed point, where distributional drift competes with contraction. It derives a speed limit showing that exponential cooling becomes incompatible with bounded tracking error as temperature vanishes.
- 3.1 Problem Setup and Sinkhorn Dynamics: The annealing process is modeled as discrete-time tracking of a moving fixed point, with error et measuring deviation from that point.
- 3.1 Problem Setup and Sinkhorn Dynamics: The analysis relies on a localized non-degeneracy assumption for the cost matrix and active support.
- 3.1 Problem Setup and Sinkhorn Dynamics: The fixed-point drift scales as ∥∇ϵP ∗ ϵ ∥= Θ(ϵ−1), while the Jacobian spectral gap vanishes as 1 −ρ(Jϵ) = Θ(ϵ).
- 3.1 Problem Setup and Sinkhorn Dynamics: Non-normal pseudospectrum contours extend beyond the unit circle, quantifying transient error amplification.
- 3.2 The Adiabatic Tracking Theorem: The thermodynamic speed limit requires annealing step size δt to be at least quadratic, δt ∝ϵ2, for bounded tracking error.
- 3.2 The Adiabatic Tracking Theorem: Exponential annealing produces δt ∝ϵt and violates the speed limit by a factor of 1/ϵt, making collapse theoretically inevitable as ϵt →0.
4 Method: Efficient Piecewise Hybrid ASC
EPH-ASC combines sensitivity-based stability control with an efficient two-phase protocol. It estimates a dataset-specific safety slope, monitors drift during training, and pauses cooling when drift exceeds basin capacity.
- EPH-ASC decouples expensive spectral diagnostics from the training loop while targeting topological stability and computational efficiency.
- 4.1 Approximating the Stability Constraint: The method approximates stability using a permissible drift threshold that follows a linear law in ϵ.
- 4.2 Two-Phase Protocol: Offline calibration estimates ksafe from the drift-to-temperature ratio at intentionally triggered topological collapse.
- 4.2 Two-Phase Protocol: At runtime, the controller monitors instantaneous shift ∥∆t∥F and enforces the stability constraint.
- 4.2 Two-Phase Protocol: When ∥∆t∥F > ksafe · ϵt, EPH-ASC holds ϵt+1 ←ϵt constant in a Thermodynamic Pause.
- 4.2 Two-Phase Protocol: The controller’s stability-braking mechanism pauses cooling when primal drift approaches or exceeds the stability threshold.
5 Experiments
Experiments on SPair-71k and FineWeb-Edu show that standard annealing can collapse, while EPH-ASC stabilizes training through drift-triggered braking and entropy preservation.
- SPair-71k: Standard annealing collapses around Epoch ≈20 on SPair-71k, causing vanishing gradients and flatlined accuracy.Aggressive sharpening forces the plan into a spurious basin.
- SPair-71k: EPH-ASC detects drift spikes and triggers Stability Braking, holding ϵ constant to preserve uncertainty before annealing resumes.This combines deterministic gradients with adaptive control.
- SPair-71k: 47 epochs and a 1.60× speedup over Gumbel-Sinkhorn were achieved after features matured and annealing resumed.Standard annealing failed, while EPH-ASC incurred 0.51% overhead.
- FineWeb-Edu: On FineWeb-Edu, the Naive schedule appeared successful for 98% of training before a catastrophic gradient explosion at Step 980.The experiment used NanoGemma with mHC over 1,000 steps.
- FineWeb-Edu: EPH-ASC detected critical distributional drift at Step 640, then locked ϵ ≈0.04 and secured a 340-step safety margin.Its entropy preservation prevented numerical underflow observed in the baseline.
6 Conclusion
The conclusion identifies Premature Mode Collapse as a thermodynamic failure caused by distributional shift exceeding inference contraction, and presents EPH-ASC as an adaptive response based on linear stability scaling.
- Conclusion: Premature Mode Collapse occurs when distributional shift exceeds the inference operator’s contraction rate.The paper describes this as a fundamental thermodynamic failure.
- Conclusion: EPH-ASC uses a lightweight adaptive schedule based on the derived O(ϵ) stability law.The method is intended to reconcile instability with practical scheduling.
- Scope: The analysis is conditional on localized non-degeneracy and support stability for sufficiently small ϵ.The active support is assumed locally stable rather than globally invariant.
- Scope: The result is a lower bound on failure: exponential annealing fails even in the best-behaved regime captured by the assumption.Support bifurcations or active-set changes lie outside this best-case regime and may aggravate instability.
A.2 Transient Stability Analysis: The Resolvent View
The resolvent analysis connects Sinkhorn sensitivity to the Jacobian’s spectral properties, showing that shrinking stability margins and curvature make collapse inevitable under exponential annealing.
- Resolvent relation: Theorem A.2 links optimal-plan sensitivity to cost perturbations with the spectral distance of the Sinkhorn Jacobian from unit gain.The bound does not require the Jacobian to be normal or diagonalizable.
- Resolvent relation: Sensitivity scaling as Ω(1/ϵ) forces the spectral gap to vanish at most linearly, with 1 −ρ(Jϵ) ≤O(ϵ).The restoring capacity is governed by a gap scaling as O(ϵ).
- Basin mismatch: Under exponential annealing, the tracking error diverges because the schedule’s step size remains incompatible with the shrinking restoring capacity.Corollary A.3 states this divergence directly.
- Basin mismatch: The linearization region shrinks as O(ϵ), while the tracking error scales as O(1/ϵ), so collapse occurs when the error exceeds that region.The mismatch becomes inevitable as ϵ approaches zero.
- Active reduced system: The active-support reduction restricts the analysis to supported variables, where the reduced operator is uniformly invertible under the stated assumption.Entries outside the active support are O(τ(ϵ)) and do not affect active stability.
A.5 Directional sensitivity lower bound
The analysis constructs a directional lower bound showing that Sinkhorn fixed-point sensitivity grows at least as O(1/ε) under active-support regularity. The result follows by factoring the implicit sensitivity through the reduced active-support linear system and a nonzero cost-action operator.
- Constructive lower bound: A nonzero composed cost-action operator M(ε) yields a constructive operator-norm lower bound for the sensitivity.A singular-vector perturbation can attain or approximate the lower bound, while a nondegeneracy assumption ensures M(ε) is not identically zero.
- Uniform scaling: Under σmin(A_S(ε)) ≥ a_min > 0 and bounded response operators, ||D_S^ε(C)||op ≥ c0/ε for sufficiently small ε.The constant c0 is independent of ε under the stated uniform spectral and operator-norm conditions.
- Constructive lower bound: The fixed-point sensitivity operator is factored through the active-support system as an inverse linearization composed with cost and primal-response maps.The reduced operators A_S(ε), B_S, and R_S capture the implicit differentiation pipeline on the gauge-fixed active support.
- Spectral implication: The resolvent obeys ||(I−J)^−1|| ≥ 1/dist(1,σ(J)) without requiring diagonalizability or normality.Thus sensitivity divergence forces the spectrum of J toward 1, although stronger eigenvalue asymptotics require additional conditioning assumptions.
A.6 Refined Sensitivity Analysis: The O(ϵ−1) Scaling
The refined analysis explains why entropic optimal-transport sensitivity scales as O(1/ε), rather than the O(1/ε^2) suggested by the Gibbs kernel alone. Adaptive dual potentials and marginal constraints cancel the higher-order singularity.
- Assumptions: The analysis conditions its scaling argument on non-degeneracy and a stable active support for sufficiently small ε.The active-support formulation avoids ill-conditioning from rows and columns whose marginals vanish in the zero-temperature limit.
- Asymptotic scaling: The derivative of the optimal transport plan with respect to ε is dominated by a 1/ε factor on the active support.The derivation combines the primal-dual representation with asymptotic behavior of the plan and dual potentials.
- Asymptotic scaling: Marginal constraints and adaptive dual potentials reduce the apparent Gibbs-kernel divergence from O(ε−2) to O(ε−1).The dual potentials shift to counteract mass displacement caused by changes in the kernel.
- Numerical validation: The paper recommends numerical evaluation of plan norms, reduced-system spectra, forcing responses, and separation conditions for instance-specific constants.For small ε, it also recommends log-domain Sinkhorn computation and implicit differentiation rather than naive finite differences.
B.3 Proof of Theorem 3.2
The proof applies tracking-error dynamics and a contraction error bound to relate annealing drift and resolvent growth. It concludes that the annealing step must scale quadratically with ε to maintain bounded tracking error.
- Error dynamics: The tracking analysis applies the equilibrium-error bound to Sinkhorn dynamics with target drift between successive annealing steps.The drift is quantified using a first-order expansion and the previously established sensitivity result.
- Resolvent growth: The resolvent grows at least as Θ(ε−1), while non-normal pseudospectral effects may make it grow faster.The Θ(ε−1) behavior is used as a conservative lower bound for the necessary stability condition.
- Stability condition: To keep the steady-state error within a basin radius R, the combined drift and resolvent bound imposes a quadratic ε scaling on the annealing step.The basin condition assumes that R shrinks slowly or remains O(1).
- Stability condition: The resulting condition is that the step size δ_t must scale quadratically with ε to maintain tracking.This is the stated conclusion of the proof for preventing escape from the local stability basin.
B.4 Proof of Corollary 3.3
The proof analyzes standard exponential annealing and shows that its stability ratio diverges as temperature approaches zero. Consequently, below a critical temperature, tracking error exceeds basin capacity and premature collapse becomes inevitable.
- Exponential annealing uses ϵt+1 = αϵt with fixed decay rate 0 < α < 1.
- The resulting stability ratio diverges to infinity as ϵt →0.
- Below a critical temperature ϵcrit, tracking error inevitably exceeds basin radius R, triggering premature collapse.
- The appendix aligns this nontrivial claim with an explicit lemma, theorem, or derivation.