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

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.

한 줄 요약

최종 레이어 대신 신뢰도 높은 중간 레이어를 선택하는 Confident Decoding이 추론 성능을 향상시킨다.

핵심 기여도

Guess-Refine-Perturb 레이어 동작 패턴을 식별하여, 최종 레이어가 항상 최적의 예측을 제공하지 않는다는 점을 밝힘.
Confident Decoding이라는 훈련 없이 사용 가능한 decoding 전략을 제안, 엔트로피 기반 보수적 역방향 탐색을 통해 신뢰도 높은 중간 레이어를 동적으로 선택.
GPQA-Diamond, Omni-MATH, HLE 등 복잡한 추론 벤치마크에서 기존 greedy decoding 대비 성능 향상 (최대 +2.6%)을 달성.
메모리 오버헤드 0%, 지연 시간 증가 2% 미만으로, 실용적 적용 가능성 확보.

핵심 아이디어

기존의 LLM 추론 방식은 최종 레이어의 출력만을 사용하지만, 이 연구는 중간 레이어가 더 정확한 추론 결과를 제공할 수 있음을 제시한다. 연구팀은 LLM의 forward pass 동작을 분석하여, Guess-Refine-Perturb라는 3단계 패턴을 발견했다. 즉, 초기 레이어는 대략적인 예측을 생성하고, 중간 레이어는 의미적 정제를 수행하지만, 최종 레이어는 정제된 정보를 방해하는 편향을 도입할 수 있다. 이는 alignment 과정에서 발생하는 "Alignment Tax"의 결과로, 추론 성능을 저하시킨다. 이를 해결하기 위해 Confident Decoding은 엔트로피가 낮은, 즉 예측 분포가 날카로운 레이어를 역방향으로 탐색하며 선택하는 방식을 제안한다. 이는 훈련 없이 적용 가능한 decoding 전략으로, 모델의 전체 forward pass를 유지하면서도 최적의 예측 레이어를 동적으로 선택한다.

기술적 접근법

**Confident Decoding**: 훈련 없이 적용 가능한 decoding 전략.
**엔트로피 기반 보수적 역방향 탐색**: 최종 레이어 근처에서 엔트로피가 낮은 레이어를 탐색하여 선택.
**Shannon entropy**: unembedded distribution의 엔트로피를 사용해 예측 신뢰도를 측정.
**KV cache 유지**: 모델 가중치, 키-벨류 캐시, 샘플링 과정은 변경 없이 유지.
**Optimal stopping problem**: 레이어 선택을 최적 정지 문제로 수식화, 투영 노이즈와 alignment 편향이 있을 때 손실을 제한.

주요 결과

**GPQA-Diamond**: 기존 greedy decoding 대비 +2.6% 성능 향상.
**Omni-MATH**: +2.1% 개선.
**HLE**: +1.8% 향상.
**모델 아키텍처**: dense 및 Mixture-of-Experts (MoE) 아키텍처에서 모두 일관된 성능 향상.
**지연 시간**: 2% 미만 증가.
**메모리 오버헤드**: 0%.

의의 및 한계

Confident Decoding은 alignment 과정에서 발생하는 추론 편향을 동적으로 제거함으로써, LLM의 내재된 추론 능력을 더 잘 발휘할 수 있도록 한다. 이는 복잡한 추론, 코드 생성, 긴 문맥 이해 등 다양한 벤치마크에서 실질적인 성능 향상을 보여준다. 또한, 훈련 없이 적용 가능한 decoding 전략으로, 기존 추론 파이프라인에 쉽게 통합할 수 있다. 그러나 이 접근법은 unembedding 행렬 $ W_U $가 중간 잔여 상태와의 구조적 정렬에 의존하기 때문에, 일부 경우 단어 집합 불일치로 인한 성능 저하가 발생할 수 있다. 또한, 이 방법은 decoding 단계에서만 alignment 편향을 완화할 뿐, 훈련 단계에서의 원인을 해결하지는 못한다.

실용적 활용

Confident Decoding은 복잡한 추론이 필요한 AI 어시스턴트, 코드 생성 도구, 긴 문맥 이해 시스템 등에 적용 가능하다. 특히, alignment 과정에서 성능 저하가 발생하는 대형 언어 모델에서 유용하며, 추론 과정을 개선하면서도 추가 훈련이나 메모리 사용 없이 적용할 수 있어 산업 현장에서 즉각적으로 활용 가능하다.