arxiv:2601.23039

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

Published on Jan 30

· Submitted by

YIZHI LIU on Feb 9

Upvote

Authors:

Yizhi Liu

Abstract

Researchers identify and address premature mode collapse in optimal transport-based structural prediction models through an adaptive stability control algorithm that prevents gradient explosions during large-scale training.

AI-generated summary

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

View arXiv page View PDF Add to collection

Community

leon0923

Paper author Paper submitter 1 day ago

Excited to share my latest work: "Avoiding Premature Collapse: Adaptive Annealing for Entropy-Regularized Structural Inference" (arXiv:2601.23039) 🚀

We identify a key failure mode in Sinkhorn annealing (ε→0): Premature Mode Collapse, caused by a "Thermodynamic Speed Limit" — standard exponential schedules violate stability due to O(1/ε) sensitivity and vanishing spectral gap.

Solution: EPH-ASC (Efficient Piecewise Hybrid Adaptive Stability Control) — lightweight adaptive scheduler that monitors primal drift and triggers "Thermodynamic Pause" when needed. Overhead <0.5%!

Results:

SPair-71k keypoint matching: 1.6× speedup over Gumbel-Sinkhorn (47 vs 75 epochs to 90%)
FineWeb-Edu + Nano Gemma w/ mHC: detects instability early, prevents late gradient explosion with 340-step safety margin

Also built an interactive demo to visualize annealing strategies: https://huggingface.co/spaces/leon0923/torch-sinkhorn-asc-demo

Would love feedback from OT, routing, or mHC folks! How does this resonate with your experiences in large-scale training? @leon0923 (me) happy to discuss / collaborate on integrations.

#OptimalTransport #Sinkhorn #MachineLearning #mHC