Más profundo no siempre es mejor: Mitigando el costo de alineación mediante decodificación de capas con confianza

Resumen

La generación autorregresiva en modelos de lenguaje grandes (LLMs) decodifica convencionalmente desde la última capa, asumiendo que las representaciones más profundas generan predicciones del siguiente token más fiables. Revisamos esta suposición al revelar una dinámica recurrente de Adivinar-Refinar-Perturbar: las capas tempranas forman conjeturas gruesas, las capas intermedias refinan la semántica relevante para el razonamiento, y las capas finales pueden perturbar estas predicciones refinadas hacia tokens genéricos o preferidos por el alineamiento. Introducimos Decodificación Confiada, una estrategia de decodificación sin entrenamiento que selecciona dinámicamente la capa cercana a la final más fiable mediante una búsqueda hacia atrás conservadora guiada por entropía. Además, proporcionamos una formulación teórica de la selección de capas como un problema de parada óptimo, mostrando que, bajo ruido de proyección acotado y perturbación de alineamiento dominante en etapas tardías, nuestra regla de búsqueda filtra la perturbación mientras acota la pérdida en relación con la capa de refinamiento oráculo. Experimentos en LLMs densos y de Mezcla de Expertos demuestran ganancias consistentes en puntos de referencia de razonamiento desafiantes, incluyendo GPQA-Diamond, Omni-MATH y HLE, con cero sobrecarga de memoria y menos del 2% de aumento en latencia. Estos resultados sugieren que sortear dinámicamente las perturbaciones de la capa final puede desbloquear un comportamiento de razonamiento más sólido en LLMs alineados.

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.