Latent Reasoning in TRMs is Secretly a Policy Improvement Operator

Resumo Técnico: O Raciocínio Latente em TRMs é Secretamente um Operador de Melhoria de Política

1. Definição do Problema

Modelos pequenos recentes que utilizam recursão latente, como os Tiny Recursive Models (TRMs) e Hierarchical Reasoning Models (HRMs), demonstraram resultados promissores em tarefas de raciocínio complexas, como o ARC-AGI. A teoria predominante sugere que esses modelos alcançam ganhos de desempenho ao aumentar a profundidade efetiva da rede, permitindo que emulem modelos maiores com menos parâmetros. No entanto, observações empíricas revelam uma discrepância: camadas adicionadas recursivamente frequentemente apresentam desempenho inferior a modelos de passagem única (one-pass) de profundidade feed-forward equivalente. Isso indica que, em arquiteturas de loop, nem todo passo recursivo contribui significativamente para a solução, resultando em "computação morta" (dead compute).

As questões centrais abordadas por este trabalho são:

• Quando e por que o raciocínio latente melhora o desempenho?
• Quando ele resulta em computação redundante?
• Como o treinamento desses modelos pode ser melhorado para garantir que cada passo seja eficaz?

Os métodos de treinamento atuais para TRMs aplicam supervisão primariamente à saída final. Isso deixa os passos de refinamento iniciais com uma atribuição de crédito (credit assignment) fraca, tornando a otimização instável, particularmente em regimes de poucos dados.

2. Metodologia

Os autores propõem um reenquadramento teórico do raciocínio latente e uma metodologia de treinamento correspondente.

Formalização Teórica: Raciocínio Latente como Melhoria de Política
O artigo formaliza o raciocínio recursivo latente como um operador de melhoria de política.

• Política de Referência ($\hat{\pi}_t$): A distribuição de saída do modelo antes de um passo de raciocínio latente (a "primeira tentativa").
• Política Melhorada ($\pi^+_t$): A distribuição de saída após uma atualização de raciocínio latente.
• Sinal de Vantagem: A diferença entre os log-probabilidades da política melhorada e da política de referência atua como um sinal de vantagem implícito ($A_t(s, a) = \log \pi^+_t(a|s) - \log \hat{\pi}_t(a|s)$).
• Mecanismo: Os autores demonstram que um único passo de TRM implementa uma atualização de política que repondera a política de referência com base em uma condição de otimalidade. Especificamente, o passo de raciocínio latente se comporta como uma política condicionada ao evento de que a resposta é melhorada. Isso conecta o algoritmo de espaço latente ao conhecido framework de melhoria de política em Aprendizado por Reforço (RL) e ao Classifier-Free Guidance (CFG) em modelos de difusão.

Esquema de Treinamento Proposto: Supervisão de Melhoria Profunda (DIS)
Baseando-se na percepção de melhoria de política, os autores introduzem a Supervisão de Melhoria Profunda (Deep Improvement Supervision - DIS) para eliminar passos de computação morta.

• Alvos Passo a Passo: Em vez de supervisionar apenas a saída final, o DIS fornece alvos intermediários $\{y^\dagger_s\}$ para cada passo de recursão $s$. Esses alvos são gerados via um esquema de corrupção discreta monotônica da saída real $y^*$.
• Objetivo de Treinamento: O modelo é treinado para mapear o alvo anterior $y^\dagger_{s-1}$ (via política de referência) para o próximo alvo menos corrompido $y^\dagger_s$ (via política melhorada).
• Função de Perda: O treinamento minimiza uma perda de entropia cruzada entre os logits pós-raciocínio e o alvo específico do passo $y^\dagger_s$. Isso força a vantagem da razão de log em cada passo a ser positiva para o token correto, garantindo que cada passo de recursão realize uma melhoria direcionada em vez de computação redundante.

Simplificações Arquiteturais
O método proposto simplifica a arquitetura TRM ao:

• Remover a necessidade de um mecanismo de interrupção aprendido (Adaptive Computation Time).
• Reduzir o número de passos de supervisão e de passos de raciocínio latente.
• Utilizar um número fixo de passos de recursão ($N_{sup}$) determinado pelo esquema de corrupção.

3. Principais Contribuições

1. Insight Teórico: O artigo estabelece que o raciocínio latente em TRMs é matematicamente equivalente a um operador de melhoria de política. Ele faz a ponte entre algoritmos de raciocínio em espaço latente e frameworks de RL, explicando o comportamento do modelo através da lente da maximização de vantagem.
2. Supervisão de Melhoria Profunda (DIS): Um novo esquema de treinamento que fornece alvos intermediários passo a passo. Isso transforma o problema difícil de atribuição de crédito de longo prazo em uma tarefa de aprendizado supervisionado tratável, onde cada passo é explicitamente treinado para melhorar o estado anterior.
3. Ganhos de Eficiência: O método alcança desempenho competitivo enquanto reduz significativamente o overhead computacional. Ele evita o treinamento de módulos de interrupção e reduz o número total de passagens (forward passes) em 18× em comparação com configurações padrão de TRM, mantendo ou melhorando a precisão.

4. Resultados Experimentais

Os autores avaliaram seu método nos problemas N-Queens e nos benchmarks ARC-AGI (ARC-AGI-1 e ARC-AGI-2).

• N-Queens: Usando um modelo compacto (0.8M de parâmetros), o método DIS alcançou 0.69 de precisão, igualando o desempenho de um TRM padrão com uma cabeça de interrupção (halting head) e 16 passos de supervisão, mas utilizando significativamente menos passos de inferência ($T=1, n=2$ vs. $T=3, n=6$).
• ARC-AGI:
  • Modelo Compacto (0.8M de parâmetros): O DIS alcançou 24.0% de precisão no ARC-AGI-1, superando o TRM-compact original (12.0%) e a maioria dos LLMs de código aberto existentes sem conhecimento externo.
  • Modelo Médio (7M de parâmetros): O DIS alcançou 41.3% no ARC-AGI-1 e 6.0% no ARC-AGI-2, superando o TRM-medium padrão (40.4% / 3.3%) e igualando o desempenho do TRM original (40.4% / 3.3%) enquanto utiliza menos passos de raciocínio latente.
• Redução de Computação: O método reduziu o número total de passagens (forward passes) por um fator de 18 em comparação com a configuração baseline do TRM, demonstrando que a supervisão explícita de passos intermediários elimina a necessidade de loops excessivos de raciocínio latente.

5. Significância e Alegações

O artigo afirma que, ao visualizar o raciocínio latente como um processo de melhoria de política, é possível explicar por que modelos recursivos padrão às vezes falham (devido à "computação morta", onde os passos não melhoram a margem de vantagem) e como corrigir isso.

A principal significância reside no framework de Supervisão de Melhoria Profunda, que:

• Fornece uma regra de treinamento construtiva para garantir que cada passo de recursão realize uma melhoria não trivial.
• Simplifica o treinamento de modelos recursivos ao remover a necessidade de mecanismos de interrupção complexos e de atribuição de crédito instáveis.
• Demonstra que modelos pequenos e eficientes em parâmetros podem alcançar resultados de estado da arte em tarefas de raciocínio complexas (como o ARC-AGI) quando o objetivo de treinamento está alinhado com o mecanismo teórico de melhoria de política.

Os autores concluem que sua abordagem oferece um caminho para modelos de raciocínio mais eficientes e eficazes, alinhando o sinal de treinamento com a natureza algorítmica intrínseca da recursão latente.