To Memorize or to Retrieve: Scaling Laws for RAG-Considerate Pretraining

Resumo Técnico: Leis de Escala para Pré-treinamento com RAG

1. O Problema

As Leis de Escala (Scaling Laws) tradicionais estabelecem como o desempenho de Modelos de Linguagem (LMs) melhora com o aumento do número de parâmetros e tokens de treinamento. No entanto, esses modelos tratam o corpus de treinamento como monolítico, assumindo que todo o conhecimento deve ser comprimido parametricamente nos pesos do modelo.

A Geração Aumentada por Recuperação (RAG) introduz uma nova variável: parte do corpus pode ser mantida como um datastore externo e acessada no momento da inferência. A relação entre o conhecimento paramétrico (adquirido no pré-treinamento) e o conhecimento não paramétrico (acessado via recuperação) é mal compreendida, especialmente sob orçamentos de dados fixos. A questão central do artigo é: dado um corpus fixo de $N$ tokens, qual é a alocação ótima entre dados para pré-treinamento (pesos) e dados para o datastore de recuperação?

2. Metodologia

Os autores realizaram um estudo sistemático variando a escala do modelo, o tamanho dos dados de pré-treinamento e o tamanho do datastore de recuperação.

• Configuração Experimental:

  • Modelos: Baseados na série OLMo-2, variando de 30M a 3B parâmetros.
  • Dados: Utilizaram até 100B tokens do corpus DCLM.
  • Variáveis:
    • Escala de pré-treinamento: 1x a 150x o número de parâmetros.
    • Tamanho do datastore de recuperação: 1x a 20x o número de parâmetros (construído a partir de fatias mantidas de fora do mesmo corpus DCLM).
  • Avaliação: Benchmarks diversificados cobrindo raciocínio, QA científica e QA de domínio aberto (ex: ARC, HellaSwag, SciQ, Natural Questions, StrategyQA).
  • Métrica Principal: Perplexidade (PPL) foi usada em vez de apenas acurácia, pois oferece uma medida contínua e suave para análise de leis de escala, evitando o comportamento "de degrau" da acurácia.

• Construção do Índice:

  • Utilizaram o FAISS com o algoritmo IVFPQ.
  • Modelo de embedding: Qwen3-Embedding-8B.
  • Os índices foram construídos de forma hierárquica (subconjuntos estritos) para permitir comparações controladas de escala.

3. Contribuições Principais

1. Descoberta de Relação Não Trivial: Demonstraram que a relação entre pré-treinamento e recuperação é estruturada, mas não monotônica. Os benefícios da recuperação dependem da escala do modelo, do tipo de tarefa e do grau de saturação do pré-treinamento.
2. Lei de Escala Tridimensional: Propuseram um novo modelo matemático que descreve a perda (loss) como uma função de três variáveis: tamanho do modelo ($N$), tokens de pré-treinamento ($D$) e tokens de recuperação ($R$).
   • A fórmula proposta é:
     $$L(N, D, R) = A \left(\frac{N}{10^9}\right)^{-\alpha} + B \left(\frac{D}{10^9}\right)^{-\beta} - C \log\left(1 + \eta \frac{R}{10^9}\right) + L_0$$
     Onde o termo logarítmico captura os ganhos com a recuperação e sua saturação.
3. Ponto de Cruzamento (Crossover Point): Identificaram um ponto de inflexão dependente da escala. Abaixo de uma certa razão de tokens por parâmetro, a recuperação não substitui eficazmente o pré-treinamento. Acima desse limiar (aprox. 4.14 tokens de pré-treinamento por parâmetro), a recuperação torna-se um substituto eficiente, onde cada token recuperado pode economizar múltiplos tokens de pré-treinamento.

4. Resultados Chave

• Substituibilidade Dependente da Escala:
  • Modelos Pequenos (ex: 30M): Beneficiam-se enormemente da recuperação, que atua como um substituto eficiente para a falta de capacidade paramétrica.
  • Modelos Grandes (ex: 3B): Os benefícios marginais da recuperação diminuem à medida que o modelo cresce e internaliza mais conhecimento. Modelos grandes exibem retornos decrescentes e sensibilidade ao excesso de alocação para recuperação.
• Natureza da Tarefa:
  • Tarefas Baseadas em Fatos (QA): A recuperação oferece ganhos significativos, especialmente quando o modelo está saturado ou é pequeno.
  • Tarefas de Raciocínio (ex: PIQA, StrategyQA): A recuperação tem impacto limitado ou nulo. O fator limitante nessas tarefas parece ser a capacidade de computação e raciocínio do modelo, e não o acesso a fatos externos.
• Qualidade da Recuperação: Ajustes na formulação da consulta (query) melhoram os resultados, mas não alteram as tendências fundamentais de escala. A recuperação não é um substituto universal para a capacidade paramétrica.

5. Significado e Implicações

O trabalho oferece uma base quantitativa para o design de sistemas de linguagem escaláveis:

• Otimização de Orçamento de Dados: Em vez de tratar pré-treinamento e RAG como escolhas de design separadas, eles devem ser otimizados conjuntamente. Para orçamentos fixos, a alocação ideal depende do tamanho do modelo e do domínio da tarefa.
• Mudança de Paradigma: Sugere que os corpora de treinamento não devem ser vistos apenas como dados para compressão em pesos, mas como recursos que podem ser particionados intencionalmente entre "memória interna" (pesos) e "memória externa" (índices).
• Eficiência: Para modelos menores ou em regimes de dados limitados, investir em um datastore de recuperação pode ser mais eficiente em termos de custo computacional do que aumentar massivamente o pré-treinamento.

Em resumo, o artigo estabelece que a recuperação é um substituto dependente da escala para o pré-treinamento, sendo mais valiosa para modelos menores e tarefas factuais, enquanto modelos grandes e tarefas complexas de raciocínio ainda dependem fortemente de capacidades paramétricas robustas.