Resumo Técnico: Raciocínio em Memória (RiM)

Declaração do Problema

Os métodos atuais para aprimorar as capacidades de raciocínio de Modelos de Linguagem de Grande Escala (LLMs) dependem fortemente da escalabilidade do cálculo no momento do teste, gerando tokens intermediários antes da resposta final, como no raciocínio em Cadeia de Pensamento (CoT). No entanto, essa abordagem acopla o cálculo interno diretamente à geração de texto autoregressiva. Consequentemente, o modelo é forçado a "pensar em voz alta", alocando uma parte significativa de seu orçamento computacional para gerar linguagem natural gramaticalmente fluente em vez de raciocínio puramente interno.

Embora as recentes abordagens de "raciocínio latente" tentem contornar as restrições da linguagem natural usando representações contínuas, elas preservam em grande parte o paradigma de geração autoregressiva passo a passo. Nessas metodologias, os cálculos intermediários ainda são externalizados sequencialmente antes que o próximo passo possa condicionar-se a eles. Este artigo argumenta que a cognição humana oferece um princípio de design superior: a memória de trabalho. Os humanos podem reter e manipular informações relevantes para a tarefa internamente, sem externalizar cada pensamento intermediário. Os autores postulam que os LLMs podem ser treinados para utilizar um espaço de trabalho interno similar, desacoplando o raciocínio do gargalo sequencial da geração de tokens.

Metodologia: Raciocínio em Memória (RiM)

Os autores introduzem o Raciocínio em Memória (RiM), um método de raciocínio latente que substitui a geração autoregressiva de etapas de raciocínio por blocos de memória fixos.

Mecanismo Central

Em vez de gerar uma sequência de tokens de raciocínio, o modelo recebe $K$ blocos de memória fixos, $m_{1:K}$. Cada bloco consiste em uma sequência predefinida de tokens especiais:
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
onde $<b>$ e $</b>$ delimitam as fronteiras do bloco, e os tokens $<m>$ formam o espaço de trabalho latente.

• Entrada Fixa: Esses blocos são anexados à pergunta de entrada como tokens fixos.
• Passada Única (Single Forward Pass): Como os blocos são entradas fixas e não saídas geradas, toda a sequência aumentada (pergunta + blocos de memória) é processada em uma única passada direta. Isso elimina o gargalo de decodificação sequencial associado à geração autoregressiva.
• Representações Aprendidas: Embora as identidades e posições dos tokens sejam fixas, suas representações contextuais são aprendidas para serem dependentes da tarefa, codificando etapas intermediárias de raciocínio.

Currículo de Duas Etapas

Para ensinar o modelo a usar esses blocos fixos como uma memória de trabalho, o RiM emprega um currículo de treinamento de duas etapas:

Etapa 1: Supervisão de Etapas de Raciocínio

• Objetivo: Fundamentar os blocos de memória atribuindo-lhes um papel computacional.
• Mecanismo: Para um rastro de raciocínio com $T$ etapas, o modelo é treinado com $T$ blocos de memória. Após cada bloco $m_t$, o modelo deve prever a próxima etapa de raciocínio explícita $r_{t+1}$.
• Restrição: Uma máscara de atenção personalizada garante que a leitura para a etapa $t+1$ possa atender à pergunta e aos blocos de memória anteriores, mas não às etapas de raciocínio explícitas anteriores. Isso força o modelo a codificar o cálculo intermediário necessário dentro do espaço de trabalho latente dos blocos de memória.
• Supervisão: Utiliza um objetivo de verossimilhança negativa ponderada com um cronograma de annealing para reduzir gradualmente a força da supervisão.

Etapa 2: Refinamento da Resposta Final

• Objetivo: Converter os blocos de memória fundamentados em um espaço de trabalho latente para a previsão da resposta final.
• Mecanismo: O rastro de raciocínio explícito é descartado. O modelo usa um número fixo de $K$ blocos de memória. Após cada bloco, o modelo é treinado para prever a resposta final diretamente.
• Objetivo: Isso desloca a supervisão da previsão de etapas intermediárias para o refinamento progressivo da resposta final à medida que mais blocos de memória (cálculo latente) se tornam disponíveis.
• Inferência: No momento do teste, o modelo processa a pergunta e os blocos de memória em uma única passada, e a leitura do bloco final fornece a resposta.

Mascaramento de Atenção

Uma máscara de atenção causal personalizada (Figura 2 no artigo) é crítica para o método. Ela previne o vazamento de informações garantindo que:

1. Os blocos de memória atendam à pergunta e aos blocos de memória anteriores.
2. As leituras-alvo (etapas de raciocínio na Etapa 1, resposta final na Etapa 2) atendam à pergunta e aos blocos de memória disponíveis, mas não a outras leituras-alvo.
   Isso permite que todos os alvos sejam previstos em uma única passada direta sem que o modelo "trambique" lendo a verdade fundamental futura.

Contribuições Principais

1. Raciocínio em Memória (RiM): Um novo método de raciocínio latente que desbloqueia a capacidade de memória de trabalho dos LLMs, fornecendo um espaço de trabalho latente dedicado e fixo para cálculo intermediário, desacoplando-o da geração autoregressiva.
2. Currículo de Duas Etapas: Uma abordagem estruturada de treinamento que primeiro fundamenta os blocos de memória via supervisão de etapas explícitas e depois os refina para a previsão da resposta final, permitindo o uso eficaz do espaço latente sem objetivos auxiliares complexos.
3. Eficiência e Desempenho: Evidências empíricas mostrando que o RiM iguala ou supera as melhores linhas de base de raciocínio explícito e latente, mantendo ao mesmo tempo a velocidade de inferência de modelos de resposta direta (sem sobrecarga de geração sequencial).

Resultados Experimentais

Os autores avaliaram o RiM nos benchmarks de raciocínio matemático GSM8K (em distribuição) e GSM-Hard (fora de distribuição) usando os modelos GPT-2 e Llama-3.2 (1B e 3B parâmetros).

• Desempenho: O RiM superou consistentemente a linha de base de raciocínio latente mais forte (Coconut) e o Ajuste Fino Supervisionado (SFT) de resposta direta.
  • No GSM8K (Llama-3.2-1B), o RiM alcançou 42,1% de precisão gananciosa, comparado a 36,9% para o Coconut e 23,9% para o SFT sem CoT.
  • No GSM-Hard, o RiM alcançou 10,5%, superando o Coconut (8,5%) e o SFT sem CoT (5,3%).
  • O RiM também mostrou pontuações Pass@8 superiores, indicando que soluções corretas estão distribuídas ao longo da trajetória latente com alta probabilidade.
• Latência: O RiM alcançou velocidades de inferência comparáveis ao SFT sem CoT (Tempo até o Primeiro Token ~16ms para Llama-3.2-1B), enquanto o Coconut foi ~7x mais lento e o SFT com CoT foi ~27x mais lento devido à geração autoregressiva.
• Análise de Representação: Projeções de PCA das representações dos blocos de memória mostraram que o modelo aprende a criar espaços de trabalho latentes distintos e dependentes da amostra. As representações evoluem sistematicamente durante o treinamento, confirmando que os blocos não são meramente espaços reservados fixos, mas estados computacionais ativos.
• Robustez: O desempenho permaneceu estável em diferentes orçamentos de memória no momento da inferência (número de blocos), sugerindo que o método é robusto a restrições de implantação.

Significado e Alegações

O artigo afirma que LLMs podem ser treinados para usar a memória de trabalho como um mecanismo eficaz para raciocínio latente.

O significado principal reside no desacoplamento do cálculo da comunicação. Ao substituir a geração de pensamentos intermediários por blocos de memória fixos, o RiM permite que os modelos realizem cálculos internos densos e paralelizáveis sem a sobrecarga de gerar linguagem natural ou representações contínuas sequenciais. Isso resulta em um método que é:

1. Computacionalmente Eficiente: Evita o gargalo sequencial da geração autoregressiva, permitindo inferência em passagem única.
2. Eficaz: Alcança desempenho de ponta entre métodos de raciocínio latente em benchmarks matemáticos.
3. Escalável: Funciona em diferentes famílias e tamanhos de modelos sem exigir arquiteturas recorrentes complexas ou objetivos de treinamento de múltiplos caminhos.

Os autores concluem que essa abordagem demonstra um caminho viável para sistemas de raciocínio mais eficientes que mimetizam a eficiência cognitiva humana ao internalizar o processo de raciocínio em vez de externalizá-lo como texto.