SAE Feature Explorer: Descodificando la Caja Negra de los LLMs
SAE Feature Explorer (activations-vis) es una herramienta avanzada de interpretabilidad diseñada para visualizar y analizar las características latentes extraídas mediante Sparse Autoencoders (SAE) en el flujo de activación de modelos de lenguaje. Permite a los investigadores mapear conceptos semánticos específicos directamente a las activaciones de las neuronas del modelo.
🚀 Propósito y Visión Técnica
El objetivo principal es transformar las activaciones abstractas de un LLM en un “mapa de características” (starfield) interpretable. Al aplicar un SAE, el sistema descompone el estado interno del modelo en miles de características dispersas, donde cada una suele representar un concepto comprensible para el humano (ej: un país, un sentimiento, una estructura gramatical o una entidad técnica).
🛠️ Arquitectura y Stack Tecnológico
El proyecto combina un backend de procesamiento de tensores con un frontend de visualización de alta densidad:
- Backend: Desarrollado en Python, utilizando PyTorch y HuggingFace Transformers para la inferencia de modelos y la extracción de activaciones.
- Servidor de API: Implementado con Flask para servir datos precomputados y gestionar solicitudes de inferencia en tiempo real.
- Reducción de Dimensionalidad: Utiliza UMAP (
umap-learn) para proyectar las miles de dimensiones de las características del SAE en un espacio 2D o 3D navegable. - Frontend: Una interfaz web interactiva construida con Three.js (WebGL) para el renderizado de hasta 32,768 puntos de datos simultáneos con alta fluidez.
- Compatibilidad: Optimizado para modelos como Arcee Trinity Nano (MoE), pero extensible a cualquier checkpoint de TopK SAE.
✨ Características Principales
- Exploración Multidimensional: Soporte para vistas 2D y 3D, donde la posición de cada característica indica su similitud semántica con las demás.
- Análisis de Activación por Token: Panel detallado que muestra qué características se activan ante un texto específico, permitiendo un desglose quirúrgico del razonamiento del modelo.
- Estadísticas de Características: Visualización de la densidad de disparo, fuerza máxima de activación y mapas de calor de contexto para cada característica seleccionada.
- Navegación por Similitud: Algoritmo integrado para encontrar las 10 características más cercanas basadas en patrones de co-activación.
- Anotación Humana: Sistema para etiquetar y guardar interpretaciones de características en archivos
.jsonl, facilitando la creación de datasets de interpretabilidad.
⚙️ Flujo de Trabajo (Workflow)
- Fase de Precomputación: Se extraen las activaciones de una capa específica del LLM, se procesan por el SAE y se generan las coordenadas de visualización y archivos de metadatos.
- Fase de Servidor: Se levanta la API que carga los pesos del SAE y, opcionalmente, el modelo base en una GPU con soporte CUDA (requiere ~14GB VRAM para inferencia completa).
- Fase de Visualización: El usuario interactúa con el “campo de estrellas” mediante clics para inspeccionar qué “conceptos” se activan cuando la IA procesa diferentes tipos de información.
Nota técnica: SAE Feature Explorer representa una pieza fundamental en la ingeniería de seguridad de la IA (AI Safety), proporcionando una ventana transparente hacia el interior de los sistemas de inteligencia artificial más complejos.