Resumo Técnico: Aprendizado de Rastreamento de Estado a partir de Código usando RNNs Lineares

Declaração do Problema

O rastreamento de estado é uma capacidade fundamental para compreender a execução de programas, estados de jogos e modelos de mundo. Embora o trabalho teórico tenha estabelecido uma divisão onde os Transformers se destacam na recuperação associativa e as Redes Neurais Recorrentes (RNNs) desempenham bem o rastreamento de estado, existe uma lacuna significativa em como esses modelos são avaliados versus como são treinados.

A maioria dos benchmarks existentes para rastreamento de estado (ex: composição de permutações) utiliza uma configuração de sequência-para-sequência: mapear uma sequência de ações diretamente para uma sequência de estados. Isso difere fundamentalmente do paradigma de predição do próximo token utilizado para treinar os Large Language Models (LLMs) modernos, onde o modelo deve prever o próximo token em um fluxo que intercala ações com revelações parciais de estado (ex: instruções print). Além disso, a execução de código no mundo real frequentemente envolve observabilidade parcial e transições probabilísticas (ex: ramificações baseadas em estados ocultos ou escolhas aleatórias), criando um cenário onde o estado real não é totalmente observável, exigindo a manutenção de uma distribuição de crença em vez de um estado determinístico.

Metodologia

1. Reformulação do Rastreamento de Estado via Código (Traços de REPL)

Para preencher a lacuna entre benchmarks abstratos de autômatos e objetivos de treinamento de LLMs, os autores convertem tarefas de composição de permutação em traços de REPL Python.

• Mecanismo: Em vez de mapear uma sequência de permutações $\sigma_1, \dots, \sigma_m$ para estados cumulativos, o modelo processa uma sequência de comandos Python (atribuições de variáveis, trocas) intercalados com instruções print que revelam informações parciais de estado.
• Objetivo: O modelo é treinado via predição do próximo token para prever o próximo token no traço, que pode ser um comando ou um valor de estado revelado. Isso força o modelo a aprender o rastreamento de estado implicitamente através da tarefa de predição.

2. Comparação Arquitetural

O estudo compara três classes arquiteturais nessas tarefas:

• Transformers: Modelos padrão baseados em atenção.
• RNNs Lineares (DeltaNet [0, 1]): RNNs lineares restritas a autovalores em $[0, 1]$, conhecidas por serem incapazes de realizar rastreamento de estado complexo como paridade ou composição de permutação.
• RNNs Lineares (DeltaNet [-1, 1]): RNNs lineares com intervalos de autovalores estendidos $[-1, 1]$, teoricamente capazes de resolver composição de permutação determinística.

3. Estrutura Teórica: PFSA-SR

Para analisar a dificuldade do rastreamento de estado em código real, os autores introduzem o Autômato de Estados Finitos Probabilístico com Revelações de Estado (PFSA-SR).

• Definição: Uma tupla $(Q, \Sigma, \delta, \rho, q_0)$ onde $\delta$ é um kernel de transição probabilística e $\rho$ é uma função de revelação que mapeia entradas para subconjuntos de estados.
• Dinâmica: O sistema evolui via transições probabilísticas, mas o observador recebe restrições "rígidas" (revelações) que podam o conjunto de estados possíveis (poda de suporte) em vez de fornecer probabilidades suaves (soft likelihoods).
• Atualização de Crença: O rastreamento requer a manutenção de uma distribuição de crença $b_t$ sobre os estados. A atualização envolve um passo de revelação (zerar estados inconsistentes e renormalizar) seguido por um passo de transição. Crucialmente, o passo de renormalização $f(x) = x/\|x\|_1$ é não linear.

4. Análise de Representação

O artigo investiga duas representações para estados de crença em RNNs lineares:

• Representação Conjunta: Enumera explicitamente todas as $n!$ permutações. Os autores argumentam que isso é numericamente instável para RNNs lineares porque o termo de injeção de entrada $B(x_t)$ não consegue repor adaptativamente a massa de probabilidade das configurações sobreviventes após uma revelação parcial, levando ao decaimento exponencial da norma.
• Representação Marginal: Rastreia a probabilidade marginal de cada posição independentemente (uma matriz duplamente estocástica $n \times n$). Embora isso permita que $B(x_t)$ injete diretamente certeza para as entradas reveladas, os autores mostram que ainda é suscetível ao decaimento exponencial sob sequências de revelação adversariais, onde entradas específicas não reveladas nunca são diretamente repostas.

Principais Resultados

1. Densidade de Supervisão e Arquitetura

• Transformers: Requerem supervisão densa de estado para aprender o rastreamento de estado. À medida que o espaçamento entre as instruções print (revelações) aumenta, o desempenho do Transformer cai rapidamente. Eles falham em generalizar para permutações completas mesmo com supervisão densa.
• RNNs Lineares (DeltaNet [-1, 1]): Destacam-se neste cenário. Elas aprendem a rastrear estados de forma confiável mesmo com supervisão esparsa e demonstram extrapolação de comprimento robusta (generalizando para sequências mais longas que o treinamento).
• RNNs Lineares (DeltaNet [0, 1]): Falham em resolver a tarefa, confirmando que o intervalo estendido de autovalores $[-1, 1]$ é necessário para o rastreamento de estado complexo.

2. Interpretabilidade do Rastreamento de Estado

A análise do DeltaNet [-1, 1] treinado revela que o rastreamento de estado é localizado em componentes específicos:

• Extensão de Autovalor: Pelo menos uma cabeça de atenção em uma camada específica (ex: Camada 12, Cabeça 3) aprende a definir seu parâmetro de decaimento $\beta$ como 2, resultando em um autovalor de $-1$.
• Estrutura de Subespaço: Esta cabeça específica opera em um subespaço linear significativamente de menor dimensão comparado a outras cabeças, consistente com os requisitos teóricos para rastreamento de permutações.
• Intervenção: Escalar o intervalo de $\beta$ desta cabeça específica de volta para $[0, 1]$ causa colapso imediato de desempenho, enquanto escalar outras cabeças não tem efeito, confirmando esta cabeça como o mecanismo dedicado de rastreamento de estado.

3. Limites de RNNs Lineares em Cenários Probabilísticos

O artigo demonstra que, embora as RNNs lineares resolvam tarefas determinísticas, elas enfrentam barreiras inerentes em cenários probabilísticos (PFSA-SR):

• Decaimento Exponencial: Sob a representação conjunta, revelações parciais fazem com que a norma do estado não normalizado decaia exponencialmente ($1/2 \to 1/4 \to 1/8 \dots$) porque a recorrência linear não pode realizar a renormalização não linear necessária.
• Vulnerabilidade Adversarial: Mesmo com a representação marginal, um adversário pode construir uma sequência de trocas probabilísticas e revelações seletivas que faz com que a massa de probabilidade das entradas não reveladas decaia exponencialmente até desaparecer em precisão finita.
• Conclusão: A manutenção estável de crença em cenários probabilísticos com observabilidade parcial parece inviável para RNNs lineares sem uma renormalização não linear na recorrência ou resets frequentes do estado completo.

Significância e Alegações

O artigo afirma fornecer um ambiente de teste mais realista para o rastreamento de estado que se alinha ao objetivo de predição do próximo token de modelos de linguagem, indo além de benchmarks sintéticos de sequência-para-sequência.

1. Preenchendo a Lacuna: Demonstra que RNNs lineares capazes de rastreamento de estado (especificamente DeltaNet com autovalores estendidos) podem aprender essas capacidades via predição do próximo token em traços semelhantes a código, enquanto os Transformers não conseguem.
2. Barreira Teórica: Identifica uma limitação fundamental das RNNs lineares em cenários de código do mundo real. O código real frequentemente envolve transições probabilísticas e observabilidade parcial (ex: ramificações ocultas, escolhas aleatórias). Os autores mostram que rastrear estados em tais ambientes exige renormalização não linear, algo que as RNNs lineares não podem realizar internamente. Isso sugere que a "lacuna de complexidade" observada em tarefas do mundo real pode advir da natureza probabilística da execução de código, e não apenas do comprimento ou complexidade da sequência.
3. Implicações para Benchmarking: Os autores sugerem que avaliar arquiteturas apenas em tarefas de Autômatos de Estados Finitos Determinísticos (DFA) obscurece o potencial de RNNs não lineares e as limitações de RNNs lineares. O framework PFSA-SR proposto e seus modos de falha associados (decaimento exponencial da norma) oferecem uma explicação teórica para a qual as RNNs lineares padrão lutam quando o raciocínio de incerteza é necessário.

O artigo conclui modestamente, observando que seus experimentos focam em grupos de permutação sintéticos e que trabalhos futuros são necessários para estender esses achados a traços de execução do mundo real e para investigar arquiteturas híbridas que intercalam recorrência linear com normalização não linear periódica.