HN abierto: Estudio de interpretabilidad de Llama 3.2 con Sparse Autoencoders

Resumen del proyecto

• Los modelos de lenguaje grandes (LLM) modernos codifican conceptos superponiendo múltiples características en una misma neurona, de modo que la activación de cada neurona puede tener varios significados interpretables según la activación de otras neuronas. A esto se le llama superposición (superposition).
• Los Sparse Autoencoders (SAE) se insertan en un LLM entrenado para proyectar activaciones en un espacio latente disperso muy grande, separando representaciones superpuestas y transformándolas en características interpretables.
• Este proyecto busca reproducir investigaciones realizadas con éxito por Anthropic, OpenAI y Google DeepMind para extraer características interpretables.
• Proporciona un pipeline completo para crear Sparse Autoencoders funcionales e interpretables para el modelo Llama 3.2-3B.

Funcionalidades clave

• Ofrece un pipeline completo de extremo a extremo, escrito en PyTorch, desde la captura de activaciones hasta el entrenamiento de SAE, la interpretación de características y su validación.
• Captura activaciones residuales de modelos de lenguaje grandes para usarlas como dataset de entrenamiento para SAE.
• Preprocesa eficientemente los datos de entrenamiento y soporta entrenamiento distribuido a gran escala con múltiples GPU.
• Implementa pérdida auxiliar durante el entrenamiento del SAE para evitar variables latentes muertas y estabilizar la dinámica de entrenamiento.
• Proporciona logging integral, visualización y checkpoints del entrenamiento de SAE mediante Weights & Biases.
• Incluye herramientas de análisis de interpretabilidad para estudiar el significado de las características aprendidas.
• Permite uso general y verificación de resultados sin dependencias externas mediante una implementación pura en PyTorch de Llama 3.1/3.2.
• Valida el impacto de los SAE en el comportamiento del modelo mediante tareas de completado de texto y chat, y permite ajustar las características semánticas extraídas.

Recursos publicados

1. Dataset de oraciones de OpenWebText:

   • Una versión personalizada del dataset OpenWebText usada para la captura de activaciones.
   • Conserva el texto original y almacena oraciones individuales en formato Parquet para acceso rápido.
   • Segmentación de oraciones usando el tokenizador "Punkt" de NLTK 3.9.1.

2. Activaciones capturadas de Llama 3.2-3B:

   • 25 millones de activaciones residuales de la capa 23 de Llama 3.2-3B.
   • 4 TB de datos en bruto comprimidos a 3.2 TB y divididos en 100 archivos.

3. Logs de entrenamiento de SAE:

   • Logs de visualización de métricas de entrenamiento, validación y depuración mediante Weights & Biases.
   • Incluye 10 épocas y 10,000 pasos registrados.

4. Modelo SAE entrenado de 65,536 latentes:

   • El modelo SAE final entrenado tras 10 épocas.

Estructura del proyecto

1. Captura de datos

• capture_activations.py: captura de activaciones residuales del LLM.
• openwebtext_sentences_dataset.py: dataset personalizado para procesamiento a nivel de oración.

2. Entrenamiento de SAE

• sae.py: implementación central del modelo SAE.
• sae_preprocessing.py: preprocesamiento de datos para entrenamiento de SAE.
• sae_training.py: implementación de entrenamiento distribuido de SAE.

3. Interpretabilidad

• capture_top_activating_sentences.py: identifica oraciones que maximizan la activación de características.
• interpret_top_sentences_send_batches.py: genera y envía lotes para interpretación.
• interpret_top_sentences_retrieve_batches.py: recopila resultados de interpretación.
• interpret_top_sentences_parse_responses.py: analiza los resultados de interpretación.

4. Validación y pruebas

• llama_3_inference.py: implementación central de inferencia.
• llama_3_inference_text_completion_test.py: prueba de completado de texto.
• llama_3_inference_chat_completion_test.py: prueba de completado de chat.
• llama_3_inference_text_completion_gradio.py: interfaz Gradio para pruebas interactivas.

[GN⁺]

Opiniones en Hacker News

• La interpretabilidad mecánica de los LLM resuelve el problema de que los modelos generen respuestas convincentes cuando se les pide explicarse a sí mismos. Cuanto más potente es el modelo, más convincente puede ser al justificar una "mentira", por lo que podría obtener peores resultados en pruebas de autodetección. El objetivo es la coherencia, no la verdad

  • La verdadera explicabilidad consiste en mostrar el proceso causal de razonamiento por el que pasa el modelo al generar una respuesta. Esto es similar a los pensamientos que aparecen brevemente en varias direcciones cuando una persona piensa

• En la investigación sobre Sparse Autoencoders (SAEs), se observó que la cota inferior de la curva de pérdida escala con una ley de potencias. Con una pérdida auxiliar se pudo resolver por completo el problema de las variables latentes muertas, y se observó un patrón suave de onda sinusoidal durante las iteraciones de entrenamiento

• Se plantea una pregunta sobre la interpretabilidad mecánica: existe preocupación por la posibilidad de que una IA futura, al supervisar su propio entrenamiento, pueda crear modelos que engañen a los observadores de interpretabilidad mecánica aprovechando la ambigüedad

• Tras leer una publicación de blog sobre la dificultad de evaluar SAEs, surge la duda de cómo resolvieron este problema. Quieren encontrar en el repositorio la parte donde pueda entenderse el enfoque

• Creen que este trabajo podría tener un efecto positivo en la alineación, pero aún no han revisado los detalles. Se preguntan cuánto habría que pagar para compensar el tiempo, el costo y el riesgo

• Agradecen que se haya dedicado tanto tiempo a la documentación

• Es un trabajo muy genial y se preguntan si hay planes de integrarlo con SAELens