Block-Based Double Decoders

Resumen Técnico: Doble Decodificador Basado en Bloques

Enunciado del Problema
El artículo aborda un compromiso fundamental en el diseño de arquitecturas de transformadores entre la eficiencia de entrenamiento y la eficiencia de inferencia. Los modelos basados únicamente en decodificador ofrecen una eficiencia de entrenamiento y escalabilidad superiores, pero sufren altos costos de inferencia, particularmente en lo que respecta a la memoria de caché KV y al cómputo por token. Por el contrario, los modelos codificador-decodificador ofrecen ahorros significativos en tiempo de inferencia (reducción de la caché KV), pero sufren de supervisión dispersa durante el preentrenamiento debido a objetivos como la corrupción de intervalos (donde solo ~15% de los tokens generan una señal de pérdida) y longitudes de secuencia dinámicas que dificultan el agrupamiento eficiente. Las alternativas existentes como PrefixLM intentan cerrar esta brecha, pero aún dejan muchos tokens sin entrenar y no logran la eficiencia completa de los codificador-decodificador.

Metodología: Doble Decodificador Basado en Bloques
Los autores proponen una arquitectura novedosa llamada Doble Decodificador Basado en Bloques, que combina dos pilas de decodificadores para lograr tanto supervisión completa de la pérdida como eficiencia en la inferencia.

1. Estructura de la Arquitectura:

   • Decodificador de Contexto: Un transformador estándar causal basado únicamente en decodificador que procesa la secuencia de entrada completa y emite latentes causales ($h_t$) para cada token.
   • Decodificador de Generación: Una segunda pila de decodificadores que toma tres entradas: los latentes causales del decodificador de contexto, la secuencia de tokens cruda y una partición en bloques.
   • Enmascaramiento Basado en Bloques Doble-Causal: La secuencia de entrada se divide en bloques contiguos. Dentro de un bloque, el decodificador de generación utiliza autoatención causal. Entre bloques, utiliza atención cruzada completa hacia los latentes del decodificador de contexto de los bloques anteriores. Esto crea una estructura "doble-causal" donde el flujo de atención avanza dentro de los bloques y desde los bloques anteriores.

2. Objetivo de Entrenamiento:

   • A diferencia de la corrupción de intervalos, que enmascara tokens, este método asegura que cada token en la secuencia de entrada contribuya a la pérdida exactamente una vez por pase hacia adelante.
   • La longitud de la secuencia permanece estática después del empaquetado, eliminando la necesidad de agrupamiento dinámico y relleno de tokens, lo que mejora el rendimiento.
   • Los autores implementan el mecanismo de atención calculando logits de autoatención y atención cruzada por separado y combinándolos mediante normalización log-sum-exp post-hoc (usando FlexAttention de PyTorch), ya que aún no está disponible una implementación unificada de atención de doble clave.

3. Mecanismo de Inferencia:

   • En tiempo de inferencia, el Decodificador de Contexto se ejecuta una vez sobre el prompt (relleno inicial). Dado que es un decodificador causal estándar, no requiere una caché KV para su propia salida (solo el Decodificador de Generación necesita cachear).
   • El Decodificador de Generación luego genera tokens de forma autoregresiva, atendiendo a los latentes de contexto en caché y a sus propias salidas anteriores.
   • Esta separación permite reutilizar la caché KV a nivel de prefijo (común en modelos basados únicamente en decodificador) mientras se mantienen los beneficios de memoria de las arquitecturas codificador-decodificador.

Contribuciones Clave

• Arquitectura Novel: Introducción del doble decodificador basado en bloques, que utiliza dos pilas de decodificadores para desacoplar el procesamiento de contexto de la generación de tokens.
• Enmascaramiento Doble-Causal: Una nueva máscara de atención que permite supervisión completa de tokens (densidad de pérdida del 100%) mientras mantiene longitudes de secuencia estáticas para un agrupamiento eficiente.
• Optimizaciones de Inferencia: La arquitectura reduce la memoria de caché KV y el cómputo por token en aproximadamente 2/3 en comparación con los modelos basados únicamente en decodificador (asumiendo una división de 2/3 para el decodificador de contexto y 1/3 para el decodificador de generación) sin sacrificar las capacidades de caché de relleno inicial.
• Eficiencia de Entrenamiento: El método evita la supervisión dispersa y los problemas de agrupamiento dinámico inherentes a los codificador-decodificador entrenados con corrupción de intervalos.

Resultados Experimentales
Los autores realizaron experimentos de leyes de escalado comparando Modelos Basados Únicamente en Decodificador, Codificador-Decodificador Estándar (SED) y Doble Decodificador Basado en Bloques a través de conteos de parámetros (6.25M–100M) y presupuestos de tokens (62.5M–1B).

• Rendimiento: El Doble Decodificador supera ampliamente la línea base SED y sigue de cerca al modelo Basado Únicamente en Decodificador. A la escala más grande (100M parámetros, 1B tokens), el Doble Decodificador logra una pérdida de evaluación solo ~0.2 nats peor que el modelo Basado Únicamente en Decodificador coincidente, mientras que el SED se queda atrás de ambos en ~0.7 nats.
• Comportamiento de Escalado: Se encontró que los modelos SED estaban en un "régimen limitado por datos", donde aumentar los parámetros no mejoraba el rendimiento debido a la supervisión dispersa. En contraste, los Doble Decodificadores y los modelos Basados Únicamente en Decodificador siguieron las leyes de escalado estándar, aplanándose al alcanzar los límites de capacidad.
• Compromiso de Cómputo: Aunque los Doble Decodificadores incurren en un ligero aumento de FLOPs de entrenamiento (~2.4% más que los Basados Únicamente en Decodificador en configuraciones específicas) debido a proyecciones KV adicionales, esto se ve superado por los beneficios de inferencia.

Significado y Afirmaciones
El artículo afirma que los Doble Decodificadores Basados en Bloques resuelven con éxito la dicotomía histórica entre la eficiencia de entrenamiento y la de inferencia. Al restaurar la supervisión completa de la pérdida a una estructura similar a la de un codificador-decodificador, los autores demuestran que es posible lograr la eficiencia de inferencia de los codificador-decodificador (específicamente reducción de caché KV y latencia por token) mientras se retienen la eficiencia de entrenamiento y las propiedades de escalado de los modelos basados únicamente en decodificador.

Los autores enfatizan que la brecha modesta de pérdida en tiempo de entrenamiento (~0.2 nats) es un costo necesario para la separación arquitectónica que permite ahorros significativos en tiempo de inferencia. Argumentan que esto hace que la arquitectura sea particularmente adecuada para dispositivos periféricos y escenarios con restricciones de memoria, al tiempo que habilita nuevas direcciones de investigación como el razonamiento de cadena de pensamiento latente durante el preentrenamiento y la aplicación de transformadores en bucle específicamente al decodificador de contexto.