PaperlayerView exact revision on arXiv

Source-linked AI summary

Deeper is Not Always Better: Mitigating the Alignment Tax via Confident Layer Decoding

Xuanming Zhang, Sining Zhoubian, Yuxuan Chen, Tianyi Tang, An Yang, Sean Du, Chujie Zheng, Fei Huang, Dayiheng Liu, Gao Huang, Jingren Zhou

arXiv:2606.21906v1cs.CL

TL;DR

Final-layer decoding can perturb reasoning-relevant predictions, so this paper introduces Confident Decoding to dynamically select a reliable near-final layer. Across dense and Mixture-of-Experts models, it consistently improves benchmark performance with zero memory overhead and less than 2% additional latency.

  • Problem

    Final-layer representations are not always the most reliable source for reasoning-sensitive next-token prediction because later layers can perturb refined semantics.

  • Method

    Confident Decoding is a training-free strategy that dynamically selects near-final logits through entropy-guided conservative backward search.

  • Results

    Confident Decoding consistently improves over greedy decoding and strong contrastive baselines across dense and Mixture-of-Experts models, with zero memory overhead and less than 2% additional latency.

  • Takeaways & Limitations

    Dynamically bypassing final-layer perturbations can unlock stronger reasoning behavior from aligned language models.

  • Takeaways & Limitations

    The approach is constrained by alignment between the unembedding matrix and intermediate residual states, and shallow layers may suffer vocabulary mismatch.

Abstract

from arXiv · show

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.

1 Introduction

The paper challenges final-layer decoding by identifying a recurring Guess–Refine–Perturb dynamic in which later layers can disrupt refined reasoning signals. It proposes Confident Decoding, a training-free strategy that selects reliable near-final logits and reports improvements across diverse benchmarks and model architectures.

  • Motivation: Intermediate layers can encode task-relevant semantics, while later layers may compress, redirect, or perturb already-refined information.This challenges the assumption that representation reliability increases monotonically with depth.
  • Motivation: Post-training objectives may encourage final-layer distributions to favor alignment-preferred outputs, contributing to late-stage perturbation.The passage identifies continued pre-training, supervised fine-tuning, RLHF, RLAIF, and DPO as relevant procedures.
  • Method: Confident Decoding dynamically selects the most reliable near-final layer at each generation step without truncating the transformer or modifying its forward pass.It computes candidate logits from a small near-final layer window and selects using local entropy.
  • Results: Across dense and Mixture-of-Experts architectures, Confident Decoding improves performance on reasoning, mathematics, long-context, coding, and safety benchmarks.The evaluation includes GPQA-Diamond, HLE, Omni-MATH, LongBench v2, LiveCodeBench v6, and Air-Bench-2024.
  • Contributions: The paper provides theoretical and empirical evidence that Confident Decoding filters late-layer perturbations while preserving refined reasoning signals.The claimed benefits span diverse reasoning, coding, long-context, and safety benchmarks.

2 Preliminaries

The preliminaries formalize layer-wise representational dynamics through contribution magnitude and directional fidelity, revealing a Guess–Refine–Perturbation progression. They motivate entropy-valley decoding as a token-dependent alternative to static early exits, which degrade accuracy when applied frequently.

  • Layer-wise dynamics: The analysis tracks Relative Contribution Norm and Residual I/O Cosine Similarity to measure layer write intensity and directional semantic fidelity.High IO-CosSim indicates preservation and refinement of the existing semantic trajectory, while lower values indicate rotation into a distinct subspace.
  • Guess–Refine–Perturbation: 1.6 and 0.67 characterize the first layer’s high-magnitude, low-fidelity update, indicating rapid construction of an initial latent representation under high uncertainty.Norm Ratio(1) ≈1.6, while IO-CosSim(1) ≈0.67.
  • Guess–Refine–Perturbation: 0.23–0.57 and 0.91–0.97 characterize the intermediate refinement phase, where layers make incremental, directionally faithful updates that integrate contextual information.The Norm Ratio remains below 1, while IO-CosSim stays consistently high across intermediate layers.
  • Guess–Refine–Perturbation: Final-layer perturbation breaks monotonic improvement: Norm Ratio(40) peaks at 2–3× the Phase II level and IO-CosSim(40) ≈0.69.This motivates extracting predictions near the end of Phase II rather than assuming the final layer is optimal.
  • Guess–Refine–Perturbation: Qwen3.5-35B-A3B exhibits three phases: shallow layers guess, intermediate layers refine, and final layers perturb the semantic trajectory.The pattern was also validated on additional architectures in Appendix B.1.
  • Dynamic layer selection: Static early exits cause a precipitous accuracy decline as execution probability increases, whereas the Entropy Valley selects a local entropy minimum by scanning backward from the final layer.Dynamic valley selection maintains high accuracy, while decoding from neighboring Valley±k layers consistently degrades performance.

3 Theoretical Grounding

The theoretical analysis models layerwise decoding as Guess–Refine–Perturb dynamics, with an intermediate representational optimum before possible alignment perturbation. It formalizes conservative backward search as an optimal stopping rule that filters perturbations while bounding semantic loss.

  • Information-Theoretic Dynamics: The analysis identifies an intermediate representational zenith where predictive entropy reaches its minimum before potential final-layer perturbation.For Qwen3.5-35B-A3B on GPQA Diamond, V_onset ≈ 28 and V* = 39, with Phase II accumulating ΔH ≈ −10.8 nats across 202,935 tokens.
  • Information-Theoretic Dynamics: The Guess phase compresses superficial input and creates unstable, high-entropy representations, whereas the Refine phase integrates context and monotonically improves predictions after V_onset.The residual compression noise carries into early refinement, while the late-refine window exhibits increasing mutual information and strictly decreasing predictive entropy.
  • Alignment Perturbation: 16.2% of tokens exhibit an entropy rise of +0.37 nats at the final layer, while 83.8% remain unperturbed and do not enter Phase III.The perturbed-token pattern is the stated alignment-tax signature, whereas unperturbed content tokens remain uncertain at V*.
  • Optimal Stopping: Theorem 1 guarantees that conservative backward search selects ˆV within [V_onset, V*], filtering alignment perturbations while bounding semantic precision loss.When final-layer alignment conflicts with task-optimal predictions by more than projection noise, the scan avoids Phase III; synergistic alignment naturally permits stopping near L.
  • Optimal Stopping: The selection is oracle-optimal under strong refinement signals and incurs bounded loss under weak local oscillations, because diminished semantic gradients make I(h(k); Y) approximately equal to I(h(V*); Y).The resulting loss is bounded by the integral of the diminished gradient over [k, V*], while projection-noise penalties remain asymptotically negligible.

4 Methodology

Confident Decoding is a training-free inference method that dynamically forwards logits from the most entropy-confident near-final layer instead of always using the final layer. It preserves the complete forward pass and standard sampling while adding candidate-layer extraction, entropy evaluation, and a bounded backward scan.

  • Confident Decoding: At each decoding step, Confident Decoding selects the near-final layer with the lowest predictive Shannon entropy and forwards its logits to the sampler.The model still computes the full forward pass; only the logits source layer changes.
  • Candidate set: The method evaluates candidate layers in a near-final window C = {L − M + 1, ..., L}, where M controls both backtracking extent and computational cost.Larger M permits more backtracking but increases computation.
  • Entropy-trough selection: The default entropy-trough rule scans backward from layer L and freezes each token’s choice at the first local entropy valley, stopping after the first non-improvement.Selection is performed independently per token using a bounded scan window K ≤ M.
  • Optional fallback: The fallback probability mixes entropy-valley selection with standard final-layer decoding: p = 0 exactly recovers final-layer decoding, while p = 1 always selects the entropy-valley layer.This stochastic fallback also provides a numerical regression check against standard decoding.
  • Production implementation: The implementation leaves transformer execution, KV caching, attention kernels, prefix caching, batching, and downstream sampling unchanged, extracting candidate residual states outside the compiled region.Normalization, unembedding, entropy computation, and stateful logic remain in the eager wrapper for graph-safe execution.
  • Cost: The overall per-step computation cost is O(MBd|V|), consisting of candidate unembedding, entropy evaluation, and a K-step vectorized trough scan.Here B is the number of positions requiring sampling, and K ≤ M bounds the scan cost.

5 Experiments

Experiments evaluate Confident Decoding across diverse reasoning benchmarks and open-weight dense and MoE models, finding consistent gains that are especially large on complex reasoning while preserving performance on creativity, long-context, and safety tasks. Instruct-versus-base ablations further attribute larger benefits to post-training alignment perturbations.

  • Cross-architecture evaluation: Confident Decoding improves performance consistently across dense and MoE architectures, including Gemma-4, Qwen3.5, and gpt-oss model families.The evaluation spans varying parameter scales and reports gains in both dense and highly sparse MoE topologies.
  • Reasoning performance: +10.1% absolute gain appears for Qwen3.5-27B on LiveCodeBench-v6, alongside +6.5% for Qwen-35B on GPQA-D and +2.4% on HLE.These results identify the largest benefits on challenging, fragile reasoning and coding tasks.
  • Creativity and safety: +0.1% to +0.5% gains on WritingBench indicate stability in creative writing and stylistic formatting.The reported results address the risk that bypassing final layers could damage creative structure.
  • Alignment-tax ablation: +2.8% average gain for the Instruct model exceeds the +1.7% gain for the Base model, supporting an alignment-induced final-layer degradation effect.The ablation compares Qwen3.5-35B-A3B-Base with its instruction-tuned counterpart and interprets the larger delta as evidence for learned post-training perturbations.
  • Alignment-tax ablation: Under greedy decoding, the Instruct model scores 7.1% on HLE versus 8.0% for the Base model, exposing reasoning degradation that Confident Decoding is designed to reverse.The passage attributes this gap to generic or safe priors overriding fragile multi-hop logic at the final layer.
  • Token-level analysis: 12.8% of Instruct tokens versus 10.4% of Base tokens contain identified entropy valleys, with substitution rates of 2.60% and 2.36%, respectively.Approximately 21% of identified valleys change the decoded token, while the remainder alter entropy without changing the argmax.

6 Discussions

Confident Decoding intervenes selectively when final-layer entropy rebounds, with the largest benefits on difficult reasoning tasks across model architectures. It achieves these gains through lightweight, cache-reusing computation with negligible latency and no additional KV-cache memory.

  • Rollback dynamics: 2.47% of all generated tokens undergo hard argmax substitution, while the remaining rebound cases reduce entropy without changing the mode.Among scan-triggering tokens, 21.4% change the argmax and 78.6% redistribute probability mass toward competing candidates.
  • Rollback dynamics: The non-trivial valley rate rises from 9.7% in the first decile to 13.4% at 60–70% of generation, then falls to 8.0% in the final decile.The substitution rate follows the same inverted-U pattern, moving from 1.82% to 3.03% to 1.47%.
  • Scaling with task difficulty: Performance gains increase with task difficulty: Confident Decoding can fall 0.4% below baseline on Level 1 MATH but protects reasoning on harder tasks.Level 1 baseline accuracy exceeds 97%, whereas Level 4 Omni-MATH reduces gpt-oss-20b to 1.1% accuracy before decoding intervention.
  • Scaling across architectures: Confident Decoding yields positive average gains on every backbone across Qwen, gpt-oss, and Gemma families, including dense, hybrid-attention, and MoE architectures.The four reported MoE gains include Qwen3.5-35B-A3B (+2.8) and gpt-oss-120B (+1.7).
  • Computational overhead: < 2% per-token latency increase and zero additional KV-cache memory overhead make the method computationally lightweight.The method reuses materialized intermediate states and performs additional unembedding projections without extra attention, FFN, or MoE routing computations.

7 Related Work

Prior work shows that intermediate layers can contain stronger, more task-relevant representations than final layers, whose alignment objectives may perturb or overfit predictions. Existing contrastive decoding and early-exit methods motivate Confident Decoding, which instead uses an entropy valley to avoid corrupted terminal distributions.

  • Layer-wise Dynamics and Mechanistic Interpretability: Intermediate layers can encode stronger factual and semantic representations, while final layers may overfit alignment objectives and contribute less to the residual stream.Prior studies localize factual knowledge to narrow mid-layer MLP windows and report substantial redundancy among deeper layers.
  • Contrastive and Adaptive Decoding: Contrastive decoding exploits layer-wise knowledge by contrasting expert and amateur logits, with DoLa extending this idea within a single model.Later methods refine the contrastive signal through gradient-style corrections and adaptive layer weighting.
  • Contrastive and Adaptive Decoding: Existing contrastive methods anchor on final-layer distributions, inheriting alignment-tax and late-stage perturbation effects that Confident Decoding avoids with an entropy valley.Confident Decoding discards dependence on the final layer as the reference distribution.
  • Test-Time Computation and Optimal Stopping: Adaptive halting and early-exit methods originally targeted efficiency by stopping computation when learned or token-wise confidence criteria were met.Universal Transformers introduced adaptive per-token halting, while CALM and LayerSkip pursued latency reduction through early exits and self-speculative decoding.
  • Test-Time Computation and Optimal Stopping: Latency-driven exits can confuse valid convergence with shallow biases, whereas test-time compute scaling improves reasoning and uppermost layers concentrate RLHF-related capability degradation.These findings motivate treating layer selection as a reasoning and alignment problem rather than solely a speed optimization.

8 Conclusion · Appendix

The paper argues that final-layer representations can degrade reasoning through late-stage alignment perturbations, and presents Confident Decoding as a training-free optimal-stopping intervention. It also identifies structural limitations and calls for training and multimodal research to address the underlying conflict.

  • 8 Conclusion: The conclusion challenges the assumption that final-layer representations always provide the optimal semantic state.The paper formalizes layerwise Guess-Refine-Perturbation dynamics to motivate this challenge.
  • 8 Conclusion: Late-stage alignment constraints can corrupt otherwise carefully constructed reasoning chains.This is characterized as the Alignment Tax arising during layerwise processing.
  • 8 Conclusion: Confident Decoding is introduced as a training-free, risk-averse optimal stopping algorithm for navigating the planning-pragmatics tradeoff.It provides an inference-time intervention rather than changing model training.
  • 8 Conclusion: Confident Decoding is constrained by the structural alignment of the unembedding matrix W_U with intermediate residual states.The theoretical framework bounds projection noise, but shallow representations may still exhibit vocabulary mismatch.
  • 8 Conclusion: The approach mitigates alignment-tax symptoms during decoding rather than resolving their root cause during training.This limitation concerns the distinction between inference-time intervention and training-phase resolution.
  • 8 Conclusion: The findings expose an architectural conflict between pre-trained reasoning and post-training pragmatics that warrants deeper mechanistic exploration.The conclusion frames this conflict as a central direction for future investigation.
  • 8 Conclusion: Future training paradigms could decouple reasoning and alignment objectives by applying alignment penalties exclusively to designated routing heads rather than the core residual stream.This is proposed as one avenue for addressing the identified architectural conflict.
  • 8 Conclusion: Future work should also examine whether Guess-Refine-Perturbation dynamics persist in multimodal foundation models.The passage identifies multimodal persistence as an additional research direction.

A Hyperparameters and Configurations · B Three-Phase Structure on More Models

The experiments use fixed, greedy-style decoding and show that CONFIDENT DECODING is highly sensitive to its valley-selection probability and temperature. Replication on Qwen3.5-9B indicates the three-phase layer dynamics persist beyond the original MoE model and scale.

  • A Hyperparameters and Configurations: K=10 and p=1.0 are the default CONFIDENT DECODING settings throughout the main experiments.K controls the backward-scan lookback window, while p selects the entropy-valley token over the standard final-layer token.
  • A Hyperparameters and Configurations: T=0.0 and top-p=1.0 provide greedy-style decoding aligned with official benchmark configurations, isolating layer selection as the source of performance changes.Benchmark prompts and evaluation settings are inherited from the corresponding official repositories, with the serving backend swapped to a vLLM-hosted endpoint.
  • A Hyperparameters and Configurations: 82.8% to 80.8% accuracy accompanies increasing temperature from 0.0 to 1.0 at p=1.0, while 82.8% to 76.3% accompanies decreasing p from 1.0 to 0.0 at T=0.0.Both changes degrade accuracy, with the p ablation producing a strict, monotone collapse and p=0.0 recovering standard greedy decoding.
  • B Three-Phase Structure on More Models: Qwen3.5-9B replicates the layer-wise and per-token analyses using a 32-layer compact hybrid architecture without sparse expert routing.The replication tests whether the three-phase structure and Phase III heterogeneity depend on MoE architecture or model scale.
  • B Three-Phase Structure on More Models: On Qwen3.5-9B, the norm ratio shows Phase I overwrite at l=1 (≈8.5), stable Phase II (0.26–0.58), and Phase III resurgence at l=32 (≈0.84).These layer-wise dynamics are measured on GSM8K.
  • B Three-Phase Structure on More Models: IO-CosSim remains high throughout Phase II (0.88–0.97) before dropping to 0.80 at l=32, making the three-phase structure qualitatively identical to the 35B-A3B result.The gray bands identify the eight full-attention layers at l ∈{4, 8, 12, 16, 20, 24, 28, 32}.

B.1 Layer-wise Dynamics · B.2 Per-Token Entropy Partitioning

Across Qwen3.5-9B-Base, layer-wise behavior preserves a three-phase progression: initial embedding overwriting, directionally faithful refinement, and late-stage directional deflection. On GPQA Diamond, final-layer entropy rises for 47.4% of 9B tokens, substantially above the 16.2% observed for 35B-A3B, consistent with a compressed refinement corridor.

  • B.1 Layer-wise Dynamics: Qwen3.5-9B-Base preserves the three-phase layer-wise progression observed in the 35B-A3B analysis.Figure 5 reports Relative Contribution Norm and Residual I/O Cosine Similarity on GSM8K.
  • B.1 Layer-wise Dynamics: Phase I at l=1 overwrites the token embedding, with Norm Ratio≈8.5 and IO-CosSim≈0.14.
  • B.1 Layer-wise Dynamics: Phase II across 5 ≤l ≤31 performs incremental, directionally faithful refinement, with Norm Ratio 0.26–0.58 and IO-CosSim 0.88–0.97.
  • B.1 Layer-wise Dynamics: Phase III at l=32 shows the largest late-regime directional deflection, as Norm Ratio resurges to ≈0.84 while IO-CosSim drops to ≈0.80.The Norm Ratio is roughly 2× the Phase II plateau, matching the pattern observed in 35B-A3B.
  • B.2 Per-Token Entropy Partitioning: Tokens are classified as perturbed when entropy rises at the final layer, H(32)t, and unperturbed otherwise.
  • B.2 Per-Token Entropy Partitioning: The higher 9B perturbation rate is consistent with a compressed refinement corridor, where L=32 versus L=40 leaves fewer layers for Phase II refinement.The final full-attention layer is therefore more likely to overshoot, although the qualitative three-phase pattern remains preserved.

C Token-Level Substitution Analysis

This section examines which tokens Confident Decoding changes when the entropy-trough layer disagrees with the final layer. It reports frequent substitutions and categorizes substitution events for Base and Instruct Qwen3.5-35B-A3B on GPQA-Diamond.

  • Token-level substitution analysis: The analysis compares tokens selected by the entropy-trough layer against those selected by the final layer when their predictions disagree.This isolates the token-level changes made by Confident Decoding.
  • Token-level substitution analysis: Table 8 lists the 20 most frequently substituted tokens for Base Qwen3.5-35B-A3B on GPQA-Diamond.The table focuses on the Base variant's most common substitutions.
  • Token-level substitution analysis: Table 8 also lists the 20 most frequently substituted tokens for Instruct Qwen3.5-35B-A3B and provides a category breakdown of all substitution events.The breakdown covers substitution events across both reported variants.

D More Details of the Baselines

The paper compares Confident Decoding with DoLa and SLED, adapting both baselines to hybrid MoE models. These contrastive methods improve over last-layer decoding but underperform Confident Decoding because hybrid architectures violate their assumptions of cross-layer representational comparability and smooth logit evolution.

  • Baseline adaptation: DoLa and SLED were originally developed for dense Transformers, with implementations relying on architecture-specific layer indexing and residual-stream access.The authors re-implemented both methods in Confident Decoding’s inference framework and validated them on LLaMA-family models.
  • Baseline comparison: Both contrastive baselines yield modest improvements over last-layer decoding, but their gains fall substantially short of Confident Decoding on Qwen3.5-35B-A3B.The comparison uses re-implemented DoLa and SLED versions that support the hybrid MoE architecture.
  • Structural limitations: DoLa’s contrastive signal is weakened on hybrid backbones because DeltaNet and full-attention layers create discontinuous shifts in representation geometry.This breaks DoLa’s assumption that early and late layers occupy comparable subspaces differing primarily in depth.
  • Structural limitations: SLED is impaired by the non-monotonic entropy profile caused by layer alternation and expert switching, which corrupts its layerwise logit corrections.Both baselines generally assume representational homogeneity across contrasted layers, unlike Confident Decoding’s independent entropy evaluation of candidate layers.

E Degradation Analysis

Confident Decoding degrades mainly when hybrid layer alternation creates probe mismatch and limited depth compresses the refinement corridor. Its effects are architecture- and task-dependent: some models improve broadly, while others regress on specific benchmarks.

  • Architecture- and task-dependent degradation: Qwen3.5-9B shows mixed results, regressing on GPQA-D and Omni-MATH but improving on LCB-v6 and Air-Bench.Reported changes include GPQA-D: 64.6 →62.1, Omni-MATH: 49.1 →47.1, and LCB-v6: 41.1 →47.7.
  • Hybrid architecture: Alternating DeltaNet and full-attention layers can make shared-unembedding entropy probes geometrically mismatched to intermediate representations.This structured projection noise can distort the observed entropy valley and cause rollback to preconvergent representations.
  • Task dependence: Degradation concentrates on GPQA-D and Omni-MATH because their shallower refinement-corridor entropy gradients are more vulnerable to projection noise.When structured noise masks the true entropy gradient, the backward scan may select a non-optimal layer.
  • Model depth: Greater depth widens the late refinement corridor, giving the backward scan more same-type layers and a larger margin to locate a reliable entropy valley.Qwen3.5-9B has L=32 with only 8 full-attention layers distributed across the stack, leaving few layers in its final refinement zone.
  • MoE architectures: MoE robustness may reflect stronger late-stage refinement signals and reduced effective representation shifts, but depth and routing remain experimentally confounded.The paper leaves disentangling these two factors to future work.
Loading 2606.21906v1…