Stuck on "A": Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization of a 0.6B Language Model

Resumen Técnico: Diagnóstico y Reparación de la Lesión de Interfaz en la Linealización de Atención-a-KDA

Declaración del Problema
El artículo investiga la naturaleza "con pérdida" (lossy) de la conversión de modelos de lenguaje con atención completa preentrenados a arquitecturas de atención lineal, centrándose específicamente en el mecanismo Kimi Delta Attention (KDA). Si bien la atención lineal ofrece ventajas estructurales para la inferencia de contexto largo y el despliegue en dispositivos periféricos (estado recurrente $O(1)$ frente a un caché KV $O(n)$), el proceso de conversión a menudo falla en preservar el rendimiento de las tareas, a pesar de que las métricas estándar sugieren éxito.

El autor se centra en la conversión de Qwen3-0.6B-Base (28 capas) en una arquitectura híbrida donde 21 capas son reemplazadas por KDA y 7 capas retienen la Atención de Consulta Agrupada (GQA). Esto se ejecuta bajo estrictas limitaciones de consumo: una única GPU de 32 GB y un presupuesto de entrenamiento de decenas de millones de tokens (en lugar de miles de millones).

La pregunta diagnóstica central es: ¿Qué es exactamente lo que la conversión rompe? La evaluación estándar revela una paradoja: tras la alineación capa por capa y la destilación KL de extremo a extremo, el modelo estudiante alcanza una perplejidad cercana a la del profesor (brecha de Entropía Cruzada de +0.128 nats), pero colapsa en la precisión de opción múltiple (manteniéndose cerca del azar con un ~25–29% en C-Eval, comparado con el 50.6% del profesor).

Metodología
El autor propone un pipeline multi-etapa para diagnosticar y reparar esta brecha de rendimiento:

1. Arquitectura Híbrida y Cirugía:

   • Diseño (Layout): Las capas 3, 7, 11, 15, 19, 23 y 27 retienen la GQA nativa; las 21 capas restantes se convierten a KDA (16 cabezales, dimensión de cabezal 128).
   • Mecánica de KDA: Implementa una regla delta con compuerta por canal con una compuerta de olvido ($\alpha_t$) y una compuerta de escritura ($\beta_t$).
   • Inicialización: Un estudio de ablación crítico identificó que los trasplantes simples de proyecciones fallaban. La configuración exitosa (g6) utiliza escalas de decaimiento y ajustes de sesgo específicos para estabilizar el estado recurrente.

2. Etapas de Entrenamiento:

   • Etapa 2 (Alineación Capa por Capa): Alinea los estados ocultos de KDA con los estados ocultos del profesor congelado mediante MSE. Esto redujo la Entropía Cruzada (CE) de 9.48 a 4.13, pero dejó la precisión de C-Eval en 25.5%.
   • Etapa 3a (Destilación KL de Extremo a Extremo): Entrenada en texto de preentrenamiento utilizando la divergencia KL hacia adelante entre las distribuciones del profesor y el estudiante. Esto cerró la brecha de CE a +0.128 nats, pero solo elevó la precisión de C-Eval a 28.8%.
   • Etapa 3b (Reparación Dirigida al Formato): Una etapa de KL de 1,000 pasos de "solo completitud" utilizando preguntas de opción múltiple (MCQs), poesía y traducción de QA verificadas por el profesor. Crucialmente, esta etapa se centró en el formato de la respuesta en lugar de solo en la generación de texto general.

3. El Diagnóstico de Cuatro Permutaciones:
   Para ir más allá de la perplejidad, el autor diseñó un protocolo de diagnóstico utilizando 161 MCQs limpios (cero solapamiento con benchmarks). Cada pregunta fue evaluada bajo cuatro rotaciones cíclicas de las opciones de respuesta (A, B, C, D).

   • Métrica: La "adherencia" (stickiness) mide qué tan seguido el modelo predice la misma etiqueta (por ejemplo, "A") independientemente de la posición del contenido de la opción.
   • Hallazgo: El modelo de la Etapa 3a exhibió una severa lesión de interfaz: predijo "A" el 81% de las veces y mantuvo la misma etiqueta a través de las cuatro rotaciones para 106/161 preguntas. El modelo había perdido la capacidad de mapear el conocimiento a las etiquetas de las opciones, ignorando efectivamente el contenido de las elecciones.

4. Alineación de Persona:
   Tras la reparación, el modelo pasó por un Ajuste Fino Supervisado (SFT) con una persona inyectada ("Qingyi") y una ronda de Optimización de Preferencia Directa (DPO) on-policy para asegurar que la conversión no degradara su identidad o sus capacidades conversacionales.

Resultados Clave

• "La Mentira de la Perplejidad": El estudio confirma que la alineación de estados ocultos y la baja divergencia KL no garantizan la capacidad de tarea. El modelo puede imitar las estadísticas de distribución del profesor mientras falla en utilizar ese conocimiento para tareas específicas.
• Diagnóstico de la Lesión de Interfaz: El test de cuatro permutaciones reveló que el modelo no sufría de una pérdida de conocimiento uniforme (azar), sino de adherencia de etiqueta. Fallaba al rastrear el contenido de la respuesta, recurriendo por defecto a una etiqueta de opción específica.
• Eficacia de la Reparación: La reparación de la Etapa 3b dirigida al formato:
  • Aumentó la precisión de C-Eval de 28.83% a 41.31%.
  • Redujo la participación de predicción de "A" de 81.06% a 58.5%.
  • Redujo a la mitad la métrica de "misma etiqueta en 4 rotaciones" (106/161 $\to$ 47/161).
  • Convirtió el margen de la opción correcta de negativo (−0.117 nats) a positivo (+0.233 nats).
  • Compromiso (Trade-off): Esta reparación tuvo un costo medible en el ajuste de la distribución general (el CE/KL válido aumentó ligeramente), lo que sugiere que el modelo se especializó en el formato de MCQ.
• Desempeño Final: Tras la alineación de la persona, el modelo final alcanzó un 41.83% en C-Eval, permaneciendo dentro de la banda de ruido de la etapa de reparación post-SFT, lo que indica que no hubo olvido catastrófico.

Lecciones de Ingeniería y Modos de Fallo
El documento registra varios errores críticos de ingeniería encontrados durante el proceso:

• Fallo del Modo Maestro FP32: El uso de bf16 tanto para los parámetros como para los momentos del optimizador Adam causó que las actualizaciones pequeñas de los pesos fueran "silenciosamente absorbidas" (actualizaciones sub-ULP), resultando en pesos congelados a pesar de tener gradientes saludables. La solución fue usar pesos maestros en FP32.
• Paridad de Caché: Las nuevas capas KDA inicialmente ignoraban las actualizaciones incrementales del caché, causando que la generación colapsara en basura desde el segundo token. Una corrección específica en la implementación del caché resolvió esto sin necesidad de reentrenamiento.
• Trasplantes Falsos: Los intentos tempranos de trasplantar las proyecciones Q/K/V/O del profesor fallaron funcionalmente; el enfoque exitoso abandonó el trasplante de proyecciones por completo.

Significancia y Reivindicaciones
El artículo reclama cuatro contribuciones primarias:

1. Herramienta de Diagnóstico: Un protocolo de cuatro permutaciones que separa la "lesión de interfaz" (incapacidad de mapear el conocimiento al formato de salida) de la "pérdida de conocimiento" (olvido uniforme), aplicable a cualquier pipeline de linealización o compresión.
2. Receta Reproducible: Una receta completa y con presupuesto limitado para convertir Qwen3 a un híbrido KDA, documentando modos de fallo específicos (ablación de compuertas, errores de caché, absorción de actualizaciones) que suelen omitirse en los registros públicos.
3. Supervivencia de la Persona: Evidencia de que la alineación de la persona (SFT + DPO) puede sobrevivir a la conversión de arquitectura esencialmente "gratis", sin requerir un reentrenamiento masivo ni causar olvido catastrófico.
4. Lanzamiento Abierto: Publicación completa de código, pesos, recetas y el rastro de auditoría de fallos.

El autor concluye que, si bien convertir un modelo de 0.6B a un híbrido KDA con un presupuesto de consumo es factible, los conjuntos de métricas estándar (perplejidad, KL, MSE de estado oculto) reportan erróneamente el daño de forma sistemática. La barrera principal para el éxito no es necesariamente la pérdida de conocimiento interno, sino la degradación de la interfaz que mapea ese conocimiento a formatos de salida específicos (como las etiquetas de opción múltiple), lo cual puede repararse con supervisión de formato dirigida.