🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

Este artigo demonstra que a aplicação de arquiteturas em laço a modelos de espaço de estados baseados em Mamba aumenta as capacidades de raciocínio e a eficiência de parâmetros através do aumento da profundidade computacional, ao mesmo tempo que introduz a seleção adaptativa de estado de saída para otimizar a profundidade de predição, embora observe que a economia real de tempo de inferência requer mecanismos adicionais de manipulação de estado.

Autores originais: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Publicado 2026-07-14

📖 6 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem um robô chef brilhante e super inteligente chamado Mamba. Este chef é famoso por ser incrivelmente rápido e eficiente em termos de memória, sendo capaz de preparar uma tempestade de pratos sem precisar de uma cozinha enorme cheia de ferramentas diferentes. Mas há um porém: quando confrontado com uma receita realmente difícil — como um quebra-cabeça matemático complexo ou um jogo de "encontrar o ingrediente oculto" — o Mamba às vezes fica travado. Ele precisa pensar mais profundamente, mas não tem "camadas" suficientes de pensamento para fazer isso sem construir uma cozinha nova e gigante (o que custa uma fortuna).

Entram em cena os pesquisadores da Universidade Rikkyo e da Universidade de Tóquio. Eles fizeram uma pergunta simples e ousada: E se não construíssemos uma cozinha maior, mas em vez disso fizéssemos o Mamba cozinhar o mesmo prato repetidamente, refinando o sabor a cada passagem?

A Cozinha "Em Loop": Um Chef, Muitas Passagens

No mundo da IA, a maioria dos modelos é como uma linha de montagem de fábrica. Um dado (como uma palavra em uma frase) move-se ao longo da linha, passando por 24 estações diferentes (camadas) e, então, está pronto. Se você quiser que o robô pense com mais intensidade, geralmente você apenas adiciona mais estações.

Os autores tentaram algo diferente. Eles construíram um Mamba em Loop (Looped Mamba). Imagine que o Mamba é uma estação única e super talentosa. Em vez de enviar o dado por uma longa linha de 24 estações diferentes, eles o enviam de volta para a mesma estação, 24 vezes seguidas.

A Magia: O robô usa o exato mesmo cérebro (parâmetros) para cada passagem. É como um estudante relendo o mesmo parágrafo de um livro didático quatro vezes para entendê-lo, em vez de ler quatro parágrafos diferentes e confusos.
O Resultado: Na lógica complexa chamada Mano (aritmética modular) e p-hop induction (encontrar padrões ocultos), essa estratégia de "releitura" funcionou maravilhas. Um modelo que tinha apenas 4 camadas, mas que as loopou 6 vezes (4 ⊗ 6), esmagou um modelo padrão com 24 camadas únicas. Isso sugere que, para o raciocínio, a profundidade do pensamento importa mais do que o número de ferramentas únicas.

O "Portão de Saída": Escolhendo a Melhor Versão

Mas espere, e se o robô continuar relendo o parágrafo para sempre? Isso é um desperdício de tempo. Os pesquisadores adicionaram um recurso inteligente chamado Portão de Saída (Exit Gate).

Pense neste portão como um selecionador inteligente. Enquanto o Mamba processa o dado, o portão pergunta: "A resposta está ficando mais clara?"

Se a resposta já estiver clara após 2 loops, o portão diz: "Pare! Usaremos o resultado desta passagem."
Se o problema for difícil, o portão diz: "Continue!" e seleciona o resultado da 3ª ou 4ª passagem.

Isso é chamado de Seleção de Estado de Saída Adaptativa (Adaptive Exit-State Selection). O artigo descobriu que, para modelos menores (cerca de 140 milhões de parâmetros), este portão foi um divisor de águas. Ele permitiu que o modelo escolhesse o tempo de pensamento "na medida certa" para cada palavra específica, muitas vezes performando melhor do que forçar o robô a sempre usar o resultado do número máximo de loops.

No entanto, há uma reviravolta. Os pesquisadores foram muito cuidadosos ao notar que isso ainda não economiza eletricidade ou tempo no computador. Embora o portão selecione um resultado anterior, o robô ainda fisicamente executa todos os loops em segundo plano. Isso ocorre porque o estado de memória do Mamba é contínuo: o resultado do loop 3 depende do estado deixado pelo loop 2, que depende do loop 1. Você não pode simplesmente "derrubar" o robô das passagens posteriores sem quebrar a corrente para todas as palavras subsequentes. Portanto, o portão seleciona a profundidade de predição (qual versão da resposta usar), mas a computação de tempo real (wall-clock computation) permanece a mesma. Para realmente economizar tempo, seria necessário um novo modo de lidar com o estado de memória do robô, o que os autores dizem ser um trabalho para o futuro.

O Choque de Realidade: O Que Não Funcionou (e O Que Ainda é um Talvez)

O artigo é refrescantemente honesto sobre o que ele não resolveu.

Maior nem sempre é melhor (para este truque): Quando compararam seu Mamba em Loop a um modelo padrão não-loopado que era tão "caro" de rodar (mesmo número de cálculos, ou "iso-FLOPs"), o modelo padrão ainda venceu no teste básico de "o quão confuso o modelo está?" (perplexidade). O modelo em loop foi ótimo em tarefas de raciocínio, mas o modelo profundo padrão ainda era ligeiramente melhor em apenas prever a próxima palavra em uma frase.
O tamanho importa: O "Portão de Saída" funcionou lindamente para os modelos de 140 milhões de parâmetros. Mas quando tentaram aplicá-lo nos modelos maiores de 370 milhões de parâmetros, o portão não melhorou muito as coisas. Os modelos maiores pareciam preferir apenas rodar todos os loops de qualquer maneira. Os autores sugerem que modelos maiores podem depender mais fortemente daquela passagem final e profunda de pensamento.
Não é uma solução mágica para tudo: Os pesquisadores testaram isso em quebra-cabeças sintéticos (como árvores matemáticas e busca de padrões). Embora estes provem que o conceito funciona, eles admitem que ainda não sabemos se este truque de "looping" funcionará tão bem em raciocínios de mundo real e abertos, como escrever um romance ou resolver um caso jurídico complexo.

A Conclusão

O artigo sugere que reutilizar o mesmo cérebro inteligente repetidamente é uma forma poderosa de tornar a IA mais inteligente no raciocínio sem construir um cérebro maior e mais caro. É um eixo de escalonamento "eficiente em parâmetros".

Provado: Em quebra-cabeças lógicos específicos, o loop de um modelo Mamba pequeno supera um modelo padrão do mesmo tamanho e iguala um modelo muito maior.
Sugerido: Esta abordagem pode ser a chave para tornar a IA do futuro mais eficiente, especialmente para tarefas que exigem pensamento profundo e passo a passo.
Não Resolvido: Ainda precisamos descobrir como fazer o "Portão de Saída" realmente economizar tempo de computação (em vez de apenas selecionar a saída) e se isso funcionará nos modelos massivos de bilhões de parâmetros que alimentam as maiores IAs de hoje.

Em resumo, os autores não construíram apenas um robô maior; eles ensinaram um robô menor a pensar mais profundamente, permitindo que ele desse uma segunda (ou terceira, ou quarta) olhada no problema. E, por enquanto, essa segunda olhada parece bastante promissora.

Resumo Técnico: Modelos de Linguagem de Espaço de Estados com Loop e Seleção Adaptativa de Estado de Saída

Declaração do Problema

Avanços recentes em Grandes Modelos de Linguagem (LLMs) têm sido impulsionados principalmente pelo escalonamento do tamanho do modelo, dados de treinamento e computação intensiva de inferência. No entanto, esse paradigma de escalonamento introduz desafios significos de implantação, incluindo o aumento do consumo de memória e o crescimento rápido dos custos de inferência devido à manutenção de grandes caches de chave-valor nos mecanismos de autoatenção baseados em Transformer. Embora os Modelos de Espaço de Estados (SSMs), particularmente o Mamba-2, ofereçam uma alternativa eficiente com um escalonamento favorável de contexto longo, permanece incerto se o princípio de "profundidade computacional sobre parâmetros independentes" — demonstrado recentemente em Transformers com loop (Looped Transformers) — aplica-se aos SSMs. Além disso, a pesquisa existente sobre arquiteturas com loop concentrou-se quase exclusivamente em backbones de Transformer, deixando uma lacuna na compreensão de como o compartilhamento recorrente de parâmetros funciona dentro das arquiteturas Mamba e híbridas de Mamba-Transformer.

Metodologia

Os autores propõem as arquiteturas Looped Mamba e Looped Hybrid Mamba–Transformer. O mecanismo central envolve o compartilhamento dos parâmetros de um único bloco Mamba (ou híbrido) através das camadas e sua aplicação repetida por meio de loops explícitos de profundidade finita.

Arquitetura

Notação: Um modelo é denotado como $N \otimes R$ , onde $N$ é o número de camadas parametrizadas independentemente e $R$ é o número de passos recorrentes. A profundidade efetiva é $NR$.
Looped Mamba: Aplica repetidamente um backbone Mamba de $N$ $N$ camadas com parâmetros compartilhados $M_\theta$ $M_{θ}$ .
- $h^{(r)} = M_\theta(h^{(r-1)})$ para $r = 1, \dots, R$ .
- A saída final é derivada do estado oculto no passo $R$ .
Looped Hybrid: Estende o loop para um backbone que contém tanto camadas Mamba-2 quanto um pequeno subconjunto de camadas de autoatenção causal. O loop aplica-se a todo o bloco híbrido, permitindo o refinamento iterativo das representações ocultas através de atualizações de espaço de estado e interações baseadas em atenção.
Protocolo de Treinamento:
1. Pré-treinamento: Otimizado usando um objetivo de loop final (perda calculada apenas no passo $R$ ).
2. Estágio de Portão de Saída (Exit-Gate) 1: Introduz um objetivo de computação adaptativa com regularização de entropia (adaptado de Ouro) para aprender uma distribuição por token sobre profundidades de saída usando as saídas de todos os passos recorrentes.
3. Estágio de Portão de Saída (Exit-Gate) 2: Congela o backbone e o cabeçalho, treinando apenas o portão de saída para incentivar a continuação apenas quando o próximo passo melhora a perda de predição.

Seleção Adaptativa de Estado de Saída

O artigo adapta o portão de saída de dois estágios de Ouro para o Mamba. Na inferência, um limiar $q$ controla a seleção do passo de saída $z_i$ para cada token com base na probabilidade de saída cumulativa.

Nota Semântica: Os autores esclarecem que, em SSMs, a cadeia de estados de todos os $R$ passos deve ser preservada devido às restrições de continuidade de estado. Portanto, a "saída" seleciona qual estado oculto do passo recorrente é passado para o cabeçalho de modelagem de linguagem; isso não reduz inerentemente o tempo de computação (wall-clock) ou o número de passos executados sem mecanismos adicionais de manipulação de estado.

Principais Contribuições

Estudo Empírico Controlado: Os autores apresentam uma comparação rigorosa dos modelos Mamba e modelos híbridos sob protocolos de iso-parâmetros (mesmo $N$ , variando $R$ ) e iso-FLOPs (mesma profundidade efetiva $NR$) contra baselines de Mamba e Transformer não recorrentes (não-looped).
Avaliação de Tarefas de Raciocínio: O estudo valida a abordagem em tarefas de raciocínio sintético (Mano e indução $p$ -hop) e pré-treinamento de modelos de linguagem com benchmarks de tarefas subsequentes.
Mecanismo de Saída Adaptativa: A introdução de um mecanismo de seleção de estado de saída controlado por limiar para modelos de espaço de estado com loop, demonstrando sua capacidade de melhorar o desempenho em tarefas subsequentes em profundidades intermediárias.

Resultados Experimentais

Tarefas de Raciocínio Sintético

Tarefa Mano (Aritmética Modular): O Looped Mamba supera consistentemente os baselines não recorrentes com parâmetros correspondentes. No cenário mais difícil ( $L=24$ ), um modelo $4 \otimes 6$ alcançou 99,35% de precisão comparado a 7,50% do baseline $4 \otimes 1$ . Sob condições de iso-FLOPs (profundidade efetiva 24), o modelo $4 \otimes 6$ superou significativamente o baseline não recorrente $24 \otimes 1$ (99,35% vs. 37,45%), sugerindo que o compartilhamento recorrente de parâmetros fornece um viés indutivo mais forte para computação composicional do que o aumento de camadas não compartilhadas.
Indução $p$ -hop (Recuperação Recursiva): Modelos Mamba não recorrentes rasos tiveram dificuldades (ex: $1 \otimes 1$ com ~~38-41%), enquanto a adição de loops melhorou dramaticamente o desempenho. O modelo $1 \otimes 12$ alcançou precisão quase perfeita (~~99,89%), indicando que a profundidade recursiva é mais crítica do que a contagem de parâmetros distintos para a recuperação recursiva.

Modelagem de Linguagem e Benchmarks de Tarefas Subsequentes

Iso-Parâmetros: A profundidade recorrente melhorou a perplexidade de validação (PPL) para os backbones Mamba-2 e LLaMA-2 sob orçamentos de parâmetros fixos.
Iso-FLOPs: Modelos não recorrentes mais profundos ( $24 \otimes 1$ ) mantiveram uma vantagem na perplexidade de validação sobre os modelos com loop com profundidade efetiva correspondente. No entanto, os modelos com loop permaneceram competitivos em benchmarks de tarefas subsequentes, às vezes superando baselines não recorrentes mais profundos.
Saída Adaptativa:
- Escala de 140M: O portão de saída adaptativo melhorou o desempenho em tarefas subsequentes. Para o modelo Mamba-2 de 140M, o portão de saída do Estágio 2 alcançou uma pontuação média de 38,87 com um passo de saída selecionado médio de 3,00, superando o baseline de loop fixo 4 (37,79).
- Escala de 370M: O melhor desempenho em tarefas subsequentes permaneceu concentrado próximo ao setting de loop fixo 4. As variantes com portão de saída não superaram os picos de pontuação do loop fixo, sugerindo que modelos maiores podem depender mais fortemente do último passo recorrente ou requerem uma melhor calibração do portão.
- Compromisso (Trade-off): O portão de saída fornece uma fronteira controlável entre profundidade selecionada e desempenho, permitindo que o modelo utilize estados recorrentes intermediários de forma eficaz em vez de forçar uma profundidade máxima fixa.

Significância e Alegações

O artigo afirma que a profundidade recorrente oferece um eixo de escalonamento promissor e eficiente em termos de parâmetros para modelos de linguagem Mamba e híbridos. As principais descobertas incluem:

Profundidade vs. Parâmetros: Muitas tarefas de raciocínio beneficiam-se de maior profundidade computacional em vez de parâmetros independentes adicionais. O reuso recorrente de camadas Mamba melhora substancialmente as capacidades de computação composicional e recursiva sem aumentar a contagem de parâmetros.
Viabilidade Híbrida: O princípio da profundidade com loop aplica-se efetivamente a arquiteturas híbridas Mamba-Transformer, onde camadas de atenção são aplicadas em loop conjuntamente com camadas Mamba.
Seleção Adaptativa: Embora as implementações atuais de SSM exijam a execução total do estado, a seleção de saída adaptativa permite a seleção de profundidades de predição ideais, melhorando o desempenho em tarefas subsequentes em profundidades intermediárias (particularmente na escala de 140M).

Os autores mantêm modéstia em relação às suas alegações, observando que os experimentos são limitados a modelos de até 370 milhões de parâmetros e que a economia real de computação durante a inferência exigirá trabalhos futuros em mecanismos de manipulação de estado para pular passos recorrentes inteiramente. Os resultados são apresentados como evidência empírica controlada, e não como uma caracterização completa do desempenho em larga escala.