Transformer em Loop Eficiente em Memória: Desacoplando a Computação da Memória em Modelos de Linguagem em Loop

Resumo

Arquiteturas recorrentes de LLM surgiram como uma abordagem promissora para melhorar o raciocínio, pois permitem a computação em múltiplas etapas no espaço de embeddings sem gerar tokens intermediários. Modelos como o Ouro realizam raciocínio atualizando iterativamente representações internas, enquanto mantêm um cache padrão de Chave-Valor (KV) através das iterações, fazendo com que o consumo de memória cresça linearmente com a profundidade do raciocínio. Consequentemente, aumentar o número de iterações de raciocínio pode levar a um uso proibitivo de memória, limitando a escalabilidade prática de tais arquiteturas. Neste trabalho, propomos o Memory-Efficient Looped Transformer (MELT), uma nova arquitetura que desacopla a profundidade do raciocínio do consumo de memória. Em vez de usar um cache KV padrão por camada e por laço, o MELT mantém um único cache KV por camada que é compartilhado entre os laços de raciocínio. Esse cache é atualizado ao longo do tempo por meio de um mecanismo de portão aprendível. Para permitir um treinamento estável e eficiente sob esta arquitetura, propomos treinar o MELT usando treinamento por blocos em um procedimento de duas fases: transição interpolada, seguida de destilação alinhada por atenção, ambas do modelo inicial LoopLM para o MELT. Empiricamente, mostramos que modelos MELT ajustados a partir de parâmetros pré-treinados do Ouro superam LLMs padrão de tamanho comparável, enquanto mantêm uma pegada de memória comparável a esses modelos e dramaticamente menor que a do Ouro. No geral, o MELT alcança raciocínio iterativo com memória constante sem sacrificar o desempenho do LoopLM, utilizando apenas um procedimento leve de pós-treinamento.

English

Recurrent LLM architectures have emerged as a promising approach for improving reasoning, as they enable multi-step computation in the embedding space without generating intermediate tokens. Models such as Ouro perform reasoning by iteratively updating internal representations while retaining a standard Key-Value (KV) cache across iterations, causing memory consumption to grow linearly with reasoning depth. Consequently, increasing the number of reasoning iterations can lead to prohibitive memory usage, limiting the practical scalability of such architectures. In this work, we propose Memory-Efficient Looped Transformer (MELT), a novel architecture that decouples reasoning depth from memory consumption. Instead of using a standard KV cache per layer and loop, MELT maintains a single KV cache per layer that is shared across reasoning loops. This cache is updated over time via a learnable gating mechanism. To enable stable and efficient training under this architecture, we propose to train MELT using chunk-wise training in a two phase procedure: interpolated transition, followed by attention-aligned distillation, both from the LoopLM starting model to MELT. Empirically, we show that MELT models fine-tuned from pretrained Ouro parameters outperform standard LLMs of comparable size, while maintaining a memory footprint comparable to those models and dramatically smaller than Ouro's. Overall, MELT achieves constant-memory iterative reasoning without sacrificing LoopLM performance, using only a lightweight post-training procedure.