🤖 machine learning

Mixture of Chapters: Scaling Learnt Memory in Transformers

Este artigo apresenta os "bancos de memória aprendíveis" (learnable sparse memory banks), uma arquitetura que integra tokens latentes treináveis e roteamento baseado em capítulos aos Transformers, permitindo escalar a capacidade de memória para 262 mil tokens e melhorar a retenção de conhecimento sem custos computacionais proibitivos.

Autores originais: Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

Publicado 2026-03-24

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

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ê está tentando ensinar um robô superinteligente (um modelo de IA) a ler milhões de livros e aprender tudo sobre o mundo. O problema é que, atualmente, esses robôs funcionam como estudantes que tentam decorar tudo de cabeça. Se você pedir para eles aprenderem algo novo (como uma nova língua ou regras de um jogo), eles muitas vezes esquecem o que sabiam antes. É como se o cérebro deles fosse um balde furado: quanto mais você enche, mais vaza.

Os autores deste artigo propuseram uma solução brilhante: dar ao robô uma "biblioteca interna".

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

1. O Problema: A Memória "Escondida"

Os modelos de IA atuais (chamados de Transformers) guardam o conhecimento de forma "escondida" nos seus pesos matemáticos. É como se o robô tentasse lembrar de um número de telefone tentando recriar a imagem do papel onde ele estava escrito, em vez de apenas olhar para o papel.

O defeito: Quando você tenta ensinar algo novo, o robô precisa reescrever toda a sua "memória de cabeça", o que apaga as informações antigas. É como tentar escrever um novo capítulo em um livro onde as páginas já estão cheias; você precisa rasurar o velho para escrever o novo.

2. A Solução: A "Biblioteca de Latentes"

Os pesquisadores criaram algo chamado Banco de Memória Aprendível.

A Analogia: Imagine que, em vez de tentar decorar tudo, o robô tem uma biblioteca física dentro de si mesmo. Essa biblioteca é composta por milhares de "cartões de memória" (tokens latentes).
Como funciona: Quando o robô precisa responder a uma pergunta, ele não tenta "lembrar" de tudo de uma vez. Ele olha para a pergunta, vai até a biblioteca e busca apenas os cartões relevantes. É como usar um índice de um livro: você não lê o livro inteiro, vai direto ao capítulo que precisa.

3. O Desafio: A Biblioteca é Gigante

Aqui entra o problema da escala. Se a biblioteca tiver 262.000 cartões, o robô teria que olhar para todos eles para cada pergunta. Isso seria lento demais e gastaria muita energia (computação). Seria como pedir para um bibliotecário procurar um livro específico em uma biblioteca do tamanho de um país, lendo o título de cada um dos milhões de livros um por um.

4. A Inovação: O "Mistério dos Capítulos" (Mixture of Chapters)

Para resolver isso, os autores criaram o Mistério dos Capítulos (Mixture of Chapters - MoC).

A Analogia: Em vez de ter uma pilha gigante de cartões, eles organizaram a biblioteca em 4.000 caixas diferentes (capítulos).
O Bibliotecário Inteligente (O Roteador): Antes de abrir qualquer caixa, o robô usa um "bibliotecário inteligente" (um roteador leve). Esse bibliotecário olha para a pergunta do usuário e diz: "Ah, essa pergunta é sobre culinária. Não precisamos abrir as caixas de história ou física. Vamos abrir apenas a caixa de culinária e talvez a de química."
O Resultado: O robô só consulta uma pequena fração da biblioteca (apenas algumas caixas) para cada pergunta. Isso torna o processo super rápido e eficiente, permitindo que a biblioteca cresça para um tamanho enorme sem deixar o robô lento.

5. Por que isso é revolucionário?

O artigo mostra dois benefícios principais:

Aprendizado Contínuo sem Esquecer: Quando o robô aprende algo novo (como ser treinado para seguir instruções), ele não precisa reescrever sua "cabeça". Ele apenas adiciona novos cartões à biblioteca ou atualiza os existentes. O conhecimento antigo fica seguro na biblioteca, protegido da "esquecimento". É como ter um caderno de anotações separado: você pode escrever novas receitas sem apagar as antigas.
Mais Inteligência com o Mesmo Esforço: Eles provaram que, usando essa biblioteca, o robô ficou mais inteligente do que outros robôs do mesmo tamanho que não tinham a biblioteca, mesmo usando a mesma quantidade de energia de computação.

Resumo em uma frase

Os autores criaram um sistema onde a IA não tenta "decorar" tudo na cabeça, mas sim organiza o conhecimento em uma biblioteca gigante e inteligente, onde um "bibliotecário" seleciona apenas os livros necessários para cada tarefa, permitindo que a máquina aprenda coisas novas sem esquecer as antigas.

É como trocar a memória de um computador (que é lenta e limitada) por uma biblioteca viva e organizada, onde o acesso é rápido e o conhecimento nunca se perde.

1. O Problema

Os Transformers atuais possuem uma capacidade limitada de armazenamento de conhecimento explícito. O conhecimento é codificado implicitamente nos parâmetros densos do modelo, o que torna difícil inspecionar, editar ou escalar a memória de forma independente. Além disso, durante o treinamento contínuo (como o ajuste fino para instruções), os modelos sofrem de esquecimento catastrófico, perdendo fatos aprendidos anteriormente ao se adaptarem a novas distribuições de dados.

Métodos existentes, como RAG (Geração Aumentada por Recuperação) ou caches de KV, geralmente dependem de recuperação externa ou gerenciamento de contexto, não oferecendo um armazenamento interno aprendido e persistente que seja eficiente em termos computacionais em grande escala. O desafio principal é criar uma memória interna que possa escalar para centenas de milhares de tokens sem tornar o custo de atenção (computação) proibitivo.

2. Metodologia: Mixture of Chapters (MoC)

Os autores propõem uma arquitetura de Transformer aumentada com um banco de memória aprendido e um mecanismo de roteamento esparsificado chamado Mixture of Chapters (MoC).

A. Banco de Memória Aprendido

Estrutura: O modelo possui um banco de memória $M$ composto por $N_m$ tokens latentes (vetores), inicializados aleatoriamente e treinados end-to-end.
Acesso: Em vez de usar atenção densa sobre todos os tokens de memória (o que seria $O(L \cdot N_m)$ ), o modelo usa atenção cruzada. O fluxo de tokens do Transformer gera as queries ( $Q$ ), enquanto o banco de memória fornece as keys ( $K$ ) e values ( $V$ ).
Integração: Esta camada de memória é inserida em blocos específicos do Transformer (entre as camadas de auto-atenção e MLP).

B. Escalabilidade via Mixture of Chapters (MoC)

Para evitar o custo computacional de atender a um banco de memória massivo (ex: 262.208 tokens) para cada entrada, os autores dividem o banco em capítulos:

Particionamento: O banco de memória é dividido em $C$ capítulos (ex: 4.097 capítulos), onde cada capítulo contém $T$ tokens (ex: 64 tokens).
Roteador Leve: Um roteador simples analisa o estado oculto da sequência (via pooling) e calcula uma distribuição de probabilidade sobre os capítulos.
Seleção Esparsa: Para cada sequência de entrada, o roteador seleciona apenas o top- $k$ capítulos (ex: $k=64$ ) mais relevantes.
Atenção Restrita: A atenção cruzada é realizada apenas sobre os tokens dos capítulos selecionados.

Isso reduz a complexidade de $O(L \cdot N_m)$ para $O(L \cdot k \cdot T)$ , permitindo escalar a capacidade de memória para centenas de milhares de tokens mantendo a computação tratável.

3. Contribuições Principais

Memória Associativa Aprendida: Introdução de um banco de memória persistente e endereçável dentro do Transformer, acessado via atenção cruzada, complementando a capacidade paramétrica implícita.
Mixture of Chapters (MoC): Um mecanismo de roteamento inspirado em Mixture of Experts (MoE) que permite o acesso esparsificado e escalável a grandes bancos de memória.
Escalabilidade sob Orçamento Fixo: Demonstração de que a arquitetura de memória supera baselines de Transformers densos em configurações iso-FLOP (mesmo custo computacional).
Retenção de Conhecimento: Evidência empírica de que a memória explícita mitiga o esquecimento catastrófico durante transições de treinamento (ex: do pré-treinamento para o ajuste fino de instruções).

4. Resultados Experimentais

Os experimentos foram conduzidos com um modelo pré-treinado por 9,6 bilhões de tokens e subsequentemente ajustado para instruções (IFT) em 230 milhões de tokens.

Desempenho no Pré-treinamento: O modelo MoC alcançou uma perda de validação menor (2.79) comparado ao baseline iso-FLOP (2.86) e ao backbone apenas (2.92), indicando melhor eficiência de aprendizado.
Resistência ao Esquecimento (IFT):
- O baseline Vanilla sofreu degradação severa em tarefas baseadas em conhecimento após o IFT (ex: ARC-Challenge caiu de 31.77% para 25.08%, aproximando-se do acaso).
- O modelo MoC manteve a estabilidade (ARC-Challenge caiu apenas para 28.76% e BoolQ manteve-se estável com +0.24%).
- Isso sugere que o banco de memória "ancora" o conhecimento factual, protegendo-o contra a interferência do ajuste fino.
Congelamento da Memória: Um achado crucial é que congelar o banco de memória durante o IFT (sem atualizá-lo) produziu resultados comparáveis a atualizá-lo. Isso indica que a memória pré-treinada é robusta e pode ser reutilizada sem custo adicional de otimização.
Custo Computacional: O modelo MoC atingiu um equilíbrio eficiente, com um custo de forward pass de ~~0.459T FLOPs, comparável ao baseline iso-FLOP (~~0.496T), mas com maior capacidade de memória.

5. Significado e Conclusão

O trabalho propõe um novo eixo de escalabilidade para modelos de linguagem: memória aprendida explícita.

Complementaridade: A memória atua como um complemento à capacidade paramétrica, permitindo que o modelo armazene fatos de forma mais eficiente e editável do que apenas através de pesos densos.
Eficiência: A abordagem Mixture of Chapters resolve o gargalo de escalabilidade da atenção densa, tornando viável o uso de bancos de memória com centenas de milhares de tokens.
Aprendizado Contínuo: A arquitetura oferece uma solução promissora para o problema do esquecimento catastrófico, sugerindo que separar o "aprendizado de instruções" (backbone) do "armazenamento de fatos" (memória) pode levar a modelos mais robustos e estáveis em cenários de treinamento contínuo.

Em suma, o MoC demonstra que a adição de uma memória interna esparsamente acessível não apenas melhora o desempenho em tarefas de conhecimento, mas também estabiliza o modelo contra a perda de informações durante fases posteriores de treinamento.