Reinforced Fast Weights with Next-Sequence Prediction

Título: Fast Weights Reforçados com Previsão da Próxima Sequência (REFINE)

1. O Problema

O modelamento de contexto longo é essencial para Grandes Modelos de Linguagem (LLMs), mas arquiteturas baseadas em attention (transformers) sofrem com custos computacionais e de memória que crescem quadraticamente com o tamanho do contexto.

• Alternativa Promissora: Arquiteturas de Fast Weights (como DeltaNet e LaCT) oferecem uma solução ao substituir a atenção global por uma memória de tamanho fixo e dinamicamente atualizada, mantendo um custo de memória constante independentemente do comprimento do contexto.
• Limitação Atual: Apesar de sua eficiência, os modelos de Fast Weights são treinados usando o paradigma padrão de Previsão do Próximo Token (NTP - Next-Token Prediction).
  • O NTP otimiza apenas a previsão de um único token por vez, ignorando a coerência semântica entre múltiplos tokens subsequentes.
  • Isso resulta em representações subótimas que falham em capturar dependências de longo alcance, pois o modelo não é incentivado a manter um estado interno coerente para prever sequências futuras, apenas o próximo token imediato.

2. Metodologia: O Framework REFINE

Os autores propõem o REFINE (Reinforced Fast Weights with Next sEquence prediction), um framework de Aprendizado por Reforço (RL) que substitui ou complementa o NTP com o objetivo de Previsão da Próxima Sequência (NSP - Next-Sequence Prediction).

O objetivo NSP incentiva o modelo a prever uma sequência semanticamente coerente de vários tokens futuros, alinhando-se melhor à função de memória de longo prazo dos Fast Weights. O processo envolve quatro etapas principais (ilustrado na Figura 3 do artigo):

1. Seleção de Tokens Baseada em Entropia:

   • Em vez de treinar em todas as posições, o modelo calcula a entropia da previsão do próximo token em cada posição.
   • O contexto é dividido em chunks (fatias). Em cada chunk, um token é amostrado com probabilidade proporcional à sua entropia (alta entropia indica alta incerteza/complexidade).
   • Isso foca o treinamento em regiões do contexto onde o modelo tem mais dificuldade, otimizando o uso de recursos computacionais.

2. Geração de Rollouts (Simulações):

   • Para cada posição de alta entropia selecionada, o modelo gera uma continuação de $k$ tokens (um rollout) a partir do prefixo atual.
   • São extraídos os estados ocultos (hidden states) tanto da continuação gerada quanto da continuação real (ground truth).

3. Atribuição de Recompensa (Reward Assignment):

   • Em vez de usar apenas a perda de entropia cruzada (que pune severamente respostas semanticamente corretas mas lexicalmente diferentes), o REFINE utiliza recompensas baseadas na similaridade de estados ocultos.
   • Recompensa de Similaridade ($R_\phi$): Calculada via similaridade de cosseno entre os estados ocultos da previsão e do ground truth. Isso permite recompensar respostas semanticamente equivalentes.
   • Recompensa Binária ($R_{binary}$): Usada especificamente para Treinamento em Tempo de Teste (TTT), baseada em correspondência exata, para garantir memorização precisa.
   • Recompensa Híbrida: Combinação das duas para fases de pós-treinamento.

4. Otimização com RL (GRPO):

   • O modelo é otimizado usando o algoritmo GRPO (Group Relative Policy Optimization).
   • A função de perda combina a perda de NSP (baseada nas recompensas de RL) com a perda padrão de NTP (SFT) para evitar o esquecimento catastrófico, ajustando os pesos ($\lambda_{RL}$ e $\lambda_{SFT}$) conforme a fase de treinamento.

3. Contribuições Principais

1. Objetivo NSP para Fast Weights: Introdução da Previsão da Próxima Sequência como um objetivo de treinamento superior ao NTP para arquiteturas de Fast Weights, fornecendo feedback em nível de sequência.
2. Framework REFINE: Proposta de um método de RL que combina seleção de tokens baseada em entropia e recompensas de nível de sequência para otimizar o NSP eficientemente.
3. Versatilidade no Ciclo de Vida: Demonstração de que o REFINE é eficaz em três fases distintas do treinamento de LLMs:
   • Mid-training: Adaptação de modelos pré-treinados para o objetivo NSP.
   • Post-training: Refinamento para tarefas específicas e alinhamento.
   • Test-Time Training (TTT): Adaptação direta no prompt durante a inferência, sem necessidade de novos dados rotulados.

4. Resultados Experimentais

Os experimentos foram conduzidos em dois modelos de Fast Weights: LaCT-760M e DeltaNet-1.3B.

• Recuperação de Contexto Longo (RULER / Needle-in-a-Haystack):
  • O REFINE superou consistentemente o ajuste fino supervisionado (SFT) com NTP.
  • No DeltaNet-1.3B, houve ganhos de 20,3% no mid-training, 11,0% no post-training e 15,0% no TTT em tarefas de recuperação de agulha no palheiro.
• Perguntas e Respostas (QA) e LongBench:
  • O REFINE demonstrou melhorias significativas em tarefas de QA multi-documento e em 12 tarefas do benchmark LongBench.
  • Exemplo: No LaCT-760M, o mid-training com REFINE melhorou a performance média no RULER em 8,5% em comparação ao SFT puro.
• Análise de Componentes:
  • Entropia: A seleção baseada em entropia superou amostragem uniforme e seleção de máxima/mínima entropia, indicando que focar em regiões de incerteza balanceada é crucial.
  • Recompensa: A recompensa baseada em similaridade de cosseno ($R_\phi$) superou a recompensa binária estrita durante o mid-training, sugerindo melhor generalização semântica.
  • Comprimento do Rollout ($k$): Um comprimento de $k=5$ foi identificado como ótimo; valores maiores degradaram a nitidez da recompensa.

5. Significado e Impacto

• Superação de Limitações Arquitetônicas: O trabalho demonstra que a limitação de modelos de Fast Weights não é apenas arquitetural, mas também de treinamento. Ao mudar o objetivo de "próximo token" para "próxima sequência", o potencial de memória de longo prazo dessas arquiteturas é desbloqueado.
• Eficiência: O framework permite que modelos com memória constante (O(1)) rivalizem ou superem modelos baseados em atenção em tarefas de contexto longo, sem o custo quadrático de inferência.
• Aplicabilidade Prática: A capacidade de aplicar o REFINE no Test-Time Training é particularmente inovadora, permitindo que modelos se adaptem dinamicamente a contextos específicos durante a inferência, melhorando a precisão em tempo real.
• Futuro: O estudo sugere que o RL para objetivos de sequência é um caminho viável e flexível para o desenvolvimento de futuros modelos de linguagem eficientes e de longo contexto.

Em resumo, o REFINE oferece uma solução robusta para o gargalo de treinamento de modelos de Fast Weights, transformando-os em sistemas de memória de longo alcance mais eficazes através de um aprendizado por reforço orientado à coerência sequencial.