Resumo Técnico: LoRIF – Funções de Influência de Baixo Rank para Atribuição de Dados de Treinamento Escalável

Declaração do Problema

A Atribuição de Dados de Treinamento (TDA) visa identificar quais exemplos de treinamento específicos mais influenciaram a previsão de um modelo em uma determinada consulta. Embora as Funções de Influência (IF) forneçam uma abordagem teoricamente fundamentada para a TDA, aproximando o retreinamento leave-one-out por meio de gradientes e matrizes Hessianas inversas, elas permaneceram efetivamente inacessíveis para modelos de escala moderna (bilhões de parâmetros) treinados em milhões de exemplos.

Variantes escaláveis existentes, como TRAK, LoGRA e TrackStar, abordam gargalos computacionais armazenando gradientes projetados por exemplo e utilizando uma aproximação Hessiana de Gauss-Newton amortecida. No entanto, esses métodos enfrentam dois gargalos críticos de escalabilidade que criam um tradeoff entre a qualidade da atribuição e a viabilidade computacional:

1. Armazenamento de Gradientes e E/S: Armazenar gradientes projetados por exemplo para $N$ exemplos de treinamento requer espaço $O(ND)$, onde $D$ é a dimensão de projeção efetiva. No momento da consulta, o carregamento desses gradientes domina a latência (aproximadamente 96% do tempo de consulta).
2. Escala da Hessiana Inversa: A formação e o armazenamento da aproximação Hessiana inversa $D \times D$ requerem memória $O(D^2)$.

Aumentar a dimensão de projeção $D$ é necessário para melhorar a qualidade da atribuição, mas fazê-lo exacerba tanto os gargalos de armazenamento/E/S quanto os de memória, tornando a atribuição de alta qualidade impraticável para modelos de escala de fronteira.

Metodologia: LoRIF

O LoRIF (Funções de Influência de Baixo Rank) aborda esses gargalos explorando a estrutura de baixo rank inerente aos gradientes de redes neurais. O método segue um pipeline de três etapas:

1. Armazenamento de Gradientes de Baixo Rank (Abordando o Gargalo de E/S)

Em vez de armazenar a matriz completa de gradientes projetados por exemplo $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$ (onde $D = d_1 d_2$), o LoRIF armazena uma fatoração de rank-$c$:
$$\tilde{G}_i \approx u_i v_i^\top$$
onde $u_i \in \mathbb{R}^{d_1 \times c}$ e $v_i \in \mathbb{R}^{d_2 \times c}$.

• Mecanismo: Essa fatoração é computada via iterações de potência em blocos.
• Impacto: O armazenamento e a E/S no momento da consulta por camada por amostra são reduzidos de $O(D)$ para $O(c(\sqrt{D}))). Resultados empíricos sugerem que, mesmo com$c=1$, a aproximação mantém qualidade de atribuição significativa, e priorizar um $D$ maior em vez de um $c$ maior produz melhor eficiência de custo.

2. Aproximação de Curvatura por SVD Truncado (Abordando o Gargalo de Memória)

Para evitar a formação da Hessiana inversa densa $D \times D$ $(G^\top G + \lambda I)^{-1}$, o LoRIF computa uma Decomposição em Valores Singulares (SVD) truncada de rank-$r$ da matriz de gradiente agregada $G$:
$$G \approx U_r \Sigma_r V_r^\top$$
Usando a identidade de Woodbury, o termo da Hessiana inversa é aproximado em um subespaço de dimensão $r$:
$$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$$

• Mecanismo: Apenas os $r$ vetores singulares superiores ($V_r$) e valores singulares ($\Sigma_r$) são armazenados.
• Impacto: A complexidade de memória para o termo de curvatura é reduzida de $O(D^2)$ para $O(Dr)$. Isso permite uma dimensão de projeção efetiva $D$ significativamente maior dentro de um orçamento de memória fixo.

3. Cálculo da Pontuação de Influência

No momento da consulta, as pontuações de influência são computadas projetando o gradiente da consulta no subespaço de dimensão $r$ e utilizando os fatores de baixo rank armazenados para os gradientes de treinamento. A pontuação de influência $I(x_{tr}, x_{te})$ é aproximada como:
$$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$$
onde $g'$ representa gradientes projetados sobre os vetores singulares $V_r$. O primeiro termo é computado eficientemente usando os fatores de rank-$c$, e o segundo termo envolve apenas operações $O(r)$ devido à natureza diagonal da matriz intermediária.

Contribuições Principais

1. Identificação de Gargalos: O artigo identifica que a dimensão de projeção $D$ é a restrição primária na TDA estilo IF escalável, criando um tradeoff qualidade-escalabilidade onde a atribuição precisa (grande $D$) é atualmente impraticável.
2. Fatoração de Rank-$c$: A proposta de armazenar gradientes projetados como fatores de rank-$c$, reduzindo o armazenamento por exemplo por camada de $O(D)$ para $O(c\sqrt{D})$.
3. Aproximação de Curvatura de Rank-$r$: A proposta de usar SVD truncado e a identidade de Woodbury para aproximar o termo de curvatura, reduzindo a memória por camada de $O(D^2)$ para $O(Dr)$.
4. Demonstração de Escalabilidade: A primeira demonstração de atribuição de dados efetiva em modelos de até 70 bilhões de parâmetros (Apertus-70B) com milhões de exemplos de treinamento, mantendo custo e latência práticos.

Resultados Experimentais

O LoRIF foi avaliado em modelos variando de 124M a 70B parâmetros (GPT2-small, OLMo-3-7B, Apertus-70B) com conjuntos de dados contendo até 3,8 milhões de exemplos.

• Armazenamento e Latência: Comparado ao LoGRA, o LoRIF alcança armazenamento de gradiente 2,3–20× menor e latência de consulta 1,3–20× menor enquanto iguala ou supera a qualidade de atribuição.
  • No GPT2-small, o LoRIF reduziu o armazenamento em até 9× e a latência em 5,6× em regimes de armazenamento médio.
  • No OLMo-3-7B, o LoRIF alcançou pontuações de tail-patch quase idênticas ao LoGRA com 20,3× menos armazenamento e 22,0× menos latência.
  • No Apertus-70B, o LoRIF melhorou as pontuações de tail-patch em 1,5–2,2× sobre o LoGRA com menor armazenamento e latência.
• Qualidade de Atribuição: Medida via Pontuação de Modelagem de Dados Linear (LDS) e Pontuação de tail-patch, o LoRIF consistentemente melhorou a fronteira de Pareto. Ao permitir dimensões de projeção efetivas $D$ maiores (que eram anteriormente inviáveis devido a restrições de memória), o LoRIF alcançou qualidade superior às baselines restritas a $D$ menores.
• Avaliação Qualitativa: Uma avaliação com LLM como juiz nos exemplos recuperados top-1 mostrou que o LoRIF recupera exemplos com maior relevância semântica (pontuação média 3,04 vs. 2,49 no OLMo-3-7B) e menos recuperações completamente irrelevantes em comparação ao LoGRA.

Significado e Alegações

O artigo alega que o LoRIF torna a Atribuição de Dados de Treinamento baseada em gradientes prática em escala de fronteira. Ao explorar a estrutura de baixo rank dos gradientes, o LoRIF resolve o tradeoff qualidade-escalabilidade que tem impedido a aplicação de funções de influência a grandes modelos de linguagem.

Os autores enfatizam que o LoRIF não depende de novos primitivos de baixo rank, mas sim demonstra como estruturas de baixo rank existentes (observadas tanto em gradientes por exemplo quanto agregados) podem ser exploradas sistematicamente em métodos de influência de gradiente projetado. Os resultados sugerem que a TDA baseada em gradientes pode agora ser usada para auditoria, depuração e verificações de segurança de LLMs em grande escala quando o índice de atribuição é reutilizado em muitas consultas, sem exigir os custos proibitivos de armazenamento ou E/S de métodos anteriores.

O artigo reconhece limitações, observando que a construção do índice de atribuição ainda requer uma passagem de gradiente sobre o corpus (custo de pré-processamento) e que os experimentos atuais em grande escala focam em dados de ajuste fino supervisionado em vez de corpora completos de pré-treinamento. No entanto, o trabalho estabelece uma nova linha de base para atribuição escalável, permitindo a análise de modelos anteriormente considerados fora de alcance para tais métodos.