Resumo Técnico: Diffract: Uma Visão Espectral da Adaptação de Domínio de LLMs

Declaração do Problema

O pré-treinamento contínuo (CPT) é um mecanismo padrão para adaptar Grandes Modelos de Linguagem (LLMs) de propósito geral para domínios especializados (ex: matemática, código, instrução). Embora o CPT proporcione um desempenho superior ao treinamento de modelos especializados do zero sob orçamentos de computação semelhantes, os mecanismos subjacentes do CPT permanecem significativamente menos compreendidos do que os do Ajuste Fino Supervisionado (SFT). Fenômenos como conectividade de modo linear e aritmética de tarefas são bem documentados em SFT, mas não está claro se dinâmicas espectrais ou estruturais semelhantes regem o CPT. Especificamente, há uma falta de compreensão sobre como as matrizes de pesos evoluem durante o CPT, se a adaptação de domínio depende de mudanças nos valores singulares ou nos vetores singulares, e se existe redundância nas atualizações de parâmetros que possa ser explorada para eficiência.

Metodologia

Os autores utilizam a Decomposição de Valor Singular (SVD) para analisar as matrizes de pesos dos LLMs ao longo das trajetórias de pré-treinamento e pré-treinamento contínuo. O estudo utiliza a arquitetura OLMo 2 em escalas de modelo de 1B, 7B, 13B e 32B parâmetros.

Componentes analíticos principais:

1. Medidas Espectrais: Análise de normas de Frobenius, normas espectrais, ranks estáveis e ranks efetivos de matrizes de peso.
2. Concordância de Vetores Singulares: Medição da similaridade de cosseno entre os vetores singulares de checkpoints pré-treinados ($W_{pre-train}$) e checkpoints adaptados ao domínio ($W_{domain}$) para quantificar mudanças direcionais.
3. Ajuste de Distribuição: Ajuste de densidades espectrais empíricas (ESD) a distribuições de Marchenko-Pastur (MP) para o espectro de massa (bulk) e modelos de Lei de Potência (PL) para regiões de cauda pesada para avaliar o alinhamento com a teoria de Autorregularização de Cauda Pesada (HTSR).
4. Experimentos de Intervenção:
   • Transplante de Valor Singular: Troca de valores singulares ($\Sigma$) entre checkpoints retendo os vetores singulares para testar a importância das magnitudes espectrais versus direções.
   • Rewinding de Cabeça (Head-wise Rewinding): Reversão sistemática de cabeças de atenção para seu estado pré-treinado com base em vários critérios de importância para medir o impacto no desempenho de tarefas downstream.
   • Truncamento de SVD: Truncamento dos menores valores singulares do delta de CPT ($\Delta W = W_{domain} - W_{pre-train}$) para avaliar a redundância de parâmetros.
   • Interpolação Linear: Interpolação de pesos entre checkpoints adaptados a diferentes domínios para testar a "conectividade de domínio".

Os autores lançam o Diffract, um toolkit de código aberto projetado para facilitar essas análises espectrais.

Principais Contribuições e Resultados

1. Dinâmica Espectral: Espectros Invariantes, Vetores Evoluindo

O estudo descobre que o CPT deixa os espectros de valores singulares amplamente invariantes. O processo de adaptação é impulsionado principalmente por mudanças nos vetores singulares, particularmente aqueles localizados próximos aos picos dos espectros de SVD.

• Evolução do Pré-treinamento: Durante a etapa de pré-treinamento, as matrizes de cabeças de atenção desenvolvem estruturas espectrais complexas caracterizadas por outliers e múltiplos picos estreitos. Essa complexidade aumenta com o orçamento de tokens de pré-treinamento e desvia dos modelos HTSR padrão, que tipicamente preveem caudas simples de lei de potência.
• Estabilidade do CPT: Em contraste, os espectros de valores singulares permanecem quase estáveis durante o CPT. Os autores hipotetizam que a estrutura espectral complexa adquirida durante o pré-treinamento é um pré-requisito para uma adaptação de domínio eficiente.

2. Heterogeneidade de Cabeça e Rewinding

Um fenômeno inédito denominado heterogeneidade de cabeça é identificado. À medida que o orçamento de tokens de pré-treinamento aumenta, as cabeças de atenção exibem comportamentos variados durante o CPT:

• Algumas cabeças passam por mudanças significativas independentemente do domínio.
• Outras mudam de maneira específica ao domínio.
• Estratégia de Rewinding: Ao ordenar as cabeças com base em um novo critério (comparando a norma de Frobenius do delta de CPT contra um domínio de texto de referência), os autores demonstram que até 60% das cabeças de atenção podem ser revertidas (rewound) ao seu estado pré-treinado sem perda mensurável de qualidade. Além disso, o rewinding seletivo de cabeças de baixa importância pode melhorar a precisão em benchmarks em até 4% (especificamente em tarefas de matemática para o modelo 7B) comparado à linha de base totalmente treinada.

3. Redundância de Parâmetros nos Deltas de CPT

O delta de CPT exibe redundância significativa, que escala com o tamanho do modelo:

• Modelos de 1B: Mostram deltas de alto rank com menor tolerância à truncagem.
• Modelos de 7B e 13B: Podem tolerar truncagem agressiva. Até 50% dos valores singulares podem ser removidos do delta de CPT de 7B, e até 70% do delta de 13B, sem degradação significativa na precisão do GSM8K. Isso sugere que modelos maiores aprendem representações mais redundantes durante o CPT.

4. Conectividade de Domínio

Os autores identificam a conectividade de domínio, definida como a capacidade de interpolar linearmente entre checkpoints adaptados a diferentes domínios (ex: matemática e código) sem esquecimento catastrófico.

• Tendência de Convexidade: A qualidade do modelo interpolado ("model soup") melhora conforme o orçamento de tokens de pré-treinamento e o tamanho do modelo aumentam. Para orçamentos menores, a interpolação é côncava (desempenho inferior aos endpoints); para orçamentos maiores (ex: 4T tokens), torna-se convexa (superando a corda linear).
• Limitação: Apesar da existência de conectividade, a interpolação linear geralmente tem desempenho inferior aos modelos treinados diretamente em misturas de dados, sugerindo que a simples fusão de vetores de tarefa (task-vector merging) possui limitações.

Significância e Alegações

O artigo afirma fornecer uma compreensão fundamental dos mecanismos de CPT através de uma lente espectral. A significância reside em:

• Desafiar a Teoria HTSR: Demonstrar que, embora as camadas MLP se alinhem com HTSR, as cabeças de atenção desenvolvem estruturas espectrais complexas que os modelos de cauda pesada padrão não conseguem descrever totalmente.
• Eficiência: Fornecer um método principiado para identificar e remover parâmetros redundantes (até 60% das cabeças) ou melhorar o desempenho via rewinding seletivo, oferecendo um caminho para uma adaptação de domínio mais eficiente.
• Ferramental: O lançamento do Diffract permite a análise espectral reprodutível de modelos de bilhões de parâmetros, abordando uma lacuna nas ferramentas para analisar a dinâmica de pesos de LLMs.
• Escalabilidade: As descobertas sobre heterogeneidade de cabeça e conectividade de domínio mostram-se consistentes através de escalas de modelo (1B a 32B) e diversos domínios (matemática, instrução, código, texto).

Os autores mantêm uma postura modesta, observando que suas conclusões são robustas dentro da família OLMo, mas podem não ser universais para todas as arquiteturas modernas de LLM. Eles reconhecem limitações na cobertura de benchmarks e no uso específico do otimizador AdamW, sugerindo que trabalhos futuros explorem outras famílias de modelos e ablações de otimizadores.