Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Avoiding Premature Collapse: Adaptive Annealing for Entropy-Regularized Structural Inference

Yizhi Liu

TL;DR

The paper addresses instability when annealing entropy-regularized OT towards hard permutations in differentiable matching. It reveals a thermodynamic speed limit that forces a fast drift of the moving fixed point beyond the solver's contraction, causing Premature Mode Collapse under standard exponential schedules. To fix this, it introduces Efficient Piecewise Hybrid ASC (EPH-ASC), which enforces a linear stability bound and uses a two-phase adaptive annealing protocol to pause cooling when needed, reducing spectral diagnostics overhead. Empirically, EPH-ASC preserves uncertainty, avoids premature locking, and yields a substantial speedup (1.60x) over Gumbel-Sinkhorn with minimal overhead on SPair-71k."

Abstract

Differentiable matching layers, often implemented via entropy-regularized Optimal Transport, serve as a critical approximate inference mechanism in structural prediction. However, recovering discrete permutations via annealing $ε\to 0$ is notoriously unstable. 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 \textbf{thermodynamic speed limit}. Under standard exponential cooling, the shift in the target posterior ($O(1)$) outpaces the contraction rate of the inference operator, which degrades as $O(1/ε)$. This mismatch inevitably forces the inference trajectory into spurious local basins. To address this, we propose \textbf{Efficient PH-ASC}, an adaptive scheduling algorithm that monitors the stability of the inference process. By enforcing a linear stability law, we decouple expensive spectral diagnostics from the training loop, reducing overhead from $O(N^3)$ to amortized $O(1)$. Our implementation and interactive demo are available at https://github.com/xxx0438/torch-sinkhorn-asc and https://huggingface.co/spaces/leon0923/torch-sinkhorn-asc-demo. bounded away from zero in generic training dynamics unless the feature extractor converges unrealistically fast.

Avoiding Premature Collapse: Adaptive Annealing for Entropy-Regularized Structural Inference

TL;DR

The paper addresses instability when annealing entropy-regularized OT towards hard permutations in differentiable matching. It reveals a thermodynamic speed limit that forces a fast drift of the moving fixed point beyond the solver's contraction, causing Premature Mode Collapse under standard exponential schedules. To fix this, it introduces Efficient Piecewise Hybrid ASC (EPH-ASC), which enforces a linear stability bound and uses a two-phase adaptive annealing protocol to pause cooling when needed, reducing spectral diagnostics overhead. Empirically, EPH-ASC preserves uncertainty, avoids premature locking, and yields a substantial speedup (1.60x) over Gumbel-Sinkhorn with minimal overhead on SPair-71k."

Abstract

Differentiable matching layers, often implemented via entropy-regularized Optimal Transport, serve as a critical approximate inference mechanism in structural prediction. However, recovering discrete permutations via annealing is notoriously unstable. 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 \textbf{thermodynamic speed limit}. Under standard exponential cooling, the shift in the target posterior () outpaces the contraction rate of the inference operator, which degrades as . This mismatch inevitably forces the inference trajectory into spurious local basins. To address this, we propose \textbf{Efficient PH-ASC}, an adaptive scheduling algorithm that monitors the stability of the inference process. By enforcing a linear stability law, we decouple expensive spectral diagnostics from the training loop, reducing overhead from to amortized . Our implementation and interactive demo are available at https://github.com/xxx0438/torch-sinkhorn-asc and https://huggingface.co/spaces/leon0923/torch-sinkhorn-asc-demo. bounded away from zero in generic training dynamics unless the feature extractor converges unrealistically fast.
Paper Structure (44 sections, 12 theorems, 69 equations, 4 figures, 2 tables)

This paper contains 44 sections, 12 theorems, 69 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. The Mechanism of Inference Collapse
  3. Geometric Picture: Basin Shrinkage
  4. Thermodynamic Sensitivity (1/epsilon)
  5. Transient Inference Error and Pseudospectra
  6. Theoretical Analysis: The Thermodynamic Speed Limit
  7. Problem Setup and Sinkhorn Dynamics
  8. The Adiabatic Tracking Theorem
  9. Method: Efficient Piecewise Hybrid ASC
  10. Approximating the Stability Constraint
  11. Two-Phase Protocol
  12. Phase I: Calibration (Offline).
  13. Phase II: Runtime Control (Adaptive Annealing).
  14. Experiments
  15. Setup and Baselines
  16. ...and 29 more sections

Key Result

Proposition 2.1

Let the cost matrix $C$ satisfy the localized non-degeneracy Assumption (stable active support $S$). For sufficiently small $\epsilon$, the spectral gap of the Sinkhorn Jacobian $J_\epsilon$ satisfies $1 - \rho(J_\epsilon) \ge \gamma \cdot \epsilon$. Consequently, to ensure the inference error $e_t$

Figures (4)

  • Figure 1: Premature Mode Collapse. Standard annealing (blue) breaches the stability threshold $R$ (dotted), causing early locking into a spurious mode. Ours (red) detects the stability violation and pauses cooling. (Simulation)
  • Figure 2: The Dual View of Inference Collapse. (a) Geometric intuition: The inference fails when the target drifts faster than the shrinking basin allows. (b) Spectral reality: This shrinkage is quantified by the non-normal pseudospectrum.
  • Figure 3: Mechanism of Adaptive Stability Control. The interplay between primal drift $||\Delta_t||$ (red) and the stability threshold (dashed). The Stability Braking zone (Yellow) visualizes the algorithm strictly enforcing the thermodynamic speed limit. The controller detects imminent divergence and pauses cooling, preventing Premature Mode Collapse.
  • Figure 4: Training Dynamics on SPair-71k.Left: Standard annealing (Blue) hits the Trap, causing gradient collapse. Gumbel-Sinkhorn (Green) is stable but converges slowly due to variance. EPH-ASC (Red) achieves the fastest convergence. Right: The Adaptive Mechanism. Drift spikes (Gray) trigger the braking zone (Yellow), holding $\epsilon$ constant (Red line) to maintain thermodynamic stability.

Theorems & Definitions (23)

  • Proposition 2.1: Linear Scaling of the Stability Basin
  • Proposition 3.1: Sinkhorn Dynamics
  • proof
  • Theorem 3.2: Thermodynamic Speed Limit
  • Corollary 3.3: Inevitability of Collapse for Exponential Schedules
  • Theorem A.2: Sensitivity-Stability Duality
  • proof
  • Corollary A.3: The Basin Mismatch
  • proof
  • Definition A.4: Active reduced system
  • ...and 13 more