RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale

Resumen Técnico: RADLADS

Declaración del Problema

Si bien las variantes de atención lineal (como RWKV y Mamba) ofrecen ventajas significativas sobre los transformadores de atención softmax tradicional —específicamente un tiempo de inferencia de $O(1)$ por token y un uso de memoria constante al evitar los cachés de Clave-Valor (Key-Value)—, el entrenamiento de modelos lineales a gran escala desde cero sigue siendo prohibitivamente costoso. Los modelos transformer de última generación (SoTA) suelen requerir el entrenamiento en más de diez billones de tokens, un costo inaccesible para la mayoría de las organizaciones. Los métodos existentes para convertir transformadores softmax preentrenados en modelos lineales o recurrentes (por ejemplo, T2R, SUPRA, LoLCats, MOHAWK) históricamente han requerido cantidades masivas de tokens (que oscilan entre 20 mil millones y 100 mil millones o más) o han dado como resultado un rendimiento deficiente en tareas posteriores, particularmente en benchmarks como MMLU. Además, muchos intentos de conversión dependen de arquitecturas híbridas (que retienen algo de atención softmax) o no logran igualar la calidad de los modelos maestros originales.

Metodología

Los autores presentan RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale), un protocolo diseñado para convertir modelos de decodificadores de atención softmax en modelos de atención lineal con un costo computacional mínimo. El proceso involucra tres etapas primarias y adaptaciones arquitectónicas específicas:

1. El Protocolo RADLADS

La conversión es un proceso de tres pasos que requiere solo 350–700 millones de tokens (menos del 0.005% de los datos de preentrenamiento del maestro):

• Configuración (Transferencia de Pesos de Atención): Los pesos relacionados con la atención ($W_q, W_k, W_v, W_o$) del modelo maestro se transfieren directamente a las capas de mezcla de secuencia del modelo estudiante. Otros pesos se inicializan mediante métodos estándar de preentrenamiento o se configuran para imitar el comportamiento del maestro sin efecto inmediato.
• Paso 1: Alineación del Estado Oculto de la Atención: Las capas de mezcla de secuencia del modelo estudiante se entrenan para aproximar las salidas de los estados ocultos de las correspondientes capas de atención del maestro. Esto se realiza utilizando un objetivo de distancia L2 (o MSE). Los autores descubrieron que el uso de un kernel de Atención Lineal con Compuerta (eliminando términos de decaimiento off-by-one y términos de bonificación) permite un ajuste más cercano a los estados ocultos del maestro que el RWKV-6 estándar.
• Paso 2: Destilación de Conocimiento: El modelo estudiante completo se entrena para aproximar los logits de salida del maestro utilizando la pérdida de divergencia de Kullback-Leibler (KL). Los autores hipotetizan que el conocimiento fáctico reside principalmente en los MLPs y los embeddings del maestro; por lo tanto, las tasas de aprendizaje para estos componentes se mantienen bajas o fijas para evitar el olvido catastrófico, mientras que el mezclador de secuencias se entrena de forma más agresiva.
• Paso 3: Extensión de la Longitud de Contexto (Opcional): El modelo se ajusta finamente (fine-tuning) en secuencias más largas (hasta 16k tokens) utilizando una pérdida de entropía cruzada estándar sin un modelo maestro para mejorar las capacidades de contexto largo. Alternativamente, los autores proponen el Paso 2a, que extiende la longitud de la secuencia a 4096 durante la propia fase de destilación, lo que potencialmente evita la necesidad de un Paso 3 separado.

2. Nuevas Arquitecturas

Los autores identificaron que las arquitecturas RWKV estándar no siempre eran óptimas para la conversión. Introdujeron dos nuevas variantes:

• RAD-RWKV6 ("RADFinch"): Una modificación de RWKV-6 que utiliza un kernel de Atención Lineal con Compuerta y técnicas de equilibrio de estado para mejorar la estabilidad y el ajuste.
• RAD-RWKV7 ("RADGoose"): Una modificación de RWKV-7 que elimina el mecanismo de "tokenshift" (que no ofreció beneficios en este contexto) y aplica incrustaciones posicionales rotativas (RoPE) directamente. Esta arquitectura demostró una convergencia más rápida y una pérdida de destilación menor en comparación con la arquitectura RWKV-7 sin modificar.

3. Hiperparámetros y Datos

• Dataset: Los autores se decidieron por DCLM (DataComp-LM) para todos los pasos de conversión, encontrándolo superior a FineWeb o FineWeb-Edu para convertir modelos Qwen.
• Tasas de Aprendizaje: Se utiliza un esquema de recocido de coseno (cosine annealing) en el Paso 1, comenzando alto ($10^{-3}$) para alinear los estados ocultos y terminando cerca de la tasa de aprendizaje final de preentrenamiento del maestro ($10^{-5}$). Los Pasos 2 y 3 utilizan una tasa de aprendizaje plana.

Contribuciones Clave

1. Receta de Destilación RADLADS: Un protocolo detallado paso a paso que incluye hiperparámetros específicos, conteos de tokens y elecciones de datasets que permite una conversión de alta calidad con un mínimo de datos.
2. Nuevas Arquitecturas: La introducción de RAD-RWKV6 y RAD-RWKV7, que están optimizadas para el proceso de conversión, produciendo una inferencia más rápida y un mejor alineamiento con los modelos maestros que sus predecesores sin modificar.
3. Modelos a Gran Escala: La conversión exitosa de modelos populares de código abierto Qwen2.5 en variantes de atención lineal a escalas de 7B, 32B y 72B de parámetros.
4. Lanzamiento de Código Abierto: El lanzamiento del código y los modelos convertidos (QRWKV6/7-7B/32B/72B) bajo la licencia Apache 2.0 (con las restricciones de la Licencia Qwen para el modelo 72B), permitiendo a otros replicar el proceso.

Resultados

Los modelos convertidos alcanzan un rendimiento de estado del arte entre los modelos puramente recurrentes de su tamaño:

• Eficiencia: Convertir un modelo de 72B cuesta menos de $2,000 USD y requiere solo ~700M de tokens.
• Rendimiento: En benchmarks estándar (Lambada, MMLU, ARC, etc.), los modelos RADLADS superan consistentemente otros métodos de conversión (por ejemplo, SUPRA, LoLCats, MOHAWK, ARWKV).
  • El QRWKV7-7B-Instruct logra una puntuación relativa de MMLU del 92.4% respecto a su maestro, superando significativamente otras conversiones de 7B.
  • El QRWKV6-72B-Instruct logra una puntuación relativa de MMLU del 89.9%, estableciendo un nuevo SoTA para modelos de lenguaje puramente RNN a esa escala.
• Velocidad de Inferencia: Debido al mecanismo de atención lineal, los modelos convertidos muestran aceleraciones significativas sobre los modelos maestros a medida que aumenta el contexto. Por ejemplo, con 8k tokens de entrada y 250 de salida, el modelo 32B QRWKV7 es 1.61x más rápido que su maestro Qwen3-32B; con 6k de entrada y 2k de salida, la aceleración alcanza 3.41x.

Limitaciones y Reclamaciones

El artículo reconoce modestamente varias limitaciones:

• Razonamiento y Contexto Largo: Aunque el rendimiento es sólido en benchmarks estándar, los modelos muestran limitaciones en tareas de razonamiento complejo (por ejemplo, Minerva Math) y tareas de contexto muy largo (por ejemplo, RULER), donde no mejoran tan eficazmente con la adición de tokens de salida como los modelos maestros.
• Sensibilidad de la Arquitectura: Cada nuevo diseño de arquitectura requiere pruebas meticulosas para asegurar la compatibilidad con el protocolo RADLADS. Por ejemplo, el uso de GroupNorm/LayerNorm, común en algunas variantes de atención lineal, causó inestabilidad en el entrenamiento a escalas de 14B+ y tuvo que ser reemplazado con técnicas de pre-escalado.
• Alineación de Datasets: Los autores señalan que los datasets de razonamiento pueden necesitar un alineamiento más cercano con la distribución del modelo maestro para prevenir comportamientos de bucle repetitivo, un desafío que abordaron parcialmente mezclando DCLM con OpenThoughts.

Significancia

El artículo afirma que RADLADS proporciona una vía rentable para democratizar el acceso a modelos de atención lineal a gran escala. Al reducir el requisito de entrenamiento de billones de tokens a cientos de millones, permite a investigadores y organizaciones más pequeñas probar, entrenar y desplegar nuevas variedades de arquitecturas RNN a escala sin los costos extremos asociados con el preentrenamiento. Los autores posicionan esto como una herramienta para acelerar el desarrollo de la próxima generación de variantes de atención de estado comprimido.