From Next Token Prediction to (STRIPS) World Models

Este artigo investiga se a previsão do próximo token pode gerar modelos de mundo capazes de planejamento em ambientes simbólicos, demonstrando que, embora o modelo "STRIPS Transformer" tenha um viés indutivo forte, um transformer padrão com atenção de quebra de vara (stick-breaking) alcança maior precisão e generalização, permitindo a extração de modelos STRIPS válidos para planejamento em estados e objetivos não vistos.

O essencial

Imagine que você está tentando ensinar um robô a jogar xadrez ou a organizar uma sala de brinquedos. O objetivo não é apenas fazer o robô memorizar movimentos, mas fazê-lo entender as regras do jogo para que ele possa planejar suas jogadas futuras, mesmo em situações que nunca viu antes.

Este artigo de pesquisa investiga se as Inteligências Artificiais modernas (os famosos "Modelos de Linguagem" ou LLMs) conseguem aprender essas regras internas apenas observando sequências de ações, sem que ninguém lhes explique a teoria por trás.

Aqui está uma explicação simples, usando analogias do dia a dia:

1. O Problema: O "Adivinhador" vs. O "Planejador"

Muitos modelos de IA atuais são como adivinhos incríveis. Se você mostrar a eles as primeiras palavras de uma frase, eles sabem exatamente qual é a próxima. Eles são ótimos em prever o futuro imediato baseado em padrões.

Mas será que eles realmente entendem como o mundo funciona? Se você pedir para eles planejar uma rota complexa de A até Z, eles muitas vezes falham, porque estão apenas "chutando" a próxima palavra, sem ter um mapa mental do mundo.

Os autores deste trabalho queriam saber: Se ensinarmos uma IA a prever a próxima ação em um mundo lógico (como um jogo de tabuleiro), ela vai aprender a criar um "mapa mental" (um modelo do mundo) que permite planejar?

2. O Cenário de Teste: O "Mundo de Blocos"

Para testar isso, eles usaram um ambiente controlado chamado STRIPS. Pense nisso como um jogo de blocos de montar muito lógico:

• Você tem blocos (fatos).
• Você tem ações (pegar um bloco, empilhar, soltar).
• Cada ação tem regras estritas: "Só posso empilhar o bloco A no B se o B estiver livre".

O desafio era: dar à IA apenas uma lista de movimentos que funcionaram (e alguns que não funcionaram) e ver se ela conseguia deduzir as regras do jogo para depois resolver novos problemas.

3. As Duas Soluções Criadas

Os pesquisadores criaram duas "arquiteturas" (tipos de cérebros artificiais) para tentar aprender essas regras:

A. O "STRIPS Transformer" (O Arquiteto Rigoroso)

Imagine um engenheiro que constrói uma casa com um manual de instruções exato.

• Este modelo foi desenhado com "vieses simbólicos". Isso significa que sua estrutura interna já sabe que o mundo é feito de "fatos" e "regras". Ele tenta mapear cada parte do cérebro para uma regra específica do jogo.
• Resultado: Ele é muito preciso quando consegue aprender, mas é difícil de treinar. É como tentar ensinar um engenheiro rigoroso a desenhar; se você der poucos exemplos, ele fica confuso e não aprende bem. Ele precisa de muitos dados para funcionar.

B. O "Stick-Breaking Transformer" (O Aprendiz Intuitivo)

Imagine um criança muito inteligente que aprende observando e quebrando o problema em pedaços.

• Este é um modelo padrão (como os que usamos em chats hoje), mas com uma "mágica" chamada atenção stick-breaking.
• A Analogia do "Stick-Breaking" (Quebrar o Palito): Imagine que você tem um palito de picolé inteiro. Você precisa escolher a parte mais recente e importante para olhar. Em vez de olhar para tudo de uma vez (o que confunde a IA em sequências longas), este método "quebra" o palito, focando apenas no pedaço mais recente e relevante que ainda não foi descartado. Isso ajuda a IA a lembrar o que aconteceu recentemente e ignorar o que já não importa mais.
• Resultado: Surpreendentemente, este modelo "genérico" aprendeu muito melhor e mais rápido que o modelo rigoroso. Ele alcançou quase 100% de precisão e generalizou muito bem para problemas novos.

4. A Grande Descoberta: O "Mapa" Escondido

O resultado mais impressionante foi o seguinte:

1. Eles treinaram essas IAs apenas para prever a próxima ação.
2. Depois, eles "extrairam" as regras que a IA aprendeu internamente.
3. O Milagre: As regras extraídas eram tão perfeitas que podiam ser usadas por planejadores clássicos (softwares antigos e confiáveis de robótica) para resolver problemas que a IA nunca viu antes.

Isso significa que a IA, ao tentar apenas "adivinhar a próxima palavra", acabou construindo um modelo do mundo interno que era matematicamente correto. Ela aprendeu a lógica do jogo, não apenas a memorizar movimentos.

5. Por que isso importa?

• Generalização: As IAs conseguiram resolver problemas com milhões de combinações diferentes de cenários (estados iniciais e metas), mesmo tendo sido treinadas em poucos exemplos. É como se você ensinasse uma criança a andar de bicicleta em um parque e ela conseguisse andar em qualquer rua do mundo.
• O Poder da Simplicidade: O modelo mais simples (o "Stick-Breaking"), sem regras complexas embutidas, funcionou melhor que o modelo complexo. Isso sugere que, às vezes, dar à IA a liberdade de aprender os padrões sozinha é mais eficaz do que forçá-la a seguir uma estrutura rígida.
• O Perigo da Memória: Modelos de IA comuns (sem a técnica de "stick-breaking") falharam em sequências longas. Eles esqueciam o que aconteceu no início da frase. A técnica nova foi essencial para que eles mantivessem o foco no longo prazo.

Resumo Final

Este artigo prova que, em ambientes lógicos, prever o futuro (próxima ação) é suficiente para aprender a entender o presente (o modelo do mundo).

Os autores mostraram que, ao usar a técnica certa de "atenção" (o método de quebrar o palito), podemos transformar modelos de linguagem em planejadores robóticos reais, capazes de entender regras complexas e resolver problemas novos sem precisar ser reprogramados para cada situação. É um passo gigante para criar IAs que realmente "pensam" e planejam, em vez de apenas repetir o que ouviram.

Resumo técnico

Aqui está um resumo técnico detalhado do artigo "From Next Token Prediction to (STRIPS) World Models", apresentado em português:

Título: De Previsão do Próximo Token para Modelos de Mundo (STRIPS)

Autores: Carlos Núñez-Molina, Vicenç Gómez, Hector Geffner.
Afilições: RWTH Aachen University (Alemanha) e Universitat Pompeu Fabra (Espanha).

---

1. Problema Investigado

O trabalho investiga se a tarefa de previsão do próximo token (realizada por modelos de linguagem autoregressivos, como Transformers) é suficiente para aprender modelos de mundo que suportem efetivamente o planejamento.

O desafio central é determinar se os Transformers aprendem representações latentes que capturam a dinâmica causal e a estrutura do ambiente (o "modelo de mundo") ou se apenas memorizam regularidades estatísticas superficiais. O foco é específico: aprender modelos de ação STRIPS (um padrão clássico em IA para planejamento simbólico) a partir de traços de ação (sequências de ações aplicáveis e não aplicáveis), sem acesso direto aos estados do mundo durante o treinamento, apenas para inferir as precondições e efeitos das ações.

2. Metodologia e Arquiteturas Propostas

Os autores propõem e comparam duas arquiteturas de Transformer para aprender modelos STRIPS a partir de traços de ações (positivos e negativos):

A. STRIPS Transformer

• Conceito: Uma arquitetura alinhada simbolicamente, baseada em resultados teóricos que conectam Transformers de atenção "hard" (rígida) e a linguagem formal B-RASP às linguagens "star-free" (livres de estrela), que descrevem a estrutura de traços válidos em domínios STRIPS.
• Funcionamento:
  • Cada cabeça de atenção é dedicada a um átomo (variável booleana) do domínio.
  • Utiliza atenção hard mascarada (implementada via stick-breaking attention) para identificar, para cada ação e átomo, a ação mais recente que afeta aquele átomo.
  • A lógica de aplicabilidade é codificada explicitamente nos parâmetros: se o átomo é uma precondição, a atenção verifica se a ação mais recente que afetou o átomo o adicionou (tornando-o verdadeiro) ou o deletou (tornando-o falso).
  • Possui um viés indutivo simbólico forte, mapeando diretamente parâmetros para precondições e efeitos.

B. Stick-Breaking (SB) Transformer

• Conceito: Uma arquitetura padrão de Transformer (estilo decoder), mas com modificações cruciais para capturar a estrutura sequencial necessária.
• Diferenças Chave:
  • Não possui estrutura simbólica explícita embutida (não mapeia cabeças para átomos).
  • Substitui a atenção softmax padrão e os positional encodings por atenção stick-breaking.
  • A atenção stick-breaking é uma mecanismo de normalização sequencial diferenciável que aproxima a atenção hard, concentrando-se no predecessor mais recente com alta pontuação, preservando o viés de recência essencial para rastrear o estado de variáveis ao longo do tempo.

Processo de Extração e Planejamento

Ambos os modelos são treinados para classificar traços como positivos (sequência de ações aplicável) ou negativos. Após o treinamento:

1. Extração do Modelo: Um modelo STRIPS simbólico ($M' = \langle F', A' \rangle$) é extraído dos parâmetros aprendidos.
   • No STRIPS Transformer, a extração é direta via binarização dos parâmetros.
   • No SB Transformer, utiliza-se um processo de "sondagem de estado" (state probing) onde o modelo prevê a aplicabilidade de ações de teste (test-p) para reconstruir o estado e inferir precondições/effects por consenso majoritário.
2. Planejamento: O modelo simbólico extraído é alimentado em planejadores clássicos off-the-shelf (como o Mimir com heurística FF) para resolver problemas de planejamento com estados iniciais e objetivos nunca vistos durante o treinamento.

3. Contribuições Principais

1. Validação de Modelos de Mundo: Demonstra empiricamente que a previsão do próximo token em Transformers pode, de fato, levar à aprendizagem de modelos de mundo simbólicos completos que suportam planejamento.
2. Generalização Combinatória: Os modelos aprendidos generalizam para um número exponencial de estados iniciais e objetivos não vistos, provando que o modelo aprendeu a dinâmica do domínio e não apenas memorizou traços.
3. Comparação de Arquiteturas:
   • O SB Transformer (padrão com atenção stick-breaking) supera o STRIPS Transformer em facilidade de otimização e generalização, alcançando precisão quase perfeita.
   • O STRIPS Transformer, apesar de ter um viés simbólico forte, é mais difícil de otimizar e requer conjuntos de dados maiores para generalizar de forma confiável.
4. Papel da Atenção Stick-Breaking: Evidencia que a atenção stick-breaking é essencial para a generalização em sequências longas. Transformers padrão (com softmax) falham em generalizar para traços longos, embora ainda consigam extrair modelos simbólicos funcionais se treinados em traços curtos.

4. Resultados Experimentais

Os experimentos foram realizados em cinco domínios clássicos de planejamento: Blocksworld, Ferry, N-puzzle, Maze e Logistics, com tamanhos variados (pequenos e grandes).

• Precisão de Previsão:
  • O SB Transformer alcançou precisão de treinamento e teste próxima de 100% em todos os domínios, mesmo em traços longos (até 400 passos).
  • Transformers padrão (baselines com softmax e RoPE) alcançaram alta precisão no treinamento, mas colapsaram na generalização para traços longos (precisão de teste caindo drasticamente, ex: de 0,99 para 0,23 no Blocksworld).
• Desempenho de Planejamento:
  • Modelos extraídos do SB Transformer permitiram que planejadores clássicos resolvessem problemas com 100% de sucesso em todos os domínios e tamanhos, mesmo para estados iniciais e objetivos não vistos.
  • Modelos extraídos de Transformers padrão treinados em traços curtos também funcionaram bem, mas falharam quando extraídos de traços longos devido à incapacidade de generalizar a dinâmica.
  • O STRIPS Transformer teve desempenho variável e, em alguns casos, inferior ao SB Transformer, devido a dificuldades de otimização (estagnação em ótimos locais), embora quando bem-sucedido, também permitisse planejamento perfeito.
• Generalização Composicional: Ambos os modelos demonstraram capacidade de raciocínio composicional, combinando ações aprendidas para resolver problemas complexos não vistos no conjunto de treinamento.

5. Significado e Conclusão

O artigo fornece evidências fortes de que Transformers podem internalizar modelos de mundo simbólicos através da simples previsão do próximo token, desde que a arquitetura seja adequada (uso de atenção stick-breaking) e o domínio seja estruturado (STRIPS).

• Implicação Teórica: Refuta a ideia de que LLMs apenas aprendem estatísticas superficiais; neste cenário controlado, eles aprendem a lógica causal subjacente.
• Implicação Prática: Sugere que é possível treinar modelos de linguagem para extrair modelos de planejamento simbólicos robustos, permitindo o uso de planejadores clássicos eficientes para resolver problemas complexos, superando as limitações de planejamento direto via LLMs (que muitas vezes falham em consistência lógica).
• Futuro: O trabalho abre caminho para aprender modelos STRIPS "levantados" (lifted) diretamente de traços, estendendo a abordagem para domínios mais complexos e escaláveis.

Em resumo, o estudo demonstra que a combinação de previsão de próximo token com atenção stick-breaking é uma via eficaz para transformar redes neurais em modelos de mundo simbólicos capazes de planejamento rigoroso e generalização exponencial.