Resumen Técnico: Bergson: Una biblioteca de código abierto para la atribución de datos

Planteamiento del problema

La atribución de datos busca explicar el comportamiento de un modelo determinando el efecto causal de elementos individuales de entrenamiento sobre salidas específicas del modelo. Si bien este campo ofrece aplicaciones en depuración, curación de conjuntos de datos y comprensión de la memorización o el razonamiento, escalar estas técnicas a los Modelos de Lenguaje Extensos (LLM) modernos presenta desafíos de ingeniería significativos. La aplicación ingenua de métodos de atribución de datos a menudo resulta en requisitos de memoria poco realistas, particularmente cuando se trata de conjuntos de datos de preentrenamiento a gran escala o modelos que requieren soporte multinodo. Además, muchos métodos de vanguardia, como aquellos que logran altas correlaciones de Leave-One-Out (LOO), carecen de implementaciones de código abierto, lo que dificulta la reproducibilidad y el progreso.

Metodología y arquitectura de la biblioteca

Bergson se introduce como una biblioteca de código abierto diseñada para permitir un progreso más rápido en la atribución de datos mediante la provisión de una arquitectura composable y escalable. La biblioteca soporta una amplia gama de técnicas, incluyendo la diferenciación a través de trayectorias completas de entrenamiento, la diferenciación desenrollada aproximada y las funciones de influencia.

Componentes principales

• Pipeline Componible: Bergson utiliza una arquitectura de pipeline de datos (Figura 1) donde componentes como las aproximaciones de la Hessiana, normalizadores, constructores de índices y evaluadores de atribución son interoperables. Los flujos de trabajo pueden definirse mediante una Interfaz de Línea de Comandos (CLI) o programáticamente utilizando archivos de configuración YAML.
• Almacenamiento e indexación de gradientes: La biblioteca soporta almacenes de gradientes en disco indexados por módulo (incluyendo cabezales de atención divididos) y soporta la recolección por token o por secuencia. Para manejar índices grandes, Bergson se integra con FAISS para búsquedas de vecino más cercano aproximado (ANN), mientras que también soporta búsquedas exactas de top-k.
• Entrenamiento distribuido: Bergson implementa paralelismo replicado y fragmentado (DDP, FSDP) para escalar la atribución a modelos con hasta 405 mil millones de parámetros. Soporta operaciones multinodo donde los grupos FSDP pueden procesar independientemente fragmentos del conjunto de datos.
• Integración con PEFT: La integración nativa con la biblioteca PEFT permite una atribución de datos restringida a los parámetros del adaptador (por ejemplo, LoRA), reduando la huella de gradiente por ejemplo en órdenes de magnitud.
• Objetivos no diferenciables: Para tareas evaluadas mediante LLM-as-a-judge, Bergson soporta la atribución a través de la pérdida de Optimización de Política Relativa de Grupo (GRPO), permitiendo la atribución de comportamientos que no son capturados por las pérdidas estándar de entropía cruzada.

Algoritmos implementados

Bergson proporciona las primeras implementaciones de código abierto para tres métodos líderes:

1. MAGIC: Un algoritmo para la diferenciación exacta a través de bucles de entrenamiento desenrollados mediante checkpointing. Logra $O(N \log N)$ de computación y $O(\log N)$ de espacio en disco (o un nuevo compromiso de $O(N)$ de cómputo/$O(\sqrt{N})$ de disco).
2. SOURCE: Un método para la diferenciación desenrollada aproximada que divide la trayectoria de entrenamiento en segmentos, tratando las Hessianas y los gradientes como constantes dentro de cada segmento.
3. TrackStar: Un método de aproximación escalable para funciones de influencia que proyecta los gradientes por ejemplo en un espacio de baja dimensión utilizando proyecciones aleatorias y utiliza aproximaciones de la Hessiana de autocorrelación.

La biblioteca también soporta variantes estándar de funciones de influencia (K-FAC, EK-FAC, TK-FAC/Shampoo) y estrategias de normalización de gradiente (Adam, Adafactor, normalización unitaria).

Contribuciones clave

1. Implementación de código abierto: Bergson introduce las primeras implementaciones públicas de MAGIC, SOURCE y TrackStar, llenando un vacío en las herramientas para métodos de vanguardia en atribución de datos.
2. Escalabilidad: La biblioteca está diseñada para manejar modelos de miles de millones de parámetros y conjuntos de datos masivos, soportando almacenes de gradientes en disco, integración con FAISS y entrenamiento distribuido multinodo.
3. Guía de selección de métodos: El artículo proporciona una guía para que los profesionales seleccionen métodos basados en restricciones de recursos, volumen de consultas y la necesidad de precisión (por ejemplo, recomendando la diferenciación desenrollada para filtrado de seguridad crítica y los índices comprimidos para depuración de alta frecuencia de consultas).
4. el Marco de evaluación: Bergson incluye herramientas para la validación utilizando el Puntaje de Modelado de Datos Lineales (LDS) y variantes para filtrado top-k/bottom-k, así como soporte para experimentos de reentrenamiento leave-k-out.

Resultados experimentales

El artículo presenta varios experimentos de validación y casos de estudio:

• Validación de precisión: En un modelo GPT-2 (124M) ajustado con fine-tuning en WikiText-2, MAGIC (diferenciación desenrollada) logró la mayor correlación de Spearman (0.979) con los efectos reales de leave-one-out, seguido de EK-FAC (0.318) y TrackStar (0.164). Sin embargo, en evaluaciones de leave-k-out de subconjuntos ordenados (más relevantes para el filtrado de datos), todos los métodos alcanzaron correlaciones de 0.7 o superiores.
• Atribución a nivel de token (Bioseguridad): Utilizando el benchmark WMDP Bio, los autores ajustaron un modelo de 7B con LoRA y utilizaron la diferenciación desenrollada por token para identificar tokens influyentes. El reentrenamiento con un aumento en los pesos de importancia de los tokens más influyentes del top 10% mejoró la precisión de la evaluación en 1.5 puntos porcentuales adicionales en comparación con el reentrenamiento sin ponderar.
• Rastreo de hechos a través de estilos: En un experimento sintético que involucra hechos biográficos escritos en estilos "Shakespeareano" vs. "Legalese", TrackStar logró identificar exitosamente a los proponentes factuales correctos en el 15.26% de los resultados del top-5, comparado con el 0.80% de una línea base de producto interno de gradiente bruto. TrackStar también redujo significativamente la "filtración de estilo" (coincidir con el estilo de la consulta en lugar del origen del hecho).
• Desempeño de escalabilidad: El tiempo de recolección de gradientes escala eficientemente con el número de GPUs (de 1 a 16 A100s) para modelos que van desde 160M hasta 12B de parámetros, utilizando estrategias de FSDP y paralelismo de datos para minimizar la sobrecarga de comunicación.

Significado y afirmaciones

Los autores afirman que Bergson reduce la barrera de entrada para la investigación en atribución de datos al proporcionar una infraestructura unificada y escalable. Al ofrecer las primeras implementaciones de MAGIC, SOURCE y TrackStar, la biblioteca permite a los profesionales experimentar con técnicas de vanguardia que anteriormente eran inaccesibles o requerían un esfuerzo de ingeniería significativo para construirse.

El artículo enfatiza que una atribución de datos efectiva puede ayudar a comprender los comportamientos de los modelos, depurar generaciones inesperadas y curar conjuntos de datos de entrenamiento para la seguridad y el alineamiento. Los autores señalan que, si bien la curación de datos es de doble uso (capaz de amplificar o reducir comportamientos), se espera que la especificidad de la atribución de datos beneficie diferencialmente los esfuerzos de alineación y seguridad al permitir un direccionamiento preciso de los datos relevantes. La biblioteca se posiciona como una herramienta para permitir que más investigadores estudien comportamientos de modelos interesantes mediante la atribcción de datos, particularmente en el contexto de modelos a gran escala donde tal análisis era previamente impractico.