Mixture of Chapters: Scaling Learnt Memory in Transformers

1. Problem Statement

Standard Transformer architectures lack an explicit, persistent mechanism for storing and organizing knowledge acquired during training. Instead, they rely on implicit knowledge encoded within dense parameters, which makes memory difficult to inspect, edit, or scale independently.

• Current Limitations: Existing solutions like KV-cache retention, retrieval-augmented generation (RAG), or external tools primarily manage processed context or retrieve external information. They do not provide a learned, internal associative memory store for factual knowledge.
• The Bottleneck: While adding a learned memory bank (latent tokens) is theoretically possible, performing dense cross-attention over a massive memory bank is computationally prohibitive ($O(L \cdot N_m)$), preventing the scaling of memory capacity to useful sizes.

2. Methodology: Mixture of Chapters (MoC)

The authors propose Mixture of Chapters (MoC), a novel architecture that integrates a learnable, sparse memory bank into standard Transformer layers.

Core Components

1. Learned Latent Memory Bank:

   • A set of $N_m$ latent tokens (e.g., 262,208 tokens) initialized randomly and trained end-to-end.
   • These tokens act as a persistent storage for factual knowledge.
   • Access is achieved via cross-attention: The token stream generates queries ($Q$), while the memory bank provides keys ($K$) and values ($V$).

2. Chapter-Based Routing (Sparse Access):

   • To avoid the cost of attending to all memory tokens, the bank is partitioned into $C$ "chapters" (e.g., 4,097 chapters).
   • Router Mechanism: A lightweight router takes a pooled representation of the input sequence ($r = \text{Pool}(H)$) and computes a probability distribution over chapters.
   • Top-K Selection: The router selects only the top-$k$ most relevant chapters (e.g., $k=64$) for a given input.
   • Computation: The model attends only to the selected chapters ($M_S$), reducing the attention complexity from $O(L \cdot N_m)$ to $O(L \cdot k \cdot T)$, where $T$ is tokens per chapter.

3. Architecture Integration:

   • The MoC layer is inserted into standard decoder-only Transformer blocks.
   • It consists of a memory cross-attention module followed by a residual connection, similar to standard self-attention but querying the external memory bank.

3. Key Contributions

1. Learned Associative Memory Bank: Introduced a persistent, addressable bank of latent tokens trained end-to-end, enabling internal retrieval in a compact latent space rather than relying on text-based retrieval.
2. Mixture of Chapters (MoC): Developed a scalable sparse access mechanism inspired by Mixture-of-Experts (MoE). By routing inputs to specific memory chapters, the model scales memory capacity to 262K tokens while maintaining tractable computation.
3. Iso-FLOP Evaluation: Evaluated the architecture against standard dense Transformers under fixed compute budgets (iso-FLOP), demonstrating that explicit memory provides complementary capacity to parameter-based learning.
4. Retention Analysis: Provided empirical evidence that explicit memory improves knowledge retention during continued training (e.g., transitioning from pre-training to instruction fine-tuning), mitigating catastrophic forgetting.

4. Experimental Results

The authors evaluated their model on pre-training and instruction fine-tuning (IFT) tasks using a 9.6B token pre-training budget and heavy IFT (2 epochs on 230M tokens).

• Pre-training Performance:

  • The MoC model achieved a lower validation loss (2.79) compared to the iso-FLOP dense baseline (2.86) and the backbone-only model (2.92).
  • This indicates that the memory bank adds effective capacity beyond what can be achieved by simply increasing model depth or width within the same compute budget.

• Knowledge Retention & Forgetting:

  • Scenario: After heavy instruction fine-tuning, dense models suffered significant "catastrophic forgetting" on knowledge-heavy benchmarks.
  • Results:
    • ARC-Challenge: The dense baseline dropped by 6.69% (approaching random guessing), whereas MoC dropped only 2.68%.
    • BoolQ: The dense baseline dropped by 6.24%, while MoC remained stable with a negligible change (+0.24%).
  • Conclusion: The memory bank acts as an anchor for factual knowledge, preventing it from being overwritten during the adaptation to new tasks (IFT).

• Freezing the Memory Bank:

  • Experiments showed that freezing the memory bank during IFT (updating only the backbone) yielded performance on par with updating the bank. This suggests the pre-trained memory is robust and reusable without further optimization, reducing VRAM requirements during post-training.

5. Significance and Implications

• New Scaling Axis: The paper establishes "learned memory" as a distinct and complementary axis for scaling LLMs, alongside model size and data volume. It proves that explicit associative memory can store facts more efficiently than dense parameters alone.
• Mitigating Catastrophic Forgetting: The results offer a promising architectural solution to the problem of catastrophic forgetting in continual learning and instruction tuning, suggesting that separating "factual anchoring" (memory) from "task adaptation" (backbone) is a viable strategy.
• Efficiency: By utilizing sparse routing, the approach makes large-scale memory storage computationally feasible, bridging the gap between the theoretical benefits of external memory and the practical constraints of Transformer inference.
• Future Directions: The work opens avenues for specialized memory editing, interpretability (analyzing which chapters store specific facts), and hybrid systems combining shared global memory with routed private memory.

In summary, Mixture of Chapters successfully demonstrates that augmenting Transformers with a sparse, routed, learned memory bank improves training efficiency, reduces forgetting, and scales knowledge capacity without prohibitive computational costs.