Geheugenefficiënte Lus-Transformer: Ontkoppeling van Rekenkracht en Geheugen in Geluste Taalmodellen

Samenvatting

Recurrente LLM-architecturen zijn naar voren gekomen als een veelbelovende benadering voor het verbeteren van redeneren, omdat ze meerstapsberekening in de inbeddingsruimte mogelijk maken zonder het genereren van tussentijdse tokens. Modellen zoals Ouro voeren redeneren uit door iteratief interne representaties bij te werken, terwijl een standaard Key-Value (KV) cache over iteraties heen behouden blijft, waardoor het geheugenverbruik lineair toeneemt met de redeneerdiepte. Dit kan ertoe leiden dat een toename van het aantal redeneeriteraties leidt tot prohibitief geheugengebruik, wat de praktische schaalbaarheid van dergelijke architecturen beperkt. In dit werk stellen we de Memory-Efficient Looped Transformer (MELT) voor, een nieuwe architectuur die de redeneerdiepte ontkoppelt van het geheugenverbruik. In plaats van een standaard KV cache per laag en per lus te gebruiken, onderhoudt MELT één enkele KV cache per laag die wordt gedeeld over redeneerlussen. Deze cache wordt in de loop van de tijd bijgewerkt via een leerbaar poortmechanisme. Om stabiele en efficiënte training onder deze architectuur mogelijk te maken, stellen we voor om MELT te trainen met behulp van chunksgewijze training in een tweefasige procedure: geïnterpoleerde overgang, gevolgd door aandacht-uitgelijnde distillatie, beide van het startmodel LoopLM naar MELT. Empirisch tonen we aan dat MELT-modellen die zijn fijnafgesteld op voorgetrainde Ouro-parameters beter presteren dan standaard LLM's van vergelijkbare grootte, terwijl ze een geheugenvoetafdruk behouden die vergelijkbaar is met die modellen en aanzienlijk kleiner dan die van Ouro. Over het algemeen bereikt MELT constant-geheugen iteratief redeneren zonder de prestaties van LoopLM op te offeren, met slechts een lichtgewicht nabehandelingsprocedure.

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.