Resumo Técnico: Bergson: Uma Biblioteca de Código Aberto para Atribuição de Dados

Declaração do Problema

A atribuição de dados visa explicar o comportamento do modelo ao determinar o efeito causal de itens individuais de treinamento em saídas específicas do modelo. Embora este campo ofereça aplicações em depuração, curadoria de conjuntos de dados e compreensão de memorização ou raciocínio, escalar essas técnicas para Large Language Models (LLMs) modernos apresenta desafios de engenharia significativos. A aplicação ingênua de métodos de atribuição de dados frequentemente resulta em requisitos de memória irreais, particularmente ao lidar com conjuntos de dados de pré-treinamento de grande escala ou modelos que exigem suporte multi-nó. Além disso, muitos métodos de ponta, como aqueles que alcançam altas correlações de Leave-One-Out (LOO), carecem de implementações de código aberto, dificultando a reprodutibilidade e o progresso.

Metodologia e Arquitetura da Biblioteca

O Bergson é introduzido como uma biblioteca de código aberto projetada para permitir um progresso mais rápido na atribuição de dados, fornecendo uma arquitetura composível e escalável. A biblioteca suporta uma ampla gama de técnicas, incluindo diferenciação através de trajetórias completas de treinamento, diferenciação desenrolada aproximada e funções de influência.

Componentes Principais

• Pipeline Componível: O Bergson utiliza uma arquitetidade de pipeline de dados (Figura 1) onde componentes como aproximações de Hessian, normalizadores, construtores de índices e pontuadores de atribuição são interoperáveis. Fluxos de trabalho podem ser definidos via uma Interface de Linha de Comando (CLI) ou programaticamente usando arquivos de configuração YAML.
• Armazenamento e Indexação de Gradientes: A biblioteca suporta armazenamentos de gradientes em disco indexados por módulo (incluindo cabeças de atenção divididas) e suporta coleta por token ou por sequência. Para lidar com grandes índices, o Bergson integra-se com o FAISS para buscas de vizinho mais próximo aproximado (ANN), enquanto também suporta buscas exatas top-k.
• Treinamento Distribuído: O Bergson implementa paralelismo replicado e fragmentado (DDP, FSDP) para escalar a atribuição para modelos com até 405 bilhões de parâmetros. Ele suporta operações multi-nó onde grupos FSDP podem processar shards de dados de forma independente.
• Integração com PEFT: A integração nativa com a biblioteca PEFT permite que a atribuição de dados seja restrita aos parâmetros do adaptador (ex: LoRA), reduzindo a pegada de gradiente por exemplo em ordens de magnitude.
• Objetivos Não Diferenciáveis: Para tarefas avaliadas via LLM-as-a-judge, o Bergson suporta atribuição via perda de Group Relative Policy Optimization (GRPO), permitindo a atribuição de comportamentos que não são capturados por perdas padrão de cross-entropy.

Algoritmos Implementados

O Bergson fornece as primeiras implementações de código aberto para três métodos líderes:

1. MAGIC: Um algoritmo para diferenciação exata através de loops de treinamento desenrolados usando checkpointing. Ele alcança $O(N \log N)$ de computação e $O(\log N)$ de espaço em disco (ou um novo trade-off de $O(N)$ computação/$O(\sqrt{N})$ disco).
2. SOURCE: Um método para diferenciação desenrolada aproximada que divide a trajetória de treinamento em segmentos, tratando Hessians e gradientes como constantes dentro de cada segmento.
3. TrackStar: Um método de aproximação escalável para funções de influência que projeta gradientes por exemplo em um espaço de baixa dimensão usando projeções aleatórias e utiliza aproximações de Hessian de autocorrelação.

A biblioteca também suporta variantes padrão de funções de influência (K-FAC, EK-FAC, TK-FAC/Shampoo) e estratégias de normalização de gradiente (Adam, Adafactor, normalização unitária).

Principais Contribuições

1. Implementação de Código Aberto: O Bergson introduz as primeiras implementações públicas de MAGIC, SOURCE e TrackStar, preenchendo uma lacuna de ferramentas para métodos de ponta de atribuição de dados.
2. Escalabilidade: A biblioteca é projetada para lidar com modelos de bilhões de parâmetros e conjuntos de dados massivos, suportando armazenamentos de gradientes em disco, integração com FAISS e treinamento distribuído multi-nó.
3. Guia de Seleção de Métodos: O artigo fornece um guia para que profissionais selecionem métodos baseados em restrições de recursos, volume de consultas e necessidade de precisão (por exemplo, recomendando diferenciação desenrolada para filtragem de segurança crítica e índices comprimidos para depuração de alta consulta).
4. Estrutura de Avaliação: O Bergson inclui ferramentas para validação usando o Linear Data Modeling Score (LDS) e variantes para filtragem top-k/bottom-k, bem como suporte para experimentos de re-treinamento leave-k-out.

Resultados Experimentais

O artigo apresenta vários experimentos de validação e estudos de caso:

• Validação de Precisão: Em um modelo GPT-2 (124M) ajustado com WikiText-2, o MAGIC (diferenciação desenrolada) alcançou a maior correlação de Spearman (0.979) com os efeitos de ground-truth leave-one-out, seguido por EK-FAC (0.318) e TrackStar (0.164). No entanto, em avaliações de subconjunto ordenado leave-k-out (mais relevantes para filtragem de dados), todos os métodos alcançaram correlações de 0.7 ou superiores.
• Atribuição ao Nível de Token (Biosecurity): Usando o benchmark WMDP Bio, os autores realizaram o ajuste fino de um modelo de 7B com LoRA e usaram diferenciação desenrolada por token para identificar tokens influentes. O re-treinamento com pesos de importância aumentados nos 10% de tokens mais influentes melhorou a precisão de avaliação em 1.5 pontos percentuais adicionais em comparação com o re-treinamento sem pesos.
• Rastreamento de Fatos entre Estilos: Em um experimento sintético envolvendo fatos biográficos escritos em estilos "Shakespeareano" vs. "Jurídico", o TrackStar conseguiu evidenciar os proponentes factuais corretos em 15.26% dos resultados top-5, comparado a 0.80% para uma baseline de produto interno de gradiente bruto. O TrackStar também reduziu significativamente o "vazamento de estilo" (combinando o estilo da consulta em vez do fato fonte).
• Desempenho de Escala: O tempo de coleta de gradiente escala eficientemente com o número de GPUs (1 a 16 A100s) para modelos variando de 160M a 12B de parâmetros, utilizando estratégias de FSDP e paralelismo de dados para minimizar o overhead de comunicação.

Significância e Alegações

Os autores alegam que o Bergson reduz a barreira de entrada para a pesquisa de atribuição de dados ao fornecer uma infraestrutura unificada e escalável. Ao oferecer as primeiras implementações de MAGIC, SOURCE e TrackStar, a biblioteca permite que profissionais experimentem técnicas de ponta que anteriormente eram inacessíveis ou exigiam um esforço de engenharia significativo para construir.

O artigo enfatiza que a atribuição de dados eficaz pode auxiliar na compreensão dos comportamentos do modelo, na depuração de gerações inesperadas e na curadoria de conjuntos de dados de treinamento para segurança e alinhamento. Os autores observam que, embora a curadoria de dados seja de uso duplo (capaz de amplificar ou reduzir comportamentos), a especificidade da atribuição de dados deve beneficiar diferencialmente os esforços de alinhamento e segurança ao permitir o direcionamento preciso de dados relevantes. A biblioteca é posicionada como uma ferramenta para permitir que mais pesquisadores estudem comportamentos interessantes de modelos através da atribuição de dados, particularmente no contexto de modelos de grande escala onde tal análise era anteriormente impraticável.