Atlas: Pré-treinamento Conjunto de Recuperador-Leitor Supera LLMs de 540B Parâmetros com 11B Parâmetros

O Atlas é o trabalho subsequente de Izacard e Grave ao seu próprio artigo sobre Fusion-in-Decoder, estendendo o FiD em um sistema totalmente treinado de forma conjunta, onde o recuperador e o leitor são co-treinados desde o início. Estou lendo-o agora porque ele fecha a linhagem arquitetônica desde o artigo original sobre RAG, passando pelo FiD e chegando à recuperação treinada conjuntamente — exatamente o espaço de decisão pelo qual qualquer sistema de QA para livros razão precisa navegar.

O artigo

"Atlas: Few-shot Learning with Retrieval Augmented Language Models" (Izacard et al., JMLR 2023) questiona se modelos aumentados por recuperação podem igualar LLMs de parâmetros massivos em tarefas de poucos exemplos (few-shot) intensivas em conhecimento. A contribuição central é um sistema aumentado por recuperação cuidadosamente pré-treinado que treina conjuntamente um recuperador denso baseado em Contriever junto com um leitor Fusion-in-Decoder baseado em T5. O insight principal é que o pré-treinamento conjunto — e não a arquitetura — é o que impulsiona o desempenho do conhecimento em poucas etapas. O sistema recupera os 20 documentos principais, codifica cada um independentemente no codificador e, em seguida, os funde na atenção cruzada (cross-attention) do decodificador, o mesmo design FiD do artigo de 2021 dos autores.

Ideias-chave

• O Atlas-11B atinge 42,4% de precisão no Natural Questions com apenas 64 exemplos de treinamento, superando o PaLM (540B parâmetros) por cerca de 3 pontos, utilizando 50 vezes menos parâmetros.
• No TriviaQA (64-shot), o Atlas-11B alcança 74,5% no conjunto filtrado e 84,7% no teste oculto não filtrado, mostrando que o componente de recuperação compensa fortemente a supervisão limitada da tarefa.
• Quatro objetivos de treinamento do recuperador são avaliados: Destilação de Atenção (ADist), EMDR2 (tratando documentos recuperados como variáveis latentes), Destilação de Perplexidade (PDist) e LOOP (leave-one-out). As diferenças de desempenho entre eles são pequenas; o PDist é adotado pela eficiência computacional.
• O pré-treinamento conjunto em texto não rotulado é o fator individual mais importante: todas as configurações de pré-treinamento aumentadas por recuperação superam fortemente o baseline apenas com ajuste fino aumentado por recuperação.
• O índice de documentos pode ser atualizado após o treinamento sem re-treinar o modelo, o que é arquiteturalmente importante para bases de conhecimento dinâmicas. Índices temporalmente incompatíveis degradam o desempenho visivelmente.
• No MMLU (5-shot), o Atlas-11B atinge 47,9%, superando os 43,9% relatados do GPT-3, apesar de ter aproximadamente 16 vezes menos parâmetros.

O que se sustenta — e o que não se sustenta

A afirmação principal — de que a recuperação permite o desempenho do conhecimento em poucas etapas com uma fração da contagem de parâmetros — se sustenta de forma convincente. O número de 42,4% no NQ com 64 exemplos é um resultado impressionante, e a comparação com o PaLM é justa porque o PaLM era o benchmark de escala de última geração na época.

Mas tenho três ressalvas. Primeiro, a precisão da recuperação não é excelente mesmo após o treinamento conjunto: análises independentes mostram que o Contriever erra pelo menos uma afirmação de referência em cerca de 85% dos casos e atinge cerca de 47% de precisão de recuperação em QA. O treinamento conjunto melhora a recuperação em relação aos baselines não treinados conjuntamente, mas o leitor está fazendo um trabalho enorme para compensar a recuperação imperfeita — os números principais de poucas etapas refletem o teto do sistema, não a qualidade do componente de recuperação. Segundo, o custo da infraestrutura é real: atualizar os índices de documentos durante o pré-treinamento adiciona aproximadamente 30% de custo computacional adicional, e o índice completo da Wikipedia + CommonCrawl requer 587 GB em fp16. Isso é gerenciável para um ambiente de pesquisa, mas é uma restrição operacional genuína para implantação em produção. Terceiro, o vazamento de dados é reconhecido, mas não resolvido: 2,8% das perguntas do MMLU aparecem literalmente no corpus CCNet usado para o pré-treinamento, inflando os resultados do MMLU por uma margem desconhecida.

Há também uma limitação arquitetônica mais sutil que o artigo não aborda totalmente: o FiD codifica cada passagem recuperada de forma independente antes da fusão, o que ajuda no paralelismo, mas significa que o codificador não tem atenção entre passagens. Cadeias de raciocínio longas de múltiplas etapas (multi-hop) que precisam conectar informações entre passagens devem fazer todo esse trabalho no decodificador — e com 20 passagens recuperadas, a atenção cruzada do decodificador está carregando uma carga pesada.

Por que isso importa para a IA financeira

Para o QA de livros razão Beancount, a contribuição mais relevante do Atlas é a demonstração empírica de que o treinamento conjunto recuperador-leitor compensa em cenários de poucas etapas — e sua prestação de contas honesta de quando isso não acontece. Um agente Beancount consultando um histórico de transações de vários anos enfrenta exatamente o problema do índice dinâmico: novas entradas chegam diariamente, e um índice com um mês de atraso produz respostas erradas. O Atlas mostra que o índice pode ser trocado a quente sem re-treinamento, o que é encorajador do ponto de vista arquitetônico.

Os números de precisão de recuperação são sóbrios, no entanto. Se o Contriever erra a entrada relevante do livro razão em 53% das tentativas de recuperação, mesmo após o treinamento conjunto em texto geral, um agente do domínio financeiro operando sobre livros razão Beancount — com seus nomes de commodities específicos do domínio, hierarquias de contas e diretivas bean — precisará de treinamento de recuperador adaptado ao domínio ou recuperação aumentada por métodos de consulta estruturada (correspondência exata de conta, filtragem de data). A recuperação no estilo RAG sozinha, mesmo treinada conjuntamente, não será suficiente para operações de livro razão de alta precisão.

A comparação com o PaLM também esclarece a troca arquitetônica: a recuperação permite comprimir o conhecimento em menos parâmetros, reduzindo o custo de inferência. Para um produto como o Beancount.io, onde o custo de inferência importa em escala, a filosofia de design do Atlas é atraente. Mas o custo do índice de 587 GB transfere a carga para a infraestrutura de armazenamento e recuperação — um tipo diferente de restrição operacional que não aparece nos números de benchmark.

O que ler a seguir

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — a estrutura anterior de pré-treinamento conjunto recuperador-leitor que o Atlas estende; essencial para entender o que o Atlas realmente melhora e o que deixa inalterado.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — alcança desempenho competitivo com o Atlas usando ajuste de instrução em vez de pré-treinamento conjunto do zero; sugere que a lacuna entre o treinamento conjunto e o independente pode ser fechada sem o custo de infraestrutura.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — abordagem da DeepMind para recuperação durante o pré-treinamento em uma escala diferente; completa o quadro das abordagens de pré-treinamento aumentadas por recuperação antes de fazer escolhas arquitetônicas para o QA de livros razão.