Recent advances in Large Language Models (LLMs) have yielded impressive successes on many language tasks. However, efficient processing of long contexts using LLMs remains a significant challenge. We introduce \textbf{EpMAN} -- a method for processing long contexts in an \textit{episodic memory} module while \textit{holistically attending to} semantically relevant context chunks. The output of \textit{episodic attention} is then used to reweigh the decoder's self-attention to the stored KV cache of the context during training and generation. When an LLM decoder is trained using \textbf{EpMAN}, its performance on multiple challenging single-hop long-context recall and question-answering benchmarks is found to be stronger and more robust across the range from 16k to 256k tokens than baseline decoders trained with self-attention, and popular retrieval-augmented generation frameworks.
@article{arxiv.2502.14280,
title = {EpMAN: Episodic Memory AttentioN for Generalizing to Longer Contexts},
author = {Subhajit Chaudhury and Payel Das and Sarathkrishna Swaminathan and Georgios Kollias and Elliot Nelson and Khushbu Pahwa and Tejaswini Pedapati and Igor Melnyk and Matthew Riemer},
journal= {arXiv preprint arXiv:2502.14280},
year = {2025}
}