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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Ma, Junyu, Fang, Tianqing, Zhang, Zhisong, Zhang, Hongming, Mi, Haitao, Yu, Dong
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910981942149120
author Ma, Junyu
Fang, Tianqing
Zhang, Zhisong
Zhang, Hongming
Mi, Haitao
Yu, Dong
author_facet Ma, Junyu
Fang, Tianqing
Zhang, Zhisong
Zhang, Hongming
Mi, Haitao
Yu, Dong
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_03320
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Recall with Reasoning: Chain-of-Thought Distillation for Mamba's Long-Context Memory and Extrapolation
Ma, Junyu
Fang, Tianqing
Zhang, Zhisong
Zhang, Hongming
Mi, Haitao
Yu, Dong
Computation and Language
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.
title Recall with Reasoning: Chain-of-Thought Distillation for Mamba's Long-Context Memory and Extrapolation
topic Computation and Language
url https://arxiv.org/abs/2505.03320