SAE en Gemma-2-2B: la proporción de rasgos coincidentes sube del 2% al 29% con 50M de tokens

Los sparse autoencoders (SAE) se han convertido en una de las principales herramientas de interpretabilidad de los modelos de lenguaje, y por lo general se entrenan con cientos de millones o miles de millones de tokens de activaciones. Casi nadie ha comprobado que esa escala sea realmente necesaria. Los autores de un nuevo trabajo publicado en Hugging Face Daily Papers lo han hecho: tomaron cuatro arquitecturas de SAE, tres modelos de lenguaje y presupuestos que van de 5M de tokens hasta los 500M de referencia.

La conclusión principal: la respuesta depende de qué se mida exactamente. En Gemma-2-2B con BatchTopK, la proporción de rasgos que coinciden de forma estricta sube del 2% con 5M de tokens al 29% con 50M.

Dos medidas de convergencia

Los investigadores separaron dos conceptos. El primero es el comportamiento funcional: la calidad de la reconstrucción y los resultados en los benchmarks estándar de SAE. El segundo es la reproducibilidad «uno a uno»: si un SAE entrenado de forma independiente encuentra los mismos rasgos que el de referencia.

La diferencia resultó notable. La reconstrucción suele alcanzar el nivel de la referencia completa con presupuestos bastante menores. La reproducibilidad de los rasgos, en cambio, sigue mejorando a medida que aumentan los datos. Un SAE puede parecer «listo» según las métricas, y aun así el diccionario de rasgos que contiene sigue siendo inestable.

Observaciones clave del trabajo:

• Escala del experimento: 4 arquitecturas de SAE, 3 modelos de lenguaje, presupuestos de 5M a 500M de tokens.
• Gemma-2-2B, BatchTopK: 2% de rasgos con coincidencia estricta con 5M de tokens y 29% con 50M.
• La necesidad de datos depende mucho de la métrica: la reconstrucción se satura antes; la reproducibilidad, no.
• El muestreo aleatorio conserva mejor la distribución original de las activaciones y suele ganar en reconstrucción.
• El muestreo por diversidad suele mejorar el resultado de AutoInterp.

Cómo elegir los datos

Importa no solo el volumen, sino también el método de selección. El muestreo aleatorio reproduce con más fidelidad la distribución de las activaciones, por lo que suele ser preferible en tareas que exigen una buena reconstrucción. El muestreo basado en la diversidad cubre más comportamientos poco frecuentes del modelo, y eso con frecuencia mejora las puntuaciones de AutoInterp, es decir, de la interpretación automática de rasgos. No existe un único muestreador «mejor»: se elige según el objetivo.

La sección teórica del trabajo explica por qué las dos métricas divergen. La consistencia estadística del entrenamiento con un subconjunto es una condición. La recuperación de rasgos uno a uno exige condiciones de identificabilidad y de optimización más estrictas. Dicho de otro modo, aunque un subconjunto ofrezca una buena aproximación al funcional objetivo, eso no garantiza que distintas ejecuciones lleguen al mismo conjunto de rasgos.

Qué hacer en la práctica

La conclusión para los equipos que entrenan SAE es bastante pragmática. Si el objetivo es obtener rápido un SAE funcional con buena reconstrucción, un presupuesto de cientos de millones de tokens resulta, a la vista de los resultados, a menudo excesivo. Eso abarata los experimentos, sobre todo al probar distintas arquitecturas e hiperparámetros.

Si, en cambio, se publican afirmaciones sobre rasgos concretos, por ejemplo «este latente codifica tal concepto», las métricas de reconstrucción no bastan. Hace falta una comprobación explícita de reproducibilidad: varias ejecuciones independientes y la comparación de rasgos entre ellas. Una proporción baja de coincidencias con un presupuesto pequeño, como ese mismo 2% con 5M de tokens, advierte de que algunos rasgos pueden ser un artefacto de un entrenamiento concreto.

Por último, conviene elegir el muestreador según el objetivo. Para la reconstrucción, empiece por el muestreo aleatorio; para la interpretabilidad medida con AutoInterp, pruebe la selección por diversidad. Los autores subrayan que, con objetivos y presupuestos bien alineados, la investigación con SAE puede ser más barata y fiable. Una salvedad: las cifras corresponden a modelos y arquitecturas concretos, y trasladarlas a otras escalas requiere una verificación aparte.