Deeper is Not Always Better: Mitigating the Alignment Tax via Confident Layer Decoding
For practitioners using aligned LLMs, this work offers a simple, low-cost way to improve reasoning performance without retraining.
The paper reveals that final-layer representations in LLMs can introduce alignment-related perturbations that degrade reasoning, and proposes Confident Decoding, a training-free method that dynamically selects a near-final layer for decoding. Experiments show consistent gains on GPQA-Diamond, Omni-MATH, and HLE with zero memory overhead and <2% latency increase.
Autoregressive generation in large language models (LLMs) conventionally decodes from the final layer, assuming that deeper representations yield more reliable next-token predictions. We revisit this assumption by revealing a recurring Guess-Refine-Perturb dynamic: early layers form coarse guesses, intermediate layers refine reasoning-relevant semantics, and final layers can perturb these refined predictions toward generic or alignment-preferred tokens. We introduce Confident Decoding, a training-free decoding strategy that dynamically selects the most reliable near-final layer through entropy-guided conservative backward search. We further provide a theoretical formulation of layer selection as an optimal stopping problem, showing that under bounded projection noise and dominant late-stage alignment perturbation, our search rule filters perturbation while bounding the loss relative to the oracle refinement layer. Experiments across dense and Mixture-of-Experts LLMs demonstrate consistent gains on challenging reasoning benchmarks, including GPQA-Diamond, Omni-MATH, and HLE, with zero memory overhead and less than 2% latency increase. These results suggest dynamically bypassing final-layer perturbations can unlock stronger reasoning behavior from aligned LLMs.