Nem sempre o mais profundo é melhor: Mitigando o Custo de Alinhamento via Decodificação de Camada Confiante

Resumo

A geração autorregressiva em modelos de linguagem de grande porte (LLMs) convencionalmente decodifica a partir da camada final, assumindo que representações mais profundas produzem previsões do próximo token mais confiáveis. Revisitamos essa suposição ao revelar uma dinâmica recorrente de Palpite-Refinamento-Perturbação: camadas iniciais formam palpites grosseiros, camadas intermediárias refinam semânticas relevantes ao raciocínio, e camadas finais podem perturbar essas previsões refinadas em direção a tokens genéricos ou preferidos por alinhamento. Apresentamos a Decodificação Confiante, uma estratégia de decodificação sem treinamento que seleciona dinamicamente a camada quase final mais confiável por meio de uma busca retroativa conservadora guiada por entropia. Além disso, fornecemos uma formulação teórica da seleção de camada como um problema de parada ótima, mostrando que, sob ruído de projeção limitado e perturbação de alinhamento dominante em estágio tardio, nossa regra de busca filtra a perturbação enquanto limita a perda em relação à camada de refinamento ideal. Experimentos em LLMs densos e de Mistura de Especialistas demonstram ganhos consistentes em benchmarks desafiadores de raciocínio, incluindo GPQA-Diamond, Omni-MATH e HLE, com zero de sobrecarga de memória e aumento de latência inferior a 2%. Esses resultados sugerem que contornar dinamicamente as perturbações da camada final pode desbloquear um comportamento de raciocínio mais forte em LLMs alinhados.

English

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.