Resumen Técnico: WriteSAE para la Descomposición de Estados Recurrentes

Enunciado del Problema

Los modelos de lenguaje recurrentes híbridos y de espacio de estados (por ejemplo, Gated DeltaNet, Mamba-2, RWKV-7, Qwen3.5) utilizan una caché matricial para la propagación de estados, donde cada token escribe un producto externo de rango 1 ($k_t v_t^\top$) en una matriz $d_k \times d_v$. Los Autoencoders Dispersos (SAE) existentes operan sobre los flujos residuales, leyendo el estado después de la emisión. En consecuencia, los SAE estándar no pueden acceder ni descomponer directamente el mecanismo de escritura nativo, el cual ocurre aguas arriba en la recurrencia. Además, dado que el estado se actualiza mediante productos externos de rango 1, un átomo decodificador de SAE basado en vectores estándar no puede reemplazar una única escritura de caché sin romper la correspondencia estructural requerida para la intervención causal (parcheo de caché).

Metodología

El artículo introduce WriteSAE, el primer autoencoder disperso diseñado para descomponer y editar la escritura de caché matricial de los modelos recurrentes.

Arquitectura y Entrenamiento

• Ajuste de Forma Nativa: Los átomos decodificadores de WriteSAE están restringidos a ser productos externos de rango 1 ($v_i w_i^\top$), coincidiendo con el primitivo de escritura nativo ($k_t v_t^\top$) de la arquitectura anfitriona. Esto permite que un solo átomo reemplace una actualización específica de una ranura de caché.
• Codificador Bilinear: Para codificar eficientemente el estado matricial $S_t$, el modelo utiliza un codificador bilineal ($a_i = v_i^\top S_t w_i$) en lugar de un codificador denso aplanado. Esto reduce la cantidad de parámetros por átomo de $d_k d_v$ (16.384 para $128 \times 128$) a $d_k + d_v$ (256), una reducción de 64 veces.
• Objetivo de Entrenamiento: El modelo se entrena sobre estados centrados en la media ($x = \text{vec}(S_t - M)$) utilizando una restricción de dispersión TopK y una pérdida auxiliar para revivir características muertas. El corpus de entrenamiento consiste en secuencias de OpenWebText procesadas a través de Qwen3.5-0.8B (y otros modelos) en capas y cabezas específicas.

Clasificación de Características

El diccionario entrenado se divide en dos clases basadas en la similitud del coseno mediana entre la dirección de escritura del átomo y la dirección de escritura nativa:

1. Registros: Átomos con alta alineación de coseno con las escrituras nativas (recuperables desde la caché).
2. Paquetes: Átomos con direcciones de escritura dispersas a través de la caché.

Pruebas de Intervención Causal

El artículo propone tres mecanismos principales de validación:

1. Sustitución de Ranura de Caché: Reemplazar la escritura nativa ($k_t v_t^\top$) con un átomo SAE de norma de Frobenius coincidente ($\hat{S}_t = S_t - \Delta_{nat} + \Delta_{atom}$) y medir la divergencia KL aguas abajo.
2. Desplazamiento de Logits de Forma Cerrada: Derivar una expresión de tres factores para predecir los desplazamientos de logits por token: $\Delta \ell \approx G_{t_0 \to t} \langle w_i, q_t \rangle \langle v_i, W_U[\text{tok}] \rangle$, donde $G$ es el producto de la puerta, $q_t$ es la consulta de lectura y $W_U$ es la fila de desincrustación.
3. Intervenciones Directas en Caché: Borrar o instalar átomos en posiciones específicas para observar cambios en las probabilidades de los tokens y el comportamiento de generación.

Contribuciones Clave

1. Diccionario Ajustado a la Arquitectura: Un diccionario donde los átomos son productos externos de rango 1, habilitando la sustitución directa de escrituras de caché nativas, una capacidad previamente no disponible para los SAE de flujo residual.
2. Validación de Sustitución: Demostración de que los átomos aprendidos pueden reemplazar escrituras nativas con alta fidelidad. En Qwen3.5-0.8B (Capa 9, Cabeza 4), los átomos superaron la ablación de norma coincidente en el 92.4% de 4.851 disparos.
3. Predicción de Forma Cerrada: Una fórmula de tres factores sin parámetros que predice los desplazamientos de logits por disparo con una $R^2$ mediana de 0.98 a través de 200 celdas de átomo-por-épsilon.
4. Transferencia Trans-Arquitectónica: Validación de que el éxito de la sustitución y la partición de características (Registro vs. Paquete) se mantienen a través de diferentes arquitecturas de recurrencia matricial (Gated DeltaNet, RWKV-7, Mamba-2), aunque las tasas de éxito varían según el "rango de escritura" del sustrato.

Resultados Experimentales

Rendimiento de Sustitución

• Qwen3.5-0.8B (Gated DeltaNet): En L9 H4, los átomos superaron la ablación en el 92.4% de los disparos. La prueba de población sobre 87 átomos se mantiene en el 89.8%. La ordenación estricta $KL_{atom} < KL_{ablate} < KL_{random}$ se cumple en el 89.5% de los disparos.
• Mamba-2-370M: La sustitución tiene éxito en el 88.1% sobre 2.500 disparos, a pesar de que Mamba-2 utiliza una actualización de estado diagonal en lugar de un producto externo completo de rango 1.
• RWKV-7: Muestra un rendimiento intermedio, consistente con su regla de escritura de rango 2.

Predicción de Desplazamiento de Logits

La expresión de forma cerrada rastrea con precisión los efectos medidos. En L9 H4, la $R^2$ mediana entre los desplazamientos de logits predichos y medidos es de 0.98. Sin embargo, esta forma cerrada falla ($R^2$ negativa) en Mamba-2 y Qwen3.5-4B, lo que indica que el factor de puerta $G$ es específico del sustrato.

Intervenciones Conductuales

• Borrado: Borrar un átomo de registro específico (F412) en sus posiciones naturales de disparo redujo el log-probabilidad del token objetivo ("espacio") en una mediana de -0.116 nats ($p=1.07 \times 10^{-6}$).
• Dirigir la Generación: La instalación sostenida de una dirección de forma cerrada en tres posiciones consecutivas aumentó la tasa de "objetivo-en-continuación" para tokens de rango medio del 33.3% al 100% bajo decodificación codiciosa.
• Amplificación a Nivel de Pasaje: Amplificar características correlacionadas con los límites en un modelo de 4B redujo la generación de saltos de línea en un 33% (de 16.8 a 11.2 por cada 400 tokens) a una dosis de 5x, demostrando el control sobre la estructura del documento.

Alcance Arquitectónico

El artículo establece una jerarquía de efectividad del "rango de escritura":

• Gated DeltaNet (Rango-1): Mayor separación de coseno de registro (0.262) y fidelidad de sustitución.
• RWKV-7 (Rango-2): Separación intermedia (0.180).
• Mamba-2 (Diagonal/Scalar): Menor separación (0.0575), aunque aún logra un éxito de sustitución de ~88%.

Significado y Afirmaciones

El artículo afirma proporcionar el primer diccionario causal para el sitio de escritura de recurrencia matricial, superando el análisis de flujos residuales. Demuestra que:

1. El Rango de Escritura es Determinante: La capacidad de descomponer el estado en "registros" interpretables se correlaciona con el rango de la regla de escritura nativa (Rango-1 > Rango-2 > Diagonal).
2. La Sustitución Causal es Posible: Los átomos aprendidos pueden reemplazar escrituras nativas con una divergencia KL aguas abajo mínima, validando el concepto de "registro" en modelos recurrentes.
3. Interpretabilidad de Forma Cerrada: Para Gated DeltaNet, el efecto de una característica sobre los logits puede predecirse analíticamente sin ajustar parámetros, vinculando el diccionario directamente con los mecanismos de puerta y lectura del modelo.

Los autores notan modestamente las limitaciones: el coeficiente de forma cerrada no se transfiere entre sustratos (por ejemplo, falla en Mamba-2), y la identidad por átomo es específica de la semilla (aunque las propiedades a nivel de clase son estables). El trabajo establece un nuevo punto de referencia para la interpretabilidad mecánica en modelos de espacio de estados, mostrando que el primitivo de "escritura" es un objetivo viable para el aprendizaje de diccionarios y la edición causal.