Resumen Técnico: Decodificación de Exploración de Profundidad (DEX)

1. Planteamiento del Problema

La decodificación autorregresiva (AR) en los Modelos de Lenguaje de Gran Escala (LLMs) es inherentemente costosa porque cada token generado debe pasar por todo el stack de capas antes de que se pueda producir el siguiente token. Estudios recientes han identificado una redundancia significativa a lo largo de la dimensión de profundidad: muchos tokens se vuelven predecibles en capas intermedias, lo que significa que no requieren estrictamente la profundidad total del modelo para ser precisos.

Los métodos actuales de adaptación de profundidad sin pérdida intentan explotar esto empleando una estrategia de selección de profundidad. Estos métodos proponen un token desde una única profundidad intermedia preseleccionada y lo verifican contra el modelo de profundidad final. Sin embargo, los autores identifican un cuello de botella de selección crítico:

• Salida Tardía (Late Exit): Si la profundidad de salida elegida es demasiado profunda, se desperdicia computación en capas que ya eran redundantes.
• Salida Temprana (Early Exit): Si la salida elegida es demasiado superficial y la predicción no coincide con la salida de la profundidad final, el sistema activa un mecanismo de retroceso (fallback). Esto no solo descarta el token actual, sino que también invalida cualquier borrador "descendiente" dependiente, lo que provoca pérdidas significativas de eficiencia.

El artículo postula que la brecha entre los métodos actuales basados en selección y un "oráculo" ideal (que conoce la profundidad más temprana exacta en la que un token se vuelve estable) es sustancial. La pregunta central es si esta redundancia puede explotarse de manera más efectiva reduciendo el riesgo de una decisión errónea de la profundidad de salida.

2. Metodología: Decodificación de Exploración de Profundidad (DEX)

Los autores proponen la Decodificación de Exploración de Profundidad (DEX), un algoritmo sin pérdida que reemplaza la selección de una única profundidad con la exploración paralela sobre múltiples profundidades candidatas.

Conceptos Clave

• Profundidad de Disponibilidad Temprana (EAD - Earliest Available Depth): Los autores definen EAD como la capa más superficial donde la predicción de un token coincide con la salida de la profundidad final y permanece estable a través de las capas subsiguientes.
• Conjunto de Exploración ($X$): En lugar de elegir una salida, DEX define un conjunto de profundidades candidatas $\{d_0, d_1, \dots, d_{N-1}\}$, donde $d_{N-1}$ es la profundidad final.
• Resolución: La sobrecarga de DEX está limitada por la "resolución" de este conjunto (la brecha máxima entre cualquier capa y la siguiente profundidad explorada). A medida que el conjunto se vuelve más fino, la sobrecarga tiende a cero.

El Algoritmo: Expansión–Compromiso–Colapso (Expand–Commit–Collapse)

DEX opera sobre un retículo de profundidad-posición gestionado por una Máquina de Estados Finitos (FSM). El proceso consta de tres fases:

1. Expansión del Retículo (Exploración Paralela):

   • El modelo se particiona en $N$ etapas de profundidad, cada una manejada por un "explorador de profundidad" separado (potencialmente en diferentes unidades de hardware).
   • En cada posición de compromiso, el sistema expande ramas candidatas en paralelo. Un token de exploración puede experimentar una expansión de profundidad (pasando su estado oculto al siguiente explorador de capa) o una expansión de posición (usando el candidato actual como prefijo para generar la siguiente posición de token).
   • Esto crea un árbol de futuros potenciales, donde diferentes ramas representan predicciones realizadas en diferentes profundidades.

2. Compromiso (Validación):

   • Una vez que la "ruta ancla" (la rama que atraviesa todas las capas) alcanza la profundidad final, produce el token de referencia de profundidad final.
   • DEX valida todos los tokens candidatos no finales de la ruta ancla contra esta referencia.
   • Selecciona la rama más superficial cuyo token coincida con la referencia de la profundidad final. Si no se encuentra ninguna coincidencia intermedia, recurre a la rama de la profundidad final.
   • Crucialmente, el token que se añade a la salida es siempre el producido por la referencia de la profundidad final, asegurando la equivalencia sin pérdida con la decodificación AR estándar.

3. Colapso (Poda de Estado):

   • Después de comprometer un token, el sistema colapsa el retículo. Las ramas que son inconsistentes con el nuevo prefijo comprometido se descartan.
   • Solo los estados (cachés KV, estados ocultos) de la rama seleccionada se preservan para el siguiente paso.
   • Esto asegura que, mientras el sistema explora múltiples profundidades, mantiene el estado exacto requerido para el siguiente paso AR.

Optimizaciones

• Muestreo Acoplado a la Profundidad (DECS - Depth-Coupled Sampling): Para evitar que diferentes profundidades en la misma cadena de ranuras propongan el mismo token (lo que desperdicia la exploración), DEX enmascara los tokens ya propuestos por profundidades más superficiales al decodificar candidatos más profundos. Esto fomenta la diversidad en las propuestas intermedias sin alterar la distribución de muestreo final.
• Adaptador de Inducción: Para LLMs estándar donde las capas intermedias no están naturalmente alineadas con la cabeza de salida, los autores proponen adjuntar pequeños adaptadores residuales a las profundidades intermedias. Estos se entrenan mediante autodestilación para alinear los estados ocultos intermedios con la cabeza LM final, mejorando la "preparación" de los tokens en profundidades más superficiales.

3. Contribuciones Clave

1. Formalización de la Preparación del Token: El artículo define la Profundidad de Disponibilidad Temprana (EAD) para cuantificar cuándo un token se vuelve predecible. Utiliza esto para demostrar matemáticamente el "cuello de botella de selección" en los métodos existentes, mostrando que su sobrecarga es impulsada por el riesgo de una selección errónea más que por el simple costo de la computación.
2. El Algoritmo DEX: Un algoritmo de decodificación sin pérdida que instancia la exploración de profundidad mediante un procedimiento de expansión–compromiso–colapso. Preserva la equivalencia exacta con la decodificación AR estándar mientras reduce el costo de comprometer tokens.
3. Validación Empírica: Los autores demuestran que DEX supera a las líneas base de selección de profundidad representativas (LayerSkip, AdaDecode, DEL) y logra un rendimiento (throughput) competitivo contra métodos de decodificación especulativa y distribuida (Lookahead, PEARL, EAGLE).

4. Resultados Experimentales

Los autores evaluaron DEX tanto en modelos entrenados para salida temprana (LayerSkip-trained CodeLlama-34B y Llama-2-70B) como en LLMs estándar (CodeLlama-34B, Llama-2-70B, Qwen3-32B) a través de tareas que incluyen razonamiento matemático (GSM8K), generación de código (HumanEval) y resumen (XSum).

• Ganancias de Rendimiento (Throughput): DEX alcanzó consistentemente más tokens por segundo (TPS) que las líneas base de selección de profundidad. En los modelos LayerSkip, DEX realizó una fracción mayor del margen teórico del "oráculo" que los métodos de selección.
• Escalabilidad con la Resolución: El rendimiento mejoró a medida que la resolución de la exploración se volvía más fina (es decir, se añadieron más exploradores). Por ejemplo, en Llama-2-70B, aumentar los exploradores de 3 a 8 redujo significativamente la brecha hacia la aceleración del oráculo.
• Competitividad: En comparaciones de extremo a extremo, DEX (con 3–4 exploradores) es competitivo con el paralelismo de tensores multi-GPU (TP) y la decodificación especulativa distribuida (PEARL, EAGLE) en hardware equivalente.
• Ablación: Eliminar las ramas intermedias (variante de salida única) mostró que la ganancia de rendimiento proviene específicamente de explorar múltiples profundidades, no solo de paralelizar las capas.

5. Significado y Reivindicaciones

El artículo afirma que la exploración de profundidad paralela proporciona una forma escalable de explotar el eje de profundidad subutilizado en la decodificación de LLM.

• Cambio de Paradigma: A diferencia de los métodos anteriores que apuestan por un único punto de salida, DEX trata la adaptación de profundidad como un problema de resolución controlada. Al explorar un conjunto de profundidades, convierte la incertidumbre de "cuándo salir" en una sobrecarga acotada determinada por el espaciado del conjunto de exploración.
• Eficiencia Sin Pérdida: DEX logra estas ganancias sin comprometer la corrección de la salida (sin pérdida) y sin requerir cambios en la arquitectura del modelo (excepto por adaptadores opcionales en modelos estándar).
• Complementariedad: Los autores señalan que DEX opera en el "eje de profundidad" (qué tan temprano se determina un token), lo cual es complementario a los métodos de aceleración del "eje de token" (cuántos tokens futuros se predicen). Sugieren que, si bien la combinación de ambos es posible, requiere una coordinación compleja.

En conclusión, el artículo establece que, al reemplazar la selección de una única profundidad con la exploración paralela, es posible reducir significamente la latencia de decodificación manteniendo el comportamiento exacto de los modelos autorregresivos estándar.