Recall with Reasoning: Chain-of-Thought Distillation for Mamba's Long-Context Memory and Extrapolation

TL;DR

This paper introduces Recall with Reasoning (RwR), a simple distillation method that enhances Mamba's ability to recall and reason over long contexts, improving performance without architectural modifications.

Contribution

The paper proposes RwR, a novel distillation technique that significantly improves long-context memory and reasoning in Mamba models through chain-of-thought summarization during fine-tuning.

Findings

RwR improves long-context performance on LONGMEMEVAL and HELMET datasets.

RwR maintains short-context capabilities while enhancing long-context reasoning.

The method does not require changes to the model architecture.

Abstract

Mamba's theoretical infinite-context potential is limited in practice when sequences far exceed training lengths. This work explores unlocking Mamba's long-context memory ability by a simple-yet-effective method, Recall with Reasoning (RwR), by distilling chain-of-thought (CoT) summarization from a teacher model. Specifically, RwR prepends these summarization as CoT prompts during fine-tuning, teaching Mamba to actively recall and reason over long contexts. Experiments on LONGMEMEVAL and HELMET show RwR boosts Mamba's long-context performance against comparable Transformer/hybrid baselines under similar pretraining conditions, while preserving short-context capabilities, all without architectural changes.