Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Resumen Técnico: Entrenamiento de Modelos de Lenguaje de Difusión de Bloques Híbridos con Bidireccionalidad Parcial

Planteamiento del Problema
La generación de contexto largo de alto rendimiento en modelos de lenguaje de gran tamaño (LLMs) está limitada principalmente por el ancho de banda de la memoria en lugar del cómputo. En la decodificación autorregresiva estándar, cada paso requiere el flujo del caché de clave/valor (KV) acumulado desde la memoria, lo que provoca que la demanda de ancho de banda escale linealmente con la longitud del contexto mientras solo se emite un token. Aunque existen dos enfoques paralelos para mitigar esto —reducir el acceso a la memoria mediante mezcladores de tiempo lineal (p. ej., Mamba) y aumentar el paralelismo mediante la generación por bloques—, la combinación de ambos presenta un obstáculo técnico específico.

Los modelos de difusión híbridos anteriores, como DiffuMamba, utilizan escaneos Mamba totalmente bidireccionales, incluyendo un paso de dirección inversa sobre toda la ventana de denoisado (denoising). Este escaneo inverso depende de los tokens activos, lo que significa que sus estados internos no son exclusivamente de prefijo. Por consiguiente, estos estados no pueden cachearse y reutilizarse con precisión entre bloques, lo que anula las ganancias de eficiencia de la difusión por bloques. El desafío central es diseñar una arquitectura híbrida que soporte la generación por bloques con un caché exacto, reteniendo al mismo tiempo los beneficios de la mezcla bidireccional.

Metodología
Los autores proponen BDLM Mamba-H, una arquitectura híbrida que combina Modelos de Lenguaje de Difusión de Bloques (BDLMs) con capas de Mamba y de atención. La innovación principal es un mecanismo de bidireccionalidad parcial que restringe el escaneo Mamba de dirección inversa únicamente al bloque de denoisado activo, en lugar de a toda la secuencia.

Componentes metodológicos clave incluyen:

1. Objetivo de Entrenamiento Alineado con el Caché: El modelo emplea un objetivo de "todos los bloques" donde los bloques de prefijo completados se almacenan como:
   • Estados KV para las capas de atención.
   • Estados de frontera de solo avance (estados recurrentes) para las capas Mamba.
     El escaneo inverso se confina al bloque activo actual ($b_t$), asegurando que el caché del prefijo ($C_\ell(c)$) sea una función únicamente del prefijo limpio, independiente del paso de tiempo de difusión inversa ($t$).
2. Factorización del Paso de Tiempo: Para mantener la validez del caché, la arquitectura factoriza el cómputo de la capa en un operador de prefijo libre de paso de tiempo y una continuación de bloque activo condicionada al paso de tiempo. Esto evita que el caché del prefijo limpio se vuelva dependiente del paso de tiempo específico de la difusión inversa, un requisito para la reutilización válida durante la inferencia.
3. Protocolo de Entrenamiento: El modelo se entrena utilizando un objetivo de difusión discreta sobre bloques. El caché del prefijo se construye mediante el escaneo del prefijo limpio, y las pérdidas de los bloques descendentes entrenan directamente las representaciones almacenadas en estos estados de prefijo, alineando el cómputo de entrenamiento con la interfaz de caché del muestreador en tiempo de inferencia.

Contribuciones Clave

• Entrenamiento Alineado con el Caché: Los autores formulan un objetivo BDLM de todos los bloques donde los objetos en caché (estados KV y estados de prefijo Mamba) coinciden exactamente con los cachés de bloque en tiempo de inferencia. Esto permite que las pérdidas de los bloques descendentes entrenen el escaneo de prefijo utilizado para la reutilización del estado recurrente.
• Validación Controlada y Escalamiento: Los autores realizan un barrido exhaustivo de 87 millones de parámetros en el conjunto de datos DCLM y escalan modelos seleccionados a 350 millones de parámetros utilizando una regla de transferencia de tasa de aprendizaje de tipo Complete(d)P.
• Ganancias de Rendimiento en Contexto Largo: El estudio evalúa el rendimiento (throughput) de inferencia en longitudes de contexto extendidas, demostrando que el enfoque híbrido supera significativamente tanto a los baselines de secuencia completa como a los modelos de difusión basados puramente en atención a medida que aumenta la longitud del contexto.

Resultados

• Calidad de Validación (87M de Parámetros): BDLM Mamba-H logró la mejor perplejidad de validación C4-en (61.6) entre los modelos de 87M, superando a BDLM atención (76.5), DiffuMamba-H de secuencia completa (83.7) y atención de secuencia completa (87.6).
• Calidad de Validación (350M de Parámetros): A la escala de 350M, BDLM Mamba-H se mantuvo competitivo con BDLM atención, logrando una perplejidad C4-en de 38.2 frente a los 36.97 de BDLM atención.
• Rendimiento de Throughput:
  • A 65K tokens, BDLM Mamba-H logró 19.7x el throughput del baseline de secuencia completa DiffuMamba-H.
  • A 262K tokens, BDLM Mamba-H logró 3.7x el throughput de BDLM atención.
  • Aunque los modelos de secuencia completa son más rápidos en contextos cortos (8192 tokens) debido a la sobrecarga del objetivo de todos los bloques, los modelos BDLM nativos con caché demuestran una escalabilidad superior en regímenes de contexto largo.

Significancia y Reivindicaciones
El artículo afirma que BDLM Mamba-H representa una arquitectura viable para el lenguaje de difusión de contexto largo. Al restringir el escaneo Mamba inverso al bloque activo, el modelo permite un caché exacto entre bloques, resolviendo el problema de la reutilización de estados inherente a los híbridos de difusión bidireccionales previos. Los autores sostienen que este enfoque permite que los modelos híbridos Mamba-atención retengan una calidad de validación competitiva mientras ofrecen mejoras sustanciales en el throughput para la inferencia de contexto largo, posicionando a los híbridos de parte-inversa por bloques como una dirección práctica para escalar los modelos de lenguaje de difusión.

El artículo se mantiene modesto respecto a la optimalidad, señalando que las reglas de transferencia de tasa de aprendizaje utilizadas para el escalamiento son candidatos sólidos basados en leyes de parametrización existentes, pero no están probadas como óptimas para los modelos de lenguaje de difusión discreta (dLLMs). Además, los autores reconocen que, si bien el objetivo por bloques mejora la generación de contexto largo, actualmente incurre en una penalización de throughput de entrenamiento (aproximadamente un 25–34% más lento) en comparación con los baselines de secuencia completa en configuraciones de hardware actuales.