z0M1Shl1mi@OpenReview

Total: 1

#1 Collapse-Aware Regularization for Reliable Reasoning Under Distribution Shift [PDF] [Copy] [Kimi] [REL]

Authors: Quynh Vo, Cong-Duy T Nguyen

Reliable reasoning requires models to generalize under distribution shifts, yet in-distribution validation loss often fails to identify checkpoints that remain reliable out of distribution. We study representation collapse in Transformer hidden states as an internal degeneration that may reduce the effective capacity needed for multi-step inference. We characterize two complementary collapse signals: spectral capacity collapse, measured by the effective rank of layerwise representation covariance, and geometric alignment collapse, measured by average cosine alignment. We then propose the Collapse Risk Criterion (CRC), an ID-computable diagnostic score estimated from in-distribution validation representations. CRC is not a formal OOD guarantee, but a practical surrogate for OOD-free checkpoint selection. Across four reasoning benchmarks and three Transformer backbones, CRC correlates more strongly with OOD reasoning error than ID validation loss and standard confidence-based reliability scores. CRC-aware checkpoint selection improves OOD performance under depth/difficulty and template/rule shifts with minimal ID degradation. As an extension, collapse-aware training with a CRC-based regularizer improves average OOD performance in all evaluated dataset--backbone settings. Our results suggest that monitoring representation collapse is a simple and useful tool for improving reasoning reliability without OOD validation data.

Subject: UAI.2026 - Poster