PaperlayerView exact revision on arXiv

Source-linked AI summary

Mela: Test-Time Memory Consolidation based on Transformation Hypothesis

Lungchuan Chen

arXiv:2605.10537v1cs.CL

TL;DR

Modern sequence models lack a memory-consolidation design grounded in how transient experiences become structured, abstract representations. The paper proposes HMM and integrates it with a Transformer decoder as Mela, using MemStack to expose multi-granularity memory features. Mela outperforms Transformer baselines across model sizes and remains effective beyond its 4K pre-trained context length.

  • Problem

    Memory consolidation remains largely unexplored as a design principle for modern sequence models despite its role in transforming transient experiences into structured representations.

  • Method

    HMM uses low- and high-frequency memory sub-modules to produce episodic-detail and gist-level representations, combines them context-dependently, and integrates them into a Transformer decoder with MemStack.

  • Results

    Mela outperforms strong Transformer baselines at all model sizes and generalizes effectively beyond the pre-trained context length.

  • Takeaways & Limitations

    Online hierarchical consolidation provides an effective alternative for extending context beyond the training window without additional length-extrapolation techniques.

Abstract

from arXiv · show

Memory consolidation, the process by which transient experiences are transformed into stable, structured representations, is a foundational organizing principle in the human brain, yet it remains largely unexplored as a design principle for modern sequence models. In this work, we leverage established neuroscientific theories of memory consolidation and cross-frequency coupling to propose the Hierarchical Memory Module (HMM), a neural memory architecture composed of two functionally distinct sub-modules that operate at different update frequencies. Inspired by the transformation hypothesis, the low-frequency sub-module produces high-level representations that capture abstract, gist-level knowledge, while the high-frequency sub-module produces fine-grained representations that preserve richer episodic detail. The final memory output is dynamically reconstructed as a context-dependent combination of both representations, analogous to the reconstructive nature of human memory retrieval. We integrate HMM into a Transformer-based language decoder to form Mela, a family of memory-augmented language models that perform online memory consolidation at test time. To further exploit the multi-granularity memory representations produced by HMM, we introduce MemStack, a method that distributes different levels of memory features across the early layers of the decoder without introducing additional tokens. Experiments on language modeling demonstrate that Mela outperforms Transformer baselines across all the model sizes. Moreover, with the pretrained context length fixed at 4K, Mela maintains performance on significantly longer contexts, whereas Transformer baselines degrade rapidly beyond their training length. Extensive ablation studies validate the contribution of each component and provide guidance for practical configuration.

1 Introduction

The paper treats memory consolidation as a design principle for long-context sequence models, drawing on transformation theory and cross-frequency coupling. It proposes hierarchical memory modules that transform episodic detail into gist-level representations and integrate them with a Transformer decoder.

  • Motivation: Standard attention has O(n^2) time and space complexity, limiting Transformer applicability to long-context scenarios.The paper motivates alternatives that avoid quadratic growth with sequence length.
  • Neuroscientific foundation: The transformation hypothesis describes consolidation as changing context-rich episodic memories into abstract, decontextualized semantic or schematic representations.The hippocampus supports detailed retrieval, while the neocortex retains gist-level knowledge.
  • Contributions: Mela combines HMM with a Transformer decoder to construct memory online during inference, while MemStack exposes intermediate memory features across decoder layers.The memory module is separated from the decoder and can be independently scaled and trained.
  • Contributions: HMM uses high-frequency low-level updates for richer episodic detail and lower-frequency high-level updates for gist-level representations.The high-level representation depends on and transforms the low-level representation rather than simply copying it.

2 Preliminaries

This section formalizes neural memory as information encoded in network weights and describes surprise-based updates, forgetting, associative objectives, and query-based retrieval. It also situates hierarchical latent recursion as related work while noting assumptions behind HRM and TRM’s alternative.

  • Neural memory: Neural memory encodes historical information in a network’s weight parameters through gradient-based updates, with gradients measuring input surprise.The module maps sequence inputs to sequence outputs, and its parameters store compressed historical information.
  • Neural memory: Weight-based memory can provide more expressive representations than fixed-size matrix states, depending on the underlying architecture and capacity.
  • Memory updates: A learnable decay factor controls how much past surprise contributes to each weight update, while an input-dependent forgetting factor regulates retained memory.The forgetting factor preserves all past memory as α approaches 1 and fully forgets it as α approaches 0.
  • Memory objectives: Reconstructing the original input can emphasize surface details and redundant contexts, making it suboptimal for retaining generalizable information.
  • Memory objectives: Associative loss instead encourages the module to learn structured key-value relations rather than a verbatim copy of the input.This formulation targets relational information, including how different input aspects correspond to one another.
  • Memory retrieval: Relevant stored content is retrieved by forming an input-dependent query vector from the current input and applying it to memory.
  • Hierarchical recursion: HRM uses low- and high-level recurrent modules operating at different timescales, but assumes recurrence convergence for its one-step gradient approximation.TRM avoids that fixed-point assumption through deep supervision, running untracked recursion steps before a final gradient-enabled step.

3 Hierarchical Memory Module: Neural Memory with Test-Time Consolidation

HMM uses hierarchical latent recursion to produce complementary low-level episodic and high-level semantic memory states at test time, then fuses them into a unified representation. Its design is motivated by memory consolidation and cross-frequency coupling, with different update frequencies and capacities for the two states.

  • Memory module: HMM patchifies input tokens into episodic memory units processed by shared memory blocks containing gated attention and a neural memory module.Each patch contains P consecutive tokens, and chunks control the temporal resolution of memory updates.
  • Complementary memories: The H-module produces abstract semantic representations, whereas the L-module preserves richer episodic detail through distinct memory-block capacities.The H-module has more memory blocks, while the L-module is associated with hippocampal-like episodic memory.
  • Hierarchical latent recursion: Hierarchical latent recursion alternates L-module refinement of the l state with H-module updates of the h state across H high-level cycles.The h state is updated after the l state has undergone its L refinement cycles.
  • Neuroscientific grounding: The h and l states are initialized from patched input and updated at different frequencies, providing a computational analogue of system consolidation under the transformation hypothesis.Their interaction is framed as an analogue of cross-frequency coupling between cortical and hippocampal processes.
  • Dual memory fusion: The final HMM output concatenates h and l states and passes them through a 2-layer MLP, retaining both high-level and low-level information.This fusion differs from approaches that discard the l state and supports the paper’s transformation-hypothesis interpretation.

4 Mela Architecture

Mela uses HMM representations as reference contexts for a language decoder and introduces MemStack to distribute intermediate consolidation states through early decoder layers. This exposes the decoder to the full consolidation trajectory without adding memory tokens.

  • Mela architecture: HMM produces memory representations from input text tokens, which Mela uses as reference contexts for language modeling.The memory representations have shape m ∈ R^Np×d for embedded input x ∈ R^N×d.
  • Memory-conditioned decoding: Mela segments text into chunks, prepends memory features to each chunk, and restricts token attention to memory and preceding tokens within that chunk.This design prevents direct attention across chunk boundaries, making HMM the information channel between chunks.
  • MemStack: MemStack injects intermediate memory representations from earlier HMM cycles into the first H−1 decoder layers.The final cycle’s representation enters the decoder input, while earlier-cycle representations are fused into early layers.
  • MemStack: MemStack lets early decoder layers encode both memory content and its consolidation trajectory, while later layers perform standard sequence modeling.The strategy exposes multiple consolidation stages without introducing extra memory tokens.

5 Experiments

Mela is evaluated across model scales, context lengths, and architectural ablations. It outperforms matched Transformer baselines, preserves performance beyond its 4K training window, and shows distinct roles for HMM components across context lengths.

  • Main Results: Beyond 4K, Transformer perplexity rises sharply to an order of magnitude higher at 32K, whereas Mela increases mildly and remains on the same order as its 4K result.HMM combines episodic detail from the L-module with semantic abstraction from the H-module during extended-context inference.
  • Cycle Ablations: H = 4 achieves the lowest perplexity across cycle-ablation columns, with its advantage over H = 2 expanding from 0.06 at 1K to 0.18 at 16K tokens.The results associate deeper H-module recurrence with larger gains as context exceeds the pre-training window.
  • Cycle Ablations: L = 4 yields the best perplexity across context lengths, improving over L = 1 from 21.24 to 19.11 at 1K and from 14.96 to 13.59 at 4K.The L-module’s refinement benefit is strongest at short contexts, while longer contexts can partly compensate for under-refined episodic representations.
  • Module Capacity: Additional H-module capacity helps increasingly at long contexts, while additional L-module capacity helps most at short contexts and diminishes as context grows.The layer-2 versus layer-8 gap for the L-module shrinks from 0.49 at 1K to 0.19 at 16K; together, the ablations identify H as the long-context bottleneck and L as the short-context bottleneck.
  • Additional Ablations: Removing gating raises perplexity by 1.24, removing neural memory by 0.43 versus 0.13 for attention, and removing MemStack causes a consistent 0.10 degradation.The ablations support selective writing, persistent cross-segment memory, and exposure to intermediate consolidation states as useful design choices.

6 Conclusion

Mela uses test-time neural memory formation to extend Transformer language modeling beyond pretrained context lengths. The authors position memory-augmented architectures as a direction for broader cognitive capabilities and release Mela’s implementation.

  • Mela combines the HMM neural memory architecture with the HLR mechanism and the MemStack decoder integration strategy.HMM generates test-time memory representations, while MemStack exposes final and intermediate consolidation states to the decoder.
  • Across model scales, Mela outperforms a strong Transformer baseline under the same pretrained context length and generalizes beyond that length.
  • The authors suggest Mela as an alternative for extending context length when compute or long-context training data are limited.
  • The paper frames neural memory as relevant to capabilities including long-horizon planning and open-ended behavior.The authors also release the full Mela implementation to support further exploration of memory-augmented architectures and test-time memory formation.

A Detailed Perplexity Results by Context Length

Table 5 reports perplexity across context lengths for Mela and Transformer models, enabling comparison of their language-modeling behavior as context increases.

  • Table 5 measures model performance using perplexity (PPL) across different context lengths.
  • The table compares Mela models with Transformer models across the reported context lengths.
  • The results are organized to examine model behavior as context length varies.

B Ablation Studies

The ablations show that increasing H- and L-module computation or depth improves Mela’s perplexity, with their benefits depending on context length.

  • Increasing H from 1 to 4 consistently lowers perplexity, with the largest gains beyond the 4K pretraining window.All variants use the 400M default configuration except for H cycle count.
  • L = 4 achieves the lowest perplexity at every context length, while L = 2 underperforms L = 1 at most lengths.The intermediate setting may be insufficient to produce a meaningfully refined episodic representation.
  • Increasing H-module depth from 2 to 8 memory blocks improves perplexity at every context length, especially beyond 4K tokens.All variants use the 400M default configuration except for H-module depth.
  • Increasing L-module depth from 2 to 8 memory blocks improves perplexity at every context length, with the largest gains at short contexts.All variants use the 400M default configuration except for L-module depth.
Loading 2605.10537v1…