Resumen Técnico: Difusión de Bloques Triplet RWKV (B3D-RWKV)

Enunciado del Problema

Los modelos de lenguaje Transformer causales enfrentan dos limitaciones estructurales primarias: la decodificación estrictamente secuencial, que impide la paralelización, y los costos de atención cuadráticos por paso, que hacen que la inferencia en contextos largos sea computacionalmente costosa. Si bien los modelos causales de tiempo lineal (por ejemplo, RWKV) y los modelos de difusión discretos abordan estas debilidades de forma independiente, su integración ha sido inherentemente inconsistente. Los modelos de difusión requieren atención bidireccional para desruidar bloques de tokens en paralelo, mientras que los modelos causales son estrictamente unidireccionales. Los intentos existentes de combinarlos a menudo requieren modificaciones arquitectónicas que comprometen la eficiencia del núcleo de tiempo lineal o hacen necesario el entrenamiento desde cero.

Metodología

Los autores proponen B3D-RWKV, una variante estilo difusión de RWKV que unifica la eficiencia de inferencia $O(L)$ con la difusión discreta bidireccional paralela sin alterar el núcleo causal subyacente. La innovación central es la Disposición de Bloque Triplet, una estrategia de entrenamiento e inferencia diseñada para simular el acceso bidireccional dentro de una arquitectura estrictamente causal.

1. Disposición de Bloque Triplet (Entrenamiento)

Para habilitar la difusión en un modelo causal, cada bloque lógico de generación de tamaño $B$ se expande en tres bloques físicos contiguos dentro de la muestra de entrenamiento:

• $b_1$ (Copia Enmascarada): Una versión enmascarada del bloque lógico.
• $b_2$ (Copia Enmascarada Sujeita a Pérdida): Una copia enmascarada idéntica donde se calcula la pérdida de desruido.
• $b_3$ (Copia Limpia de Verdad Terrena): Una versión sin máscara que actualiza el estado recurrente para el bloque subsiguiente.

Acceso Pseudo-Bidireccional:
Dado que el núcleo causal procesa los tokens estrictamente de izquierda a derecha, el estado oculto que llega a cualquier posición enmascarada en $b_2$ ya ha procesado la totalidad de $b_1$. Como $b_1$ y $b_2$ comparten el mismo patrón de máscara, los tokens sin máscara en $b_1$ son efectivamente "absorbidos" en el estado oculto antes de que el modelo intente predecir los tokens enmascarados en $b_2$. Esto otorga al modelo acceso pseudo-bidireccional a su propio contexto sin máscara dentro del bloque lógico, satisfaciendo los requisitos bidireccionales de la difusión mientras mantiene la naturaleza unidireccional del núcleo RNN/SSM.

Objetivo de Entrenamiento:
El modelo se entrena utilizando un objetivo de predicción de máscara. La pérdida consiste en:

• Pérdida de Entropía Cruzada ($L_{CE}$): Calculada en posiciones supervisadas donde los tokens están enmascarados.
• Pérdida Paralela Consciente de la Confianza ($L_{CAP}$): Siguiendo a LLaDA-2.0, este término afina la distribución de probabilidad en posiciones ya predichas correctamente, incentivando al muestreador de inferencia a comprometer más tokens por paso.

2. Inferencia: Desruido Iterativo por Bloques

Durante la inferencia, el modelo genera un bloque lógico a la vez. El proceso implica:

1. Inicializar un bloque con todos los tokens enmascarados.
2. Ejecutar pasos iterativos de desruido donde el modelo predice probabilidades top-1 para posiciones enmascaradas.
3. Comprometer tokens cuya probabilidad top-1 excede un umbral $\tau$.
4. Un mecanismo de respaldo asegura el progreso comprometiendo los $k_{min}$ tokens más confiables si menos de $k_{min}$ superan el umbral.
5. Una vez que todas las posiciones están comprometidas, el bloque limpio se agrega al prefijo y comienza el siguiente bloque lógico.

Contribuciones Clave

1. B3D-RWKV-7.2B: El primer modelo de lenguaje RNN de tiempo lineal estilo difusión entrenado a escala de 7B con un objetivo de predicción de máscara. Integra la selección paralela de tokens en el núcleo RWKV-7 sin modificar sus parámetros originales.
2. Marco de Bloque Triplet: Una receta de entrenamiento universal que convierte núcleos estrictamente causales (que satisfacen la causalidad estricta y las propiedades de propagación hacia adelante del estado) en modelos de difusión por bloques. Este enfoque evita la complejidad arquitectónica de modificar el núcleo en un bloque bidireccional.
3. Eficiencia y Rendimiento: El método logra mejoras significativas en el rendimiento de transmisión (throughput) en comparación con la decodificación causal estándar, manteniendo capacidades de razonamiento competitivas.

Resultados Experimentales

El modelo fue evaluado en una suite de 8 tareas que incluye conocimiento general (MMLU, ARC, PIQA, RACE) y tareas de razonamiento/matemáticas (GSM8K, MATH, GPQA).

• Precisión: B3D-RWKV-7.2B logra una precisión comparable a la línea base RWKV-7 y otros modelos de difusión de escala similar (por ejemplo, LLaDA-8B, Dream-7B). Destaca notablemente sobre las líneas base en tareas de razonamiento como ARC-Challenge y RACE, atribuido a la percepción pseudo-bidireccional que mejora el razonamiento.
• Rendimiento de Transmisión: El modelo demuestra una aceleración promedio de 1.6× en el rendimiento de decodificación en comparación con la línea base RWKV-7 a lo largo de longitudes de contexto que van desde 1K hasta 512K tokens. Con parámetros de muestreo ajustados, es posible lograr una aceleración de 2.02× con un ligero compromiso de calidad.
• Razonamiento Matemático: El artículo señala una ligera caída en el rendimiento en tareas matemáticas complejas (por ejemplo, MATH) en comparación con la línea base causal. Los autores atribuyen esto a la naturaleza de la decodificación paralela, que lucha con la verificación altamente estructurada y paso a paso requerida para problemas de matemáticas simbólicas donde no se otorga crédito parcial.

Significado y Afirmaciones

El artículo afirma ofrecer una vía eficiente para transformar modelos de lenguaje causales preentrenados en modelos de lenguaje de difusión sin cambios arquitectónicos. Al aprovechar la disposición de bloque triplet, B3D-RWKV unifica con éxito la eficiencia de inferencia $O(L)$ de los modelos de tiempo lineal con las capacidades de decodificación paralela de la difusión.

Los autores posicionan el trabajo como una solución estructural a la incompatibilidad entre los paradigmas causal y de difusión. Reconocen limitaciones, incluida la validación actual en un solo núcleo (RWKV-7), el costo de longitud de secuencia física 3× durante el entrenamiento y el uso de un conjunto de datos SFT relativamente pequeño sin alineación RL, lo cual podría contribuir a regresiones de precisión en tareas específicas. Sin embargo, la contribución central sigue siendo la demostración de que la generación paralela de difusión de alto rendimiento es alcanzable en núcleos estrictamente causales y de tiempo lineal a través de la disposición de datos en lugar de la reingeniería arquitectónica.