Not all tokens are needed(NAT): token efficient reinforcement learning

O artigo apresenta o NAT (Not All Tokens Are Needed), um framework de aprendizado por reforço que otimiza o custo computacional ao atualizar a política apenas com um subconjunto de tokens gerados via reponderamento de Horvitz-Thompson, mantendo o desempenho em raciocínio matemático enquanto reduz significativamente o uso de memória e tempo de treinamento.

O essencial

Imagine que você está ensinando um aluno muito inteligente (uma Inteligência Artificial) a resolver problemas de matemática complexos. Para isso, você deixa o aluno tentar resolver o problema sozinho, gera uma resposta longa e cheia de passos (o "raciocínio"), e depois você verifica se a resposta final está correta. Se estiver, você dá um "bom trabalho"; se não, um "tente de novo".

O problema é que, para o computador aprender, ele precisa revisar cada palavra que o aluno escreveu, uma por uma, para entender onde errou ou acertou. Se o aluno escrever um texto gigante com 1.000 palavras, o computador tem que "ler" e "revisar" todas as 1.000 palavras para cada tentativa. Isso consome muita energia, memória e tempo, como se você tivesse que reler um livro inteiro só para corrigir uma única vírgula no final.

Os autores deste artigo (chamado NAT) perguntaram: "Será que precisamos ler e revisar todas as palavras para ensinar o aluno?"

A resposta deles é um sonoro NÃO.

A Ideia Principal: O "Filtro Inteligente"

A equipe criou um método chamado NAT (Not All Tokens Are Needed - "Nem Todos os Tokens São Necessários"). A ideia é simples, mas genial:

1. O Aluno continua escrevendo tudo: O computador gera a resposta completa, com todos os passos, para garantir que a nota final (a recompensa) seja justa e baseada no resultado completo.
2. O Professor revisa apenas uma parte: Em vez de revisar as 1.000 palavras, o computador escolhe aleatoriamente apenas 500 palavras (metade) para fazer a correção e o aprendizado.
3. O Truque Mágico (A Balança): Como o computador está ignorando metade das palavras, ele precisa garantir que a lição aprendida não fique "viciada". Para isso, eles usam uma técnica matemática chamada reponderação de Horvitz-Thompson.

A Analogia da Balança:
Imagine que você está pesando uma sacola de frutas. Você decide pesar apenas 5 maçãs aleatórias em vez de todas as 100. Para saber o peso total da sacola, você não pode simplesmente somar as 5 maçãs. Você precisa multiplicar o peso de cada uma delas por um "fator de ajuste" (como se cada maçã representasse 20 maçãs na sua mente).

O NAT faz exatamente isso: ele pega as palavras que não foram revisadas e "dá peso" extra às palavras que foram revisadas. Assim, matematicamente, o computador aprende exatamente a mesma coisa que aprenderia se tivesse lido tudo, mas gastando metade do tempo e da energia.

Duas Formas de Fazer Isso

O artigo testa duas maneiras de escolher quais palavras revisar:

1. Amostragem Aleatória (URS): O computador escolhe palavras aleatórias espalhadas pelo texto (como pegar 500 palavras soltas de um livro).

   • Problema: Mesmo escolhendo palavras aleatórias, o computador ainda precisa "ler" o livro inteiro do início até a palavra escolhida para entender o contexto. É como se você lesse todo o livro, mas só anotasse em um caderno metade das páginas. Você ainda gasta o tempo de leitura.

2. Corte de Prefixo Aleatório (RPC) - O Vencedor: O computador escolhe um ponto aleatório no texto e diz: "Vou revisar apenas os primeiros 500 passos e ignorar o resto".

   • Vantagem: Aqui, o computador realmente para de ler após o ponto de corte. Ele não precisa processar o final da frase. É como se você lesse apenas os primeiros capítulos de um livro para entender a trama principal, sem precisar chegar ao final.
   • Resultado: Isso economiza memória (o computador não precisa guardar o final da frase na memória) e tempo (ele para de processar mais cedo).

Por que isso é importante?

Hoje, os modelos de IA estão ficando cada vez mais "pensadores", gerando raciocínios longos e complexos (como resolver equações de física ou escrever códigos). Isso está ficando tão caro e lento que quase não dá para treinar novos modelos.

O método NAT, especialmente a técnica de Corte de Prefixo (RPC), permite que:

• O treinamento seja 29% mais rápido.
• O computador use 18% menos memória (o que evita que ele "trave" por falta de espaço).
• A qualidade do aprendizado não caia. O aluno continua ficando tão inteligente quanto antes.

Resumo em uma frase

O NAT é como um professor esperto que, em vez de corrigir cada linha de um trabalho longo do aluno, corrige apenas a metade inicial de forma inteligente, ajustando a nota final para garantir que o aluno aprenda tudo o que precisa, mas gastando metade do tempo e da energia do computador.

Resumo técnico

Título: Not All Tokens Are Needed (NAT): Aprendizado por Reforço Eficiente em Tokens

1. O Problema

O Aprendizado por Reforço (RL) tornou-se um motor fundamental para o avanço de Modelos de Linguagem de Grande Escala (LLMs), especialmente em tarefas de raciocínio complexo que exigem cadeias de pensamento longas (Chain-of-Thought ou CoT). No entanto, escalar o RL para trajetórias longas enfrenta um gargalo crítico:

• Custo Computacional Excessivo: Os pipelines padrão de RL (como GRPO - Group Relative Policy Optimization) realizam a retropropagação (backpropagation) através de todos os tokens gerados em cada trajetória.
• Memória e FLOPs: À medida que as trajetórias de raciocínio ficam mais longas, o custo de memória de ativação e o número de operações de ponto flutuante (FLOPs) aumentam drasticamente, muitas vezes levando a falhas de memória (OOM) ou otimização instável.
• Ineficiência: Nem todos os tokens contribuem igualmente para o sinal de aprendizado. Muitos são apenas continuação mecânica ou "boilerplate", enquanto poucos representam pontos de decisão de alto impacto.
• Limitação Atual: Mesmo com engines de inferência otimizados para gerar rollouts mais rápido, a fase de aprendizado (forward/backward passes) permanece limitada pela memória e computação, criando um teto para o escalonamento.

2. Metodologia: Framework NAT

Os autores propõem o NAT (Not All Tokens Are Needed), um framework unificado que trata o orçamento de tokens como uma primitiva de otimização de primeira classe. A ideia central é atualizar a política usando apenas um subconjunto selecionado de tokens, mantendo o sinal de recompensa calculado sobre a resposta completa.

Componentes Chave:

• Mascaramento de Tokens: Durante o cálculo do gradiente da política, apenas um subconjunto de tokens é incluído na retropropagação.
• Correção de Viés via Horvitz-Thompson (HT): Para garantir que a atualização seja estatisticamente correta (não viesada) apesar da subamostragem, o NAT utiliza o estimador de Horvitz-Thompson.
  • Cada token $t$ recebe uma probabilidade de inclusão $p_{i,t}$.
  • Os gradientes são reponderados pelo inverso dessa probabilidade ($1/p_{i,t}$).
  • Proposição Teórica: O artigo prova que esse estimador corrigido por HT é não viesado em relação ao gradiente de RL completo, preservando o sinal de aprendizado esperado.

Esquemas de Seleção Propostos:
O NAT é implementado com duas estratégias de seleção de tokens:

1. Amostragem Aleatória Uniforme (URS - Uniform Random Sampling):
   • Seleciona tokens independentemente com uma probabilidade constante $p$.
   • Vantagem: Reduz a computação de retropropagação.
   • Limitação: Em Transformers causais, o custo do forward pass (passagem direta) não diminui significativamente, pois o modelo ainda precisa processar todos os tokens anteriores para calcular a probabilidade do token atual, devido à atenção causal.
2. Corte Aleatório de Prefixo (RPC - Random Prefix Cutting):
   • Seleciona um prefixo contíguo da resposta (ex: os primeiros $L$ tokens).
   • Vantagem Principal: Permite a truncagem real do forward pass. O modelo só precisa processar o prefixo selecionado, reduzindo tanto a memória de ativação quanto os FLOPs de forward e backward.
   • Diferença Crítica: Diferente do truncamento determinístico (que sempre corta o final e introduz viés sistemático), o RPC é estocástico, garantindo que todos os posições tenham probabilidade não nula de serem incluídas, mantendo a imparcialidade estatística via HT.

3. Contribuições Principais

1. Framework Unificado NAT: Uma abordagem para RLVR (Reinforcement Learning from Verifiable Rewards) que desacopla a avaliação de recompensa (em toda a sequência) da otimização da política (em subconjunto de tokens).
2. Fundamentação Teórica de Imparcialidade: Demonstração rigorosa de que a reponderação Horvitz-Thompson permite atualizações parciais de tokens sem introduzir viés sistemático no gradiente.
3. Eficiência Prática sem Perda de Qualidade: Evidência empírica de que o NAT (especialmente o RPC) pode reduzir o volume de retropropagação em até 50% mantendo o desempenho equivalente ao GRPO completo.

4. Resultados Experimentais

Os experimentos foram conduzidos em benchmarks de raciocínio matemático (MATH, AIME24, AIME25) utilizando modelos como Qwen2.5-Math-7B e Qwen3-8B.

• Desempenho de Precisão:
  • O RPC e o URS alcançaram desempenho estatisticamente equivalente ao GRPO completo (full-token) em todas as métricas de precisão (Acc@16 e Pass@16).
  • O Truncamento Determinístico (cortar sempre 50% do final) resultou em queda significativa de desempenho e alta variância, confirmando a necessidade da correção estocástica do NAT.
• Eficiência de Memória (GPU):
  • O RPC reduziu o pico de memória da GPU em ~18% (ex: de 47.72 GB para 39.23 GB no Qwen3-8B).
  • O URS não reduziu significativamente a memória de pico, pois não encurta o prefixo processado no forward pass.
• Tempo de Treinamento:
  • O RPC reduziu o tempo de treinamento (forward + backward, sem inferência) em ~29% para o Qwen3-8B.
  • O tempo total por passo (incluindo inferência) também melhorou, embora o ganho seja menor devido ao tempo fixo de geração de rollouts.

5. Significado e Impacto

O trabalho NAT oferece uma via ortogonal às otimizações de sistema existentes (como engines de inferência mais rápidos ou speculative decoding).

• Quebra de Gargalo: Resolve o gargalo de consumo de trajetórias longas durante a fase de aprendizado, permitindo treinar modelos com contextos maiores sem estourar a memória.
• Escalabilidade: Proporciona um caminho viável para escalar o RL em tarefas de raciocínio complexo (CoT longo), onde o custo de treinamento era proibitivo.
• Generalidade: O framework é compatível com qualquer algoritmo de RL baseado em gradiente e pode ser combinado com outras otimizações de sistema.

Em resumo, o NAT demonstra que não é necessário processar todos os tokens para aprender, desde que a seleção seja feita de forma inteligente e corrigida estatisticamente, oferecendo ganhos substanciais de eficiência computacional e de memória sem sacrificar a inteligência do modelo.