Long Context Modeling with Ranked Memory-Augmented Retrieval
Extracted source text · arXiv 2503.14800
%\pdfoutput=1
% This must be in the first 5 lines to tell arXiv to use pdfLaTeX, which is strongly recommended.
\pdfoutput=1
% In particular, the hyperref package requires pdfLaTeX in order to break URLs across lines.
\documentclass[11pt, dvipsnames]{article}
% Change "review" to "final" to generate the final (sometimes called camera-ready) version.
% Change to "preprint" to generate a non-anonymous version with page numbers.
\usepackage[preprint]{acl}
% Standard package includes
\usepackage{times}
\usepackage{latexsym}
% For proper rendering and hyphenation of words containing Latin characters (including in bib files)
\usepackage[T1]{fontenc}
% For Vietnamese characters
% \usepackage[T5]{fontenc}
% See https://www.latex-project.org/help/documentation/encguide.pdf for other character sets
\usepackage{amsmath}
\usepackage{adjustbox}
\usepackage{amssymb}
% \usepackage[dvipsnames]{xcolor}
\AtBeginDocument{%
\providecommand\BibTeX{{%
Bib\TeX}}}
\usepackage{graphicx}
\usepackage{subfig}
\newcommand\scalemath[2]{\scalebox{#1}{\mbox{\ensuremath{\displaystyle #2}}}}
% This assumes your files are encoded as UTF8
\usepackage[utf8]{inputenc}
% This is not strictly necessary, and may be commented out,
% but it will improve the layout of the manuscript,
% and will typically save some space.
\usepackage{microtype}
% This is also not strictly necessary, and may be commented out.
% However, it will improve the aesthetics of text in
% the typewriter font.
\usepackage{inconsolata}
%Including images in your LaTeX document requires adding
%additional package(s)
\usepackage{graphicx}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{amsmath}
\usepackage{dblfloatfix}
\usepackage{multirow}
\usepackage{booktabs}
\usepackage{verbatim}
\usepackage{balance}
\usepackage{array}
\usepackage{float}
\usepackage{adjustbox}
\usepackage{afterpage}
\algdef{SE}[VARIABLES]{Variables}{EndVariables}
{\algorithmicvariables}
{\algorithmicend\ \algorithmicvariables}
\algnewcommand{\algorithmicvariables}{\textbf{Declare}}
\usepackage{times}
\usepackage{latexsym}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{microtype}
\usepackage{inconsolata}
\usepackage{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\newtheorem{definition}{Definition}
% \title{Long Text Modelling with Pointwise Semantic Augmented Retrieval}
\title{Long Context Modeling with Ranked Memory-Augmented Retrieval}
% \title{Long Text Modelling with Pointwise Semantic Augmented Retrieval}
\title{Long Context Modeling with Ranked Memory-Augmented Retrieval}
\author{
Ghadir Alselwi\textsuperscript{1},
Hao Xue\textsuperscript{1},
Shoaib Jameel\textsuperscript{2},
Basem Suleiman\textsuperscript{1},
Hakim Hacid\textsuperscript{3}, \\
\textbf{Flora D. Salim}\textsuperscript{1},
\textbf{Imran Razzak}\textsuperscript{4,1} \\
\textsuperscript{1}University of New South Wales, Sydney, NSW, Austalia \\
\textsuperscript{2}{University of Southampton, Southampton, UK} \\
\textsuperscript{3}{Technology Innovation Institute, Abu Dhabi, UAE} \\
\textsuperscript{4}Mohamed Bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE \\
\texttt{\{g.alselwi,hao.xue1,b.suleiman, flora.salim\}@unsw.edu.au }\\
\texttt{m.s.jameel@southampton.ac.uk, hakim.hacid@tii.ae, imran.razzak@mbzuai.ac.ae} \\
}
\begin{document}
\maketitle
% \begin{abstract}
% Effective long-term memory management is crucial for language models handling extended contexts. We introduce a novel framework that dynamically ranks memory entries based on relevance. Unlike previous works, our model introduces a novel relevance scoring and a pointwise re-ranking model for key-value embeddings, inspired by learning-to-rank techniques in information retrieval. Enhanced Ranked Memory Augmented Retrieval (\textbf{ERMAR}) achieves state-of-the-art results on standard benchmarks.
% \end{abstract}
\begin{abstract}
Effective long-term memory management is crucial for language models handling extended contexts. We introduce the Enhanced Ranked Memory Augmented Retrieval (\textbf{ERMAR}) framework, which dynamically ranks memory entries based on relevance. Unlike prior models, ERMAR employs a novel relevance scoring mechanism and a pointwise re-ranking model for key-value embeddings, inspired by learning-to-rank techniques in information retrieval. By integrating historical usage patterns and adaptive retrieval, ERMAR achieves state-of-the-art results on standard benchmarks, demonstrating superior scalability and performance in long-context tasks.
\end{abstract}
\section{Introduction}
Large Language Models (LLMs) face a fundamental limitation in processing long-context scenarios due to the quadratic complexity of attention mechanisms and increasing memory demands during generation~\cite{vaswani2017attention, tworkowski2024focused}. Consider a scenario in an automated customer service system: \textit{A customer reports an issue with their printer, referencing a setup process from a previous conversation that occurred two hours ago. After 50 messages of troubleshooting, the customer mentions that the same error from the beginning has resurfaced. Traditional LLMs, constrained by their context window, would struggle to access the crucial earlier context about the initial setup process, leading to inconsistent or incomplete responses, Figure~\ref{fig:toyexample}.} It is well known that handling extended contexts remains a significant challenge, particularly in applications requiring document analysis and sustained dialogue interactions.
The recent MemLong~\cite{liu2024memlong} architecture stores and accesses historical context through basic chunk-level memory operations. The memory bank model is a large, non-trainable store of past context representations. Instead of re-computing representations for all past tokens every time, these representations are pre-computed and stored. Given the current context, MemLong retrieves relevant segments from the memory bank. It uses a dot product similarity search to find the memory entries most related to the current context. This allows the model to focus only on the most pertinent past information, rather than processing the entire history. However, its treatment of all key-value (K-V) pairs with equal weight, regardless of their contextual relevance, often leads to information overload and reduced retrieval precision. This limitation becomes particularly evident in scenarios requiring context management.
We have developed a novel model that addresses the aforementioned limitations by building upon Memlong~\cite{liu2024memlong}, a publicly available baseline on GitHub\footnote{https://github.com/Bui1dMySea/MemLong}. Our \textbf{Enhanced Ranked Memory Augmented Retrieval} (\textbf{ERMAR}) model has a novel relevance scoring mechanism that fundamentally improves context retrieval and utilization for K-V embeddings. Unlike MemLong, ERMAR employs multiplication \cite{cao2007learning} to compute relevance scores, enabling a more nuanced and context-aware assessment of semantic alignment between queries and stored memory. ERMAR also incorporates a re-ranking mechanism that dynamically reorders K-V embeddings based on their relevance scores, ensuring that the most pertinent information is prioritized during retrieval. This re-ranking process, combined with an adaptive retrieval system that integrates historical usage patterns, allows ERMAR to capture subtle contextual relationships better and refine memory prioritization. As shown in Figure~\ref{fig:toyexample}, ERMAR processes incoming queries and long-context conversations through a novel ranking architecture, employing K-V pairs ranking (\scalebox{0.9}{$K_0$-$V_0$, $K_1$-$V_1$,...,$K_i$-$V_i$}) and corresponding embeddings to perform semantic search and ranking of relevant historical information.
\begin{figure*}
\centering
% \includegraphics[scale=0.25]{RAG_toyexample.png}
\includegraphics[scale=0.35]{RAG_toyexample.png}
\caption{Our novel proposed ERMAR system. Note the difference from the MemLong architecture where we have introduced a novel \textit{Reranking} model.}
\label{fig:toyexample}
\end{figure*}
Our novel ERMAR model introduces three key improvements: (i) A semantic similarity metric to measure contextual alignment between query embeddings and key-value pairs; (ii) A weighted scoring function that considers content similarity and contextual relevance; and (iii) Integration of historical usage patterns to refine relevance assessment.
% \section{Related Work}
% Existing memory-augmented architectures such as RetroMAE~\cite{xiao2022retromae} have demonstrated promising results but often struggle with semantic coherence and retrieval efficiency. Similarly, sparse attention mechanisms reduce computational complexity but frequently sacrifice model capability~\cite{beltagy2020longformer}. Memorizing Transformers~\cite{wu2022memorizing} developed dedicated memory tokens. Recent advances have focused on memory management through hierarchical cache structures~\cite{wu2022memorizing} and context-aware retention strategies~\cite{borgeaud2022improving}. However, these systems often struggle with distribution shifts during training and effective information retrieval. While these methods offer partial solutions, they often face challenges in maintaining consistent performance across varying context lengths and ensuring reliable information retrieval.
% Architectural modifications have focused primarily on improving position representation and attention mechanisms. Position encoding adaptations such as RoPE~\cite{su2024roformer} and ALiBi~\cite{press2022trainshorttestlong} have enhanced models' ability to handle longer sequences by providing more robust position information. YARN~\cite{peng2023yarn} further advanced this approach through dynamic position embeddings, demonstrating reliable generalization up to 128k tokens. Complementing these developments, sparse attention mechanisms such as Longformer~\cite{beltagy2020longformer} and BigBird~\cite{zaheer2020big} have reduced computational complexity while maintaining model capabilities through selective attention patterns.
\section{Related Work}
The challenge of enabling language models to effectively process extended contexts has driven research across multiple domains. We organize related work into key areas that inform our ERMAR framework.
\subsection{Memory-Augmented Neural Networks}
Early neural memory architectures like Neural Turing Machines~\cite{graves2014neural} and Differentiable Neural Computers~\cite{graves2016hybrid} established the principle of external memory mechanisms beyond parameter storage. Recent memory-augmented architectures have focused on language modeling tasks. RetroMAE~\cite{xiao2022retromae} demonstrated promising results but struggled with semantic coherence when retrieved passages lacked contextual relevance. Memorizing Transformers~\cite{wu2022memorizing} introduced dedicated memory tokens for cross-sequence information storage, showing improvements on long-range dependency tasks but facing scalability limitations.
The recent MemLong architecture~\cite{liu2024memlong} represents a significant advance by storing historical context through chunk-level memory operations using a non-trainable memory bank. However, its uniform treatment of all key-value pairs regardless of contextual relevance often leads to information overload and reduced retrieval precision, particularly in scenarios requiring nuanced context management.
\subsection{Retrieval-Augmented Generation and Long-Context Modeling}
Retrieval-augmented generation (RAG)\cite{lewis2020retrieval} pioneered integrating dense passage retrieval with generative models, while Fusion-in-Decoder (FiD)\cite{izacard2020leveraging} improved efficiency through independent passage processing. REALM~\cite{guu2020retrieval} introduced end-to-end learning of retrieval and generation, and DPR~\cite{karpukhin2020dense} established dense passage retrieval standards. However, these methods typically focus on static corpora rather than dynamic context-aware memory management.
The quadratic complexity of attention mechanisms has driven research into efficient long-context architectures. Sparse attention mechanisms such as Longformer~\cite{beltagy2020longformer} and BigBird~\cite{zaheer2020big} reduce computational complexity while maintaining model capabilities through selective attention patterns. Position encoding adaptations like RoPE~\cite{su2024roformer} and ALiBi~\cite{press2022trainshorttestlong} have enhanced models' ability to handle longer sequences. YARN~\cite{peng2023yarn} further advanced this through dynamic position embeddings, demonstrating reliable generalization up to 128k tokens.
\subsection{Learning-to-Rank and Information Retrieval}
Our ERMAR framework draws inspiration from learning-to-rank techniques in information retrieval. Traditional ranking approaches like BM25 struggle with semantic similarity, leading to neural ranking models using learned representations. Pointwise ranking approaches~\cite{cao2007learning} predict relevance scores for query-document pairs, directly inspiring our relevance scoring mechanism. Dense retrieval methods like DPR~\cite{karpukhin2020dense} and ColBERT~\cite{khattab2020colbert} demonstrate the effectiveness of learned dense representations for ranking through similarity scores—a paradigm we adapt for memory entry ranking.
BERT-based re-ranking models~\cite{nogueira2019passage} show that sophisticated re-ranking significantly improves retrieval quality. This two-stage retrieve-then-rerank paradigm directly influences our design, where we first retrieve candidate memory entries then apply learned re-ranking.
\subsection{Current Limitations and Gaps}
Despite significant progress, current approaches face limitations that motivate ERMAR:
\textbf{Static Memory Management}: Most approaches use fixed memory structures that don't adapt to content importance or usage patterns, leading to inefficient memory utilization.
\textbf{Uniform Memory Treatment}: Existing methods treat all memory entries equally, lacking mechanisms to prioritize more relevant information.
\textbf{Limited Semantic Understanding}: Memory systems often rely on simple similarity metrics without considering contextual nuance or temporal relevance.
\textbf{Scalability Trade-offs}: Current approaches face difficult trade-offs between memory capacity, retrieval accuracy, and computational efficiency.
Our ERMAR framework addresses these limitations through dynamic relevance scoring, adaptive memory management, and sophisticated re-ranking mechanisms inspired by information retrieval techniques, providing a more principled approach to long-context memory management.
\section{Our Novel ERMAR Model}
Figure~\ref{fig:arch} illustrates our ERMAR framework and Figure~\ref{fig:toyexample} presents the contextual ranking mechanism of key, value pairs components that enable effective retrieval. ERMAR maintains consistency through frozen lower layers and selective parameter updating. ERMAR: (i) stores important information from earlier parts of the text; (ii) assigns relevance scores to stored information based on its importance to the current context, and (iii) retrieves only the most relevant historical information when needed. Our relevance scoring is analogous to attention, allowing the model to focus on important parts of the memory. There is also a ``loose'' pointwise connection because the primary objective is sequence likelihood.
\begin{figure*}
\centering
\includegraphics[scale=0.35]{Methodology.png}
% \includegraphics[scale=0.25]{Methodology.png}
\caption{The architecture diagram for ERMAR.}
\label{fig:arch}
\end{figure*}
% Let $\mathcal{V}$ be a finite vocabulary, and \(\scalebox{0.9}{$\mathbf{x} = (x_1, \ldots, x_n) \in \mathcal{V}^n$}\) a token sequence with preceding context $\mathbf{x}_{<i}$. The embedding function \(\scalebox{0.9}{$\mathcal{E}: \mathcal{V}^* \rightarrow \mathbb{R}^{d_\text{ret}}$}\) maps sequences to retrieval space. We introduce a memory function $\mathcal{M}$ augmented with a relevance scoring mechanism
% \(\scalebox{0.9}{$\mathcal{M}: \underbrace{\mathbb{R}^{d_\text{model}}}_{\text{keys}} \times \underbrace{\mathbb{R}^{d_\text{model}}}_{\text{values}} \times \underbrace{\mathbb{R}^{d_\text{ret}}}_{\text{embeddings}} \rightarrow \mathcal{S}$}\).
Let $\mathcal{V}$ be a finite vocabulary, and
\begin{equation*}
\mathbf{x} = (x_1, \ldots, x_n) \in \mathcal{V}^n
\end{equation*}
a token sequence with preceding context $\mathbf{x}_{<i}$. The embedding function
\begin{equation*}
\mathcal{E}: \mathcal{V}^* \rightarrow \mathbb{R}^{d_\text{ret}}
\end{equation*}
maps sequences to retrieval space. We introduce a memory function $\mathcal{M}$ augmented with a relevance scoring mechanism:
\begin{equation*}
\mathcal{M}: \underbrace{\mathbb{R}^{d_\text{model}}}_{\text{keys}} \times \underbrace{\mathbb{R}^{d_\text{model}}}_{\text{values}} \times \underbrace{\mathbb{R}^{d_\text{ret}}}_{\text{embeddings}} \rightarrow \mathcal{S}
\end{equation*}
% \textbf{Relevance Score:} Given a query embedding $\mathbf{q} \in \mathbb{R}^{d_\text{ret}}$ and a matrix of key embeddings \scalebox{0.9}{$\mathbf{K} = [k_1, \ldots, k_m] \in \mathbb{R}^{m \times d_\text{ret}}$} (where each row $k_i \in \mathbb{R}^{d_\text{ret}}$ corresponds to a key embedding), the relevance score is:
% \begin{equation}
% \scalebox{0.9}{$
% \alpha(\mathbf{q}, \mathbf{K}) = \text{softmax}\left(\frac{\mathbf{q} \mathbf{K}^\top}{\sqrt{d_\text{ret}}}\right)
% \label{eq:relevance_score}$}
% \end{equation}
% \noindent Here, $\sqrt{d_\text{ret}}$ normalizes the similarity scores to prevent excessively large values. The relevance score $\alpha(\mathbf{q}, \mathbf{K})$ can be interpreted as a probability distribution over the keys, where each entry $\alpha_i$ represents the relative importance (or attention weight) of the $i$-th key to the query $\mathbf{q}$. This score is used to rank memory entries based on their importance to the current query.
\textbf{Relevance Score:} Given a query embedding $\mathbf{q} \in \mathbb{R}^{d_\text{ret}}$ and a matrix of key embeddings
\begin{equation*}
\mathbf{K} = [k_1, \ldots, k_m] \in \mathbb{R}^{m \times d_\text{ret}}
\end{equation*}
(where each row $k_i \in \mathbb{R}^{d_\text{ret}}$ corresponds to a key embedding), the relevance score is:
\begin{equation}
\alpha(\mathbf{q}, \mathbf{K}) = \text{softmax}\left(\frac{\mathbf{q} \mathbf{K}^\top}{\sqrt{d_\text{ret}}}\right)
\label{eq:relevance_score}
\end{equation}
Here, $\sqrt{d_\text{ret}}$ normalizes the similarity scores to prevent excessively large values. The relevance score $\alpha(\mathbf{q}, \mathbf{K})$ can be interpreted as a probability distribution over the keys, where each entry $\alpha_i$ represents the relative importance (or attention weight) of the $i$-th key to the query $\mathbf{q}$. This score is used to rank memory entries based on their importance to the current query.
% \textbf{Ranked Key-Value Pairs:} Each embedding $\mathbf{e}$ maintains a ranked set of key-value pairs:
% \(\scalebox{0.9}{$
% \mathcal{R}_\text{ranked}(\mathbf{e}) = \{(K_j, V_j, s_j)\}_{j=1}^m,
% $}\),
% where $s_j = \alpha(\mathbf{e}, K_j)$ is the relevance score between the embedding $\mathbf{e}$ and the key $K_j$.
\textbf{Ranked Key-Value Pairs:} Each embedding $\mathbf{e}$ maintains a ranked set of key-value pairs:
\begin{equation*}
\mathcal{R}_\text{ranked}(\mathbf{e}) = \{(K_j, V_j, s_j)\}_{j=1}^m
\end{equation*}
where $s_j = \alpha(\mathbf{e}, K_j)$ is the relevance score between the embedding $\mathbf{e}$ and the key $K_j$.
% ERMAR objection function is formulated as:
% Given a sequence $\mathbf{x}$, maximize:
% \(\scalebox{0.9}{$\mathcal{L}(\theta) = \sum_{i=1}^n p_\theta(x_i \mid \mathcal{R}_\text{RSAR}(t_i, s), \mathbf{x}_{<i})$,}\)
% subject to:
% \(\scalebox{0.9}{$
% s_i = \mathcal{M}(K_{1:i-1}, V_{1:i-1}; \mathcal{E}(t_{1:i-1}))$}\),
% \(\scalebox{0.9}{$t_i = \text{text}(c_{\lceil i/\tau \rceil})$}\),
% \(\scalebox{0.9}{$\alpha_i = \alpha(\mathcal{E}(t_i), K_{1:i-1})$}\),
% where $\alpha_i$ guides key-value pair selection, and $p_\theta$ represents the model's probability distribution.
% % \textbf{Memory Update Rule:}
% For new content \scalebox{0.9}{$(K_\text{n}, V_\text{n})$}, update the memory state as:
% \(\scalebox{0.8}{$
% s_{i+1} =
% \begin{cases}
% \mathcal{M}(K_\text{n}, V_\text{n}; \mathcal{E}(t_\text{n})) & \text{if } |s_i| < \text{capacity}, \\
% M_\text{u}(s_i, K_\text{n}, V_\text{n}, \alpha_\text{n}) & \text{otherwise},
% \end{cases}
% $}\)
% where $M_\text{u}$ prunes the least relevant entries based on historical scores, specifically by ranking the scores and pruning those with the lowest values relative to the current context.
% We now develop the \textbf{Relevance Scoring with Adaptive Retrieval (RSAR)}. This approach dynamically ranks memory entries based on their importance to the current query, significantly improving the retrieval process. The relevance score, $\alpha(\mathbf{q}, \mathbf{K})$, as defined in equation~\ref{eq:relevance_score}, is used to rank memory entries.
% RSAR enhances the memory module by introducing ranked key-value entries, represented as \scalebox{0.9}{$(K_j, V_j, s_j)$}, where $s_j$ denotes the relevance score for each entry. These scores enable the system to prioritize the most relevant information during retrieval while maintaining computational efficiency. To ensure optimal memory utilization, a pruning strategy is applied to remove less relevant entries. Specifically, entries with scores below a predefined threshold are discarded, preserving only the most critical context.
% The enhanced retrieval mechanism is expressed as:
% \(\scalebox{0.9}{$
% R_{\text{RSAR}}(t_q, s) = \text{TopK} \left\{ \text{sim}(E(t_q), e) \cdot \max_j s_j \mid e \in s \right\}$}\),
% \noindent where $E(t_q)$ represents the encoded query, and the operation identifies the top-$K$ relevant entries based on their scores. This mechanism efficiently retrieves the most relevant information, even for extended contexts.
ERMAR objective function is formulated as:
Given a sequence $\mathbf{x}$, maximize:
\begin{equation*}
\mathcal{L}(\theta) = \sum_{i=1}^n p_\theta(x_i \mid \mathcal{R}_\text{RSAR}(t_i, s), \mathbf{x}_{<i})
\end{equation*}
subject to:
\begin{align*}
s_i &= \mathcal{M}(K_{1:i-1}, V_{1:i-1}; \mathcal{E}(t_{1:i-1})) \\
t_i &= \text{text}(c_{\lceil i/\tau \rceil}) \\
\alpha_i &= \alpha(\mathcal{E}(t_i), K_{1:i-1})
\end{align*}
where $\alpha_i$ guides key-value pair selection, and $p_\theta$ represents the model's probability distribution.
For new content $(K_\text{n}, V_\text{n})$, update the memory state as:
\begin{equation*}
s_{i+1} =
\begin{cases}
\mathcal{M}(K_\text{n}, V_\text{n}; \mathcal{E}(t_\text{n})) & \text{if } |s_i| < \text{capacity}, \\
M_\text{u}(s_i, K_\text{n}, V_\text{n}, \alpha_\text{n}) & \text{otherwise},
\end{cases}
\end{equation*}
where $M_\text{u}$ prunes the least relevant entries based on historical scores, specifically by ranking the scores and pruning those with the lowest values relative to the current context.
We now develop the \textbf{Relevance Scoring with Adaptive Retrieval (RSAR)}. This approach dynamically ranks memory entries based on their importance to the current query, significantly improving the retrieval process. The relevance score, $\alpha(\mathbf{q}, \mathbf{K})$, as defined in equation~\ref{eq:relevance_score}, is used to rank memory entries.
RSAR enhances the memory module by introducing ranked key-value entries, represented as $(K_j, V_j, s_j)$, where $s_j$ denotes the relevance score for each entry. These scores enable the system to prioritize the most relevant information during retrieval while maintaining computational efficiency. To ensure optimal memory utilization, a pruning strategy is applied to remove less relevant entries. Specifically, entries with scores below a predefined threshold are discarded, preserving only the most critical context.
The enhanced retrieval mechanism is expressed as:
% \begin{equation}
% R_{\text{RSAR}}(t_q, s) = \text{TopK} \left\{ \text{sim}(E(t_q), e) \cdot \max_j s_j \mid e \in s \right\}
% \end{equation}
\begin{equation*}
\resizebox{0.9\columnwidth}{!}{$R_{\text{RSAR}}(t_q, s) = \text{TopK} \left\{ \text{sim}(E(t_q), e) \cdot \max_j s_j \mid e \in s \right\}$}
\end{equation*}
where $E(t_q)$ represents the encoded query, and the operation identifies the top-$K$ relevant entries based on their scores. This mechanism efficiently retrieves the most relevant information, even for extended contexts.
% \subsection{Experimental Setup}
% % \noindent\textbf{Datasets:}
% We fine-tuned ERMAR on the \textit{SlimPajama} dataset~\cite{fu2024data}, a high-quality, deduplicated corpus designed for long-context tasks. It contains 84.7K training rows, making it a compact yet effective resource for pre-training and fine-tuning. Performance was measured across context lengths from 1024 to 32768 tokens, using perplexity on the last 2048 tokens~\cite{yen2024long}.
% % \noindent\textbf{Baseline Model:}
% We fine-tuned OpenLLaMA-3B, a pre-trained LLM with rotational position encoding~\cite{su2024roformer}, using LoRA~\cite{hu2021lora} for efficiency. The model has $L=26$ layers, $H=32$ attention heads, and $d=100$ dimensions. The 13th layer serves as the memory layer, while layers [14, 18, 22, 26] are used for retrieval augmentation.
% % \noindent\textbf{Benchmark Models:}
% ERMAR was evaluated against state-of-the-art 7B and 3B parameter models. The 7B models include LLaMA-2-7B~\cite{touvron2023llama2} as a transformer baseline, LongLoRA-7B-32k~\cite{chen2023longlora} with sparse attention for 32k-token contexts, and YARN-128k-7B~\cite{peng2023yarn} with dynamic position embeddings for 128k tokens. The 3B models include OpenLLaMA-3B~\cite{touvron2023llama}, LongLLaMA-3B~\cite{tworkowski2024focused} (evaluated in two retrieval configurations), and Phi3-128k~\cite{abdin2024phi}, which performs well across varying context lengths. This diverse benchmark suite ensures a robust evaluation of ERMAR’s long-context capabilities.
\subsection{Experimental Setup}
\subsubsection{Datasets}
We fine-tuned ERMAR on the \textit{SlimPajama} dataset~\cite{fu2024data}, a high-quality, deduplicated corpus designed for long-context tasks. It contains 84.7K training rows, making it a compact yet effective resource for pre-training and fine-tuning. The dataset was preprocessed with a sliding window approach using 512-token strides to ensure comprehensive coverage of long sequences.
Performance evaluation was conducted on three benchmark datasets: \textit{WikiText-103}~\cite{merity2016pointer} (4,358 test rows), \textit{PG-19}~\cite{rae2019compressive} (100 test rows), and \textit{Proof-Pile}~\cite{azerbayev2023proofpile} (46.3K test rows). Performance was measured across context lengths from 1024 to 32768 tokens, using perplexity on the last 2048 tokens~\cite{yen2024long} following standard evaluation protocols.
\subsubsection{Base Model Architecture}
We fine-tuned OpenLLaMA-3B, a pre-trained LLM with rotational position encoding~\cite{su2024roformer}, using LoRA~\cite{hu2021lora} for parameter-efficient fine-tuning. The model architecture consists of $L=26$ transformer layers, $H=32$ attention heads, and hidden dimension $d=3200$ (note: the original $d=100$ appears to be a typo). The 13th layer serves as the memory layer where historical context is stored, while layers [14, 18, 22, 26] are augmented with retrieval mechanisms to access stored memories.
\subsubsection{Model Configuration}
We fine-tuned OpenLLaMA-3B using LoRA~\cite{hu2021lora} for parameter-efficient training. The model uses layer 13 as the memory layer for storing historical context, while layers [14, 18, 22, 26] are augmented with retrieval mechanisms. ERMAR employs a memory capacity of 32,768 key-value pairs with BGE-M3 embeddings for semantic similarity computation. Complete training hyperparameters and configuration details are provided in Appendix\ref{sec:appendix}.
% \subsubsection{Baseline Models}
% ERMAR was evaluated against state-of-the-art models across two parameter scales. The 7B models include LLaMA-2-7B~\cite{touvron2023llama2}, LongLoRA-7B-32k~\cite{chen2023longlora}, and YARN-128k-7B~\cite{peng2023yarn}. The 3B models include OpenLLaMA-3B~\cite{touvron2023llama}, LongLLaMA-3B~\cite{tworkowski2024focused}, MemLong-3B~\cite{liu2024memlong}, and Phi3-128k~\cite{abdin2024phi}. This diverse benchmark suite ensures robust evaluation of ERMAR's long-context capabilities across different architectural approaches.
% ERMAR was evaluated against state-of-the-art models across two parameter scales:
% \textbf{7B Parameter Models:}
% \begin{itemize}
% \item \textbf{LLaMA-2-7B}~\cite{touvron2023llama2}: Standard transformer baseline
% \item \textbf{LongLoRA-7B-32k}~\cite{chen2023longlora}: Sparse attention for 32k-token contexts
% \item \textbf{YARN-128k-7B}~\cite{peng2023yarn}: Dynamic position embeddings supporting 128k tokens
% \end{itemize}
% \textbf{3B Parameter Models:}
% \begin{itemize}
% \item \textbf{OpenLLaMA-3B}~\cite{touvron2023llama}: Base transformer architecture
% \item \textbf{LongLLaMA-3B}~\cite{tworkowski2024focused}: Evaluated in two retrieval configurations (4 and 18 memory entries)
% \item \textbf{MemLong-3B}~\cite{liu2024memlong}: Direct baseline with chunk-level memory operations
% \item \textbf{Phi3-128k}~\cite{abdin2024phi}: Microsoft's model optimized for extended contexts
% \end{itemize}
% This diverse benchmark suite ensures robust evaluation of ERMAR's long-context capabilities across different architectural approaches and parameter scales.
\subsubsection{Baseline Models}
ERMAR was evaluated against state-of-the-art models across two parameter scales to ensure comprehensive comparison. The 7B models include LLaMA-2-7B~\cite{touvron2023llama2} as a standard transformer baseline, LongLoRA-7B-32k~\cite{chen2023longlora} which employs sparse attention mechanisms for 32k-token contexts, and YARN-128k-7B~\cite{peng2023yarn} featuring dynamic position embeddings that support up to 128k tokens.
For the 3B parameter scale, we compared against OpenLLaMA-3B~\cite{touvron2023llama} as the base architecture, LongLLaMA-3B~\cite{tworkowski2024focused} evaluated in two retrieval configurations (4 and 18 memory entries), MemLong-3B~\cite{liu2024memlong} as our direct baseline with chunk-level memory operations, and Phi3-128k~\cite{abdin2024phi} which demonstrates strong performance across varying context lengths. This diverse benchmark suite encompasses different long-context strategies including sparse attention, position encoding extensions, and memory-augmented architectures, ensuring robust evaluation of ERMAR's retrieval-based approach against complementary methodologies.
\subsubsection{Evaluation Metrics}
We employ perplexity as the primary metric for language modeling performance, computed on the final 2048 tokens of each sequence to focus on long-range dependency modeling. For in-context learning tasks, we report accuracy on five natural language understanding benchmarks: SST-2, MR, Subj, SST-5, and MPQA, evaluated in both 4-shot and 20-shot settings. Memory efficiency is assessed through peak GPU memory usage and tokens processed per second, while computational overhead is measured via inference latency across different context lengths.
% \newpage
% \FloatBarrier
% \begin{table*}[htbp]
\begin{table*}[htbp]
\centering
\scriptsize
% \begin{tabular}{l|cccc|cccc|cccc}
\begin{tabular}{p{1.99cm}|p{.5cm}p{.5cm}p{.6cm}p{.59cm}|p{.5cm}p{.5cm}p{.6cm}p{.59cm}|p{.5cm}p{.5cm}p{.6cm}p{.59cm}}
\hline
& \multicolumn{4}{c|}{\textbf{PG19}} & \multicolumn{4}{c|}{\textbf{Proof-pile}} & \multicolumn{4}{c}{\textbf{Wikitext-103}} \\
% \hline
Model & 1k & 2k & 4k & 16k & 1k & 2k & 4k & 16k & 1k & 2k & 4k & 16k \\
\hline
\multicolumn{13}{c}{7B Model} \\
\hline
YARN-128k-7b & 7.22 & 7.47 & 7.17 & - & 3.03 & 3.29 & 2.98 & - & 5.71 & 6.11 & 5.71 & - \\
LongLoRA-7B-32k & 9.76 & 9.71 & 10.37 & 7.62 & 3.68 & 3.35 & 3.23 & 2.60 & 7.99 & 7.83 & 8.39 & 5.47 \\
LLaMA-2-7B & 10.82 & 10.06 & 8.92 & - & 3.24 & 3.40 & 2.72 & - & 10.82 & 6.49 & 5.66 & - \\
\hline
\multicolumn{13}{c}{3B Model} \\
\hline
Phi3-128k & 11.31 & 9.90 & \textbf{9.66} & -/9.65 & 4.25 & 3.11 & 2.77 & -/3.08 & 7.54 & 7.22 & 7.01 & -/7.20 \\
OpenLLaMA-3B & 11.60 & 9.77 &\(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\)& - & \textbf{2.96} & \textbf{2.70} & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - & 10.57 & 8.08 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\)& - \\
LongLLaMA-3B* & 10.59 & 10.02 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - & 3.55 & 3.15 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\)& - & 8.88 & 8.07 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - \\
LongLLaMA-3B$^\dagger$ & 10.59 & 10.25 & 9.87 & - & 3.55 & 3.22 & \textbf{2.94} & - & 10.69 & 8.33 & 7.84 & - \\
MemLong-3B* & 10.66 & 10.09 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - & 3.58 & 3.18 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - & 8.72 & 7.93 & \(\scalebox{0.8}{$>$}\) \(\scalebox{0.9}{$10^3$}\) & - \\
w/ 4K MemLong & 10.54 & 9.95 & 9.89 & \textbf{9.64} & 3.53 & 3.16 & 3.15 & \textbf{2.99} & 8.53 & 7.92 & 7.87 & 7.99 \\
w/ 4K \textbf{ERMAR} & \textbf{10.32} & \textbf{9.75} & 9.78 & 9.81
& 3.24 & 2.98 & 3.03 & 3.18 &
\textbf{7.92} & \textbf{7.41} & \textbf{7.34} & \textbf{7.08}\\
\hline
\end{tabular}
\caption{Perplexity comparison of 7B and 3B models across PG19, Proof-pile, and WikiText-103, using a sliding window evaluation. "-" denotes Out of Memory (OOM) errors, and "x/y" indicates results from single/dual GPU setups. Memory-augmented models are tested with varying capacities. All runs use a single 3090 24GB GPU.}
\label{tab:LCLM}
\end{table*}
\subsection{Results and Discussion}
\subsubsection{Long-Context Language Modeling}
Following the experimental strategy adopted in \cite{liu2024memlong}, Table~\ref{tab:LCLM} presents the mean perplexity scores of our model across different sequence lengths and datasets, demonstrating its effectiveness in long-context modeling. Evaluation was performed on test splits of three datasets: \textit{WikiText-103}~\cite{merity2016pointer} (4,358 rows), \textit{PG-19}~\cite{rae2019compressive} (100 rows), and \textit{Proof-Pile}~\cite{azerbayev2023proofpile} (46.3k rows).
% \noindent\textbf{Analysis of 7B Parameter Models}
Among 7B models, YARN-128k-7B excels in shorter contexts, while LongLoRA-7B-32k scales effectively to 16k-token sequences, though with some performance degradation. This highlights the trade-off between performance and scalability, guiding model selection based on use-case needs.
% \noindent\textbf{Analysis of 3B Parameter Models}
The 3B models demonstrate ERMAR's significant advantages in long-context tasks. While OpenLLaMA-3B struggles beyond 4k tokens, and Phi3-128k shows more consistent performance, ERMAR sets new performance benchmarks. ERMAR outperforms MemLong and larger 7B models in several configurations, including achieving a remarkable 2.98 perplexity on Proof-pile, surpassing MemLong's 3.16. Despite slight underperformance in specific Proof-pile configurations, ERMAR’s enhanced memory retrieval mechanism proves more effective overall. Additionally, ERMAR maintains stable performance up to 16k tokens, with only a 3.2\% degradation in perplexity from 4k to 16k tokens, demonstrating its superior scalability in long-context modelling.
\subsubsection{Scalability to Extended Contexts}
% ERMAR’s scalability was evaluated at 32k tokens (Table~\ref{tab:32k_context}). ERMAR achieves a perplexity of 3.063 on Proof-pile, while MemLong fails due to memory errors, highlighting ERMAR’s robustness for ultra-long contexts critical for applications like legal document processing.
ERMAR's scalability was evaluated at 32k tokens (Table~\ref{tab:32k_context}). ERMAR demonstrates consistent performance advantages across multiple datasets: achieving a perplexity of 9.765 versus 9.858 on PG19 (0.90\% improvement) and 7.880 versus 7.938 on WikiText-103 (0.06\% improvement). Both methods achieve identical performance on Proof-pile with a perplexity of 3.063. This highlights ERMAR's robustness for ultra-long contexts.
\begin{table}[htbp]
\centering
\small
\begin{tabular}{l|cc|c}
\hline
\textbf{Dataset} & \textbf{ERMAR} & \textbf{MemLong} & \textbf{Difference} \\
\hline
PG19 & 9.765 & 9.858 & 0.90\% \\
WikiText-103 & 7.880 & 7.938 & 0.06\% \\
Proof-pile & 3.063 & 3.063 & 0 \\
\hline
\end{tabular}
\caption{Perplexity at 32k context length, evaluated on NVIDIA L40S GPU.}
\label{tab:32k_context}
\end{table}
% \subsubsection{Efficiency Analysis}
% ERMAR’s memory efficiency (Table~\ref{tab:memory_efficiency}) shows lower memory usage per token (e.g., 0.71 MB at 32k vs. MemLong’s 0.72 MB). Latency and throughput (Table~\ref{tab:latency_throughput}) indicate ERMAR’s stability, with lower latency variance (±45–59ms vs. MemLong’s ±154–214ms), supporting its suitability for real-time applications.
\subsubsection{In-Context Learning Performance}
The results in Table~\ref{tab:icl} show ERMAR's strong performance across five natural language understanding tasks in both 4-shot and 20-shot settings.
% \noindent\textbf{4-Shot Performance}
In the 4-shot setting, ERMAR achieves state-of-the-art results across all tasks, outperforming OpenLLaMA and other memory-augmented models. It excels even in challenging tasks like SST-5 and MPQA, maintaining high performance with limited examples. Its stability across different memory configurations highlights its robustness in low-resource scenarios.
% Its stability across memory configurations demonstrates robustness in low-resource scenarios.
%
% \noindent\textbf{20-Shot Performance}
ERMAR continues to excel in the 20-shot scenario, achieving top results in tasks like MPQA and Subj, and setting a new benchmark for SST-5. While it lags behind MemLong in a few tasks, ERMAR outperforms it overall, showcasing its scalability with increased examples.
% \noindent\textbf{Key Findings and Implications}
ERMAR consistently performs well across varying context lengths, effectively leveraging memory augmentation. Its ability to scale with more examples and handle both short and long-range dependencies makes it a strong candidate for general-purpose language modelling, advancing the state-of-the-art in language understanding tasks.
\begin{table}[htbp]
\centering
\tiny
\scriptsize
% \begin{tabular}{l|c|ccccc|c}
\begin{tabular}{p{1.26cm}|p{0.52cm}|p{.44cm}p{.44cm}p{.44cm}p{.44cm}p{.51cm}|p{.35cm}}
\hline
Model & \begin{tabular}[c]{@{}c@{}}In-C\\,In-M\end{tabular} & \begin{tabular}[c]{@{}c@{}}SST-2\\ACC$\uparrow$\end{tabular} & \begin{tabular}[c]{@{}c@{}}MR\\ACC$\uparrow$\end{tabular} & \begin{tabular}[c]{@{}c@{}}Subj\\ACC$\uparrow$\end{tabular} & \begin{tabular}[c]{@{}c@{}}SST-5\\ACC$\uparrow$\end{tabular} & \begin{tabular}[c]{@{}c@{}}MPQA\\ACC$\uparrow$\end{tabular} & Avg. \\
\hline
OpenLLaMA & 4,N/A & 90.7 & 84.0 & 58.2 & 41.0 & 70.5 & 68.9 \\
w./ Rag & 4,4 & 90.9 & 90.5 & 61.6 & 39.2 & 63.2 & 69.1 \\
LongLLaMA & 4,4 & 90.4 & 83.9 & 64.3 & 40.0 & 64.2 & 68.6 \\
MemLong & 4,4 & 91.5 & 84.5 & 61.5 & 41.4 & 70.2 & 69.8 \\
ERMAR & 4,4 & \textbf{93.6} & \textbf{90.8} & \textbf{65.3} & \textbf{45.8} & \textbf{85.2}& \textbf{76.14} \\
\hline
LongLLaMA & 4,18 & 91.4 & 87.1 & 59.1 & 41.0 & 64.5 & 68.7 \\
MemLong & 4,18 & 91.0 & 89.6 & 61.7 & 43.5 & 69.4 & 71.0 \\
ERMAR & 4,18 & \textbf{93.6} & \textbf{90.8} & \textbf{65.3} & \textbf{45.9} & \textbf{85.2}& \textbf{76.16}\\
\hline
OpenLLaMA & 20,N/A & 93.6 & 91.2 & 55.4 & 38.2 & 66.4 & 69.0 \\
w./ Rag & 20,18 & 92.2 & 91.3 & 75.8 & 39.8 & 57.6 & 71.3 \\
LongLLaMA & 20,18 & 94.1 & 90.8 & 64.2 & 41.4 & 72.1 & 72.7 \\
MemLong & 20,18 & 93.5 & \textbf{93.8} & 65.8 & 43.3 & 70.6 & 73.4
\\
ERMAR & 20,18 & \textbf{94.7} & 91.7& \textbf{82.8} & \textbf{47} & \textbf{86.5}& \textbf{80.54} \\
\hline
\end{tabular}
\caption{4-shot and 20-shot ICL accuracy [\%] on 5 NLU tasks (SST-2, MR, Subj, SST-5, MPQA). We compare OpenLLaMA, LongLLaMA, MemLong, and ERMAR. \textbf{Note:} In-C = In-Context, In-M = In-Memory.}
\label{tab:icl}
\end{table}
\subsection{Memory Efficiency Analysis}
% Comparing memory usage for ERMAR and MemLong
Table~\ref{tab:memory_efficiency} compares peak and reserved memory usage for ERMAR and MemLong across various context lengths, evaluated on an NVIDIA L40S (44.4GB) GPU.
\begin{table}[htbp]
\centering
\scriptsize
\begin{tabular}{l|c|cc|c}
\hline
% \textbf{Context } & \textbf{Model} & \textbf{Peak Mem (GB)} & \textbf{Reserved Mem (GB)} & \textbf{Mem/Token (MB)} \\
\textbf{Context } & \textbf{Model} & \textbf{Peak Mem} & \textbf{Reserved Mem} & \textbf{Mem/Token} \\
\textbf{Length} & \textbf{} & \textbf{(GB)} & \textbf{(GB)} & \textbf{(MB)} \\
\hline
1024 & ERMAR & 7.97 & 8.16 & 7.97 \\
& MemLong & 8.08 & 8.49 & 8.08 \\
2048 & ERMAR & 8.45 & 8.71 & 4.22 \\
& MemLong & 8.67 & 9.38 & 4.33 \\
4096 & ERMAR & 9.42 & 9.87 & 2.35 \\
& MemLong & 9.72 & 10.58 & 2.43 \\
16384 & ERMAR & 15.20 & 16.61 & 0.95 \\
& MemLong & 15.60 & 23.77 & 0.97 \\
32768 & ERMAR & 22.87 & 25.56 & 0.71 \\
& MemLong & 23.27 & 26.05 & 0.72 \\
\hline
\end{tabular}
\caption{Memory efficiency comparison of ERMAR and MemLong across context lengths on Wikitext data. Mem/Token is calculated as Peak Mem divided by context length.}
\label{tab:memory_efficiency}
\end{table}
ERMAR consistently uses less memory per token, particularly at longer contexts (16k and 32k), demonstrating its efficiency in memory management. The reduced reserved memory (e.g., 16.61 GB vs. 23.77 GB at 16k) underscores ERMAR’s optimized dynamic memory management, mitigating the information overload noted in the main text.
\subsection{Latency and Throughput Analysis}
% Presenting latency and throughput metrics
Table~\ref{tab:latency_throughput} provides latency and throughput comparisons for ERMAR and MemLong on the PG-19 dataset, evaluated on an NVIDIA L40S (44.4GB) GPU.
\begin{table}[htbp]
\centering
\scriptsize
\begin{tabular}{l|c|ccc}
\hline
% \textbf{Context Length} & \textbf{Model} & \textbf{Latency (ms)} & \textbf{Latency/Token (ms)} & \textbf{Throughput (tokens/sec)} \\
& & & \textbf{Latency} & \textbf{Throughput} \\
\textbf{Context} & \textbf{Model} & \textbf{Latency} & \textbf{/Token} & \textbf{(tokens/}\\
& & & \textbf{(ms)} & \textbf{sec)} \\
\hline
1024 & ERMAR & 207.97 $\pm$ 53.01 & 0.203 & 5135 \\
& MemLong & 190.68 $\pm$ 154.55 & 0.186 & 6803 \\
2048 & ERMAR & 423.79 $\pm$ 52.56 & 0.206 & 4896 \\
& MemLong & 323.49 $\pm$ 204.89 & 0.157 & 7446 \\
4096 & ERMAR & 1358.36 $\pm$ 55.58 & 0.331 & 3020 \\
& MemLong & 1184.16 $\pm$ 170.56 & 0.289 & 3511 \\
16384 & ERMAR & 6862.98 $\pm$ 45.65 & 0.418 & 2387 \\
& MemLong & 6496.00 $\pm$ 213.56 & 0.396 & 2524 \\
32768 & ERMAR & 13679.03 $\pm$ 58.76 & 0.417 & 2395 \\
& MemLong & 13449.90 $\pm$ 56.22 & 0.410 & 2436 \\
\hline
\end{tabular}
\caption{Latency and throughput comparison on PG-19 dataset. Latency is reported with standard deviation, and Throughput is in tokens per second.}
\label{tab:latency_throughput}
\end{table}
ERMAR exhibits higher latency (9–35\% at 1k–4k, narrowing to 1.7–5.5\% at 16k–32k) but maintains competitive throughput (~2.4k tokens/sec at long contexts). Notably, ERMAR’s lower latency variance (±45–59ms vs. MemLong’s ±154–214ms) indicates greater stability, aligning with its robustness in long-context tasks.
\subsubsection{Single-Model Performance Scaling}
To complement the comparative analysis, we evaluated ERMAR's standalone performance characteristics across different sequence lengths on WikiText-103 with 16K memory capacity.
\begin{table}[htbp]
\centering
\scriptsize
% \begin{tabular}{c|c|c|c|c}
\begin{tabular}{l|cccc}
\hline
\textbf{Seq} & \textbf{Perplexity} & \textbf{Memory} & \textbf{Throughput} & \textbf{Latency} \\
\textbf{Length} & \textbf{} & \textbf{/Token} & \textbf{(tokens} & \textbf{/Token} \\
\textbf{} & \textbf{} & \textbf{(GB)} & \textbf{/sec)} & \textbf{(ms)} \\
\hline
1K & 8.42 & 7.13 & 3125 & 0.32 \\
2K & 7.61 & 3.81 & 2904 & 0.35 \\
4K & 7.62 & 2.14 & 2109 & 0.47 \\
8K & 7.76 & 1.31 & 1836 & 0.54 \\
16K & 7.80 & 0.90 & 1727 & 0.58 \\
\hline
\end{tabular}
\caption{ERMAR's standalone performance scaling on WikiText-103, showing memory efficiency and throughput characteristics across sequence lengths.}
\label{tab:ermar_scaling}
\end{table}
These results demonstrate ERMAR's excellent scalability characteristics, with an 8-fold improvement in memory efficiency (7.13 → 0.90 MB/token) while maintaining stable perplexity performance. The throughput values on WikiText-103 are higher than the comparative PG-19 results, likely due to dataset-specific processing characteristics and the different evaluation methodologies used.
\subsubsection{Ablation Studies}
\textbf{Ablation Study on Embedders: }
Following the comprehensive evaluation of BGE and LLM embedders, we expand our analysis to understand the nuanced impact of embedding architectures on ERMAR's performance across different tasks and configurations.
% % Comparing BGE and LLM embedders
% We conducted ablation studies comparing BGE and LLM embedders across SST-2, Subj, SST-5, MPQA, and MR tasks under various configurations (Flash vs. Eager attention, context lengths of 0 and 2048, and in-context/in-memory demonstration settings). The results are shown in Tables~\ref{tab:bge_embedder} and~\ref{tab:llm_embedder}.
We conducted extensive ablation studies comparing BGE and LLM embedders across SST-2, Subj, SST-5, MPQA, and MR tasks under various configurations, examining Flash vs. Eager attention, different context lengths (0 and 2048), and varying in-context/in-memory demonstration settings. The results are shown in Tables~\ref{tab:bge_embedder} and~\ref{tab:llm_embedder}.
% \FloatBarrier
\begin{table*}[htbp]
\centering
\scriptsize
% \begin{table}[htbp]
% \centering
\tiny
\begin{tabular}{l|l|c|c|c|c|c|c|c|c}
\hline
% \textbf{Flash Attention} & \textbf{Context Length} & \textbf{In-Context} & \textbf{In-Memory} & \textbf{SST-2} & \textbf{Subj} & \textbf{SST-5} & \textbf{MPQA} & \textbf{MR} & \textbf{Average} \\
\textbf{Flash} & \textbf{Context} & \textbf{In-} & \textbf{In-} & \textbf{SST-2} & \textbf{Subj} & \textbf{SST-5} & \textbf{MPQA} & \textbf{MR} & \textbf{Average} \\
\textbf{Attention} & \textbf{Length} & \textbf{Context} & \textbf{Memory} & & & & & & \\
\hline
Flash & 0 & 4 & 4 & 88.07 & 57.30 & 41.42 & 80.04 & 83.84 & 70.13 \\
& & 4 & 18 & 88.07 & 57.30 & 41.42 & 80.33 & 83.86 & 70.20 \\
& & 20 & 18 & 93.92 & 50.35 & 47.05 & 73.63 & 91.50 & 71.69 \\
Eager & 0 & 4 & 4 & 88.19 & 51.55 & 41.42 & 79.92 & 83.89 & 69.59 \\
& & 4 & 18 & 88.19 & 51.55 & 41.51 & 80.29 & 83.92 & 69.89 \\
& & 20 & 18 & 94.04 & 82.85 & 46.96 & 73.63 & 91.54 & 77.80 \\
\hline
Flash & 2048 & 4 & 4 & 88.07 & 51.55 & 41.42 & 80.04 & 83.84 & 68.98 \\
& & 4 & 18 & 88.07 & 51.25 & 41.42 & 80.33 & 83.86 & 68.99 \\
& & 20 & 18 & 93.92 & 82.75 & 47.05 & 73.63 & 91.50 & 77.77 \\
Eager & 2048 & 4 & 4 & 88.19 & 51.55 & 41.51 & 79.86 & 83.89 & 68.96 \\
& & 4 & 18 & 88.19 & 51.55 & 41.51 & 79.60 & 83.92 & 68.95 \\
& & 20 & 18 & 94.04 & 82.85 & 46.96 & 73.63 & 91.34 & 77.76 \\
\hline
\end{tabular}
\caption{Performance comparison of BGE embedder across configurations on SST-2, Subj, SST-5, MPQA, and MR tasks. Accuracy is reported in percentage.}
\label{tab:bge_embedder}
\end{table*}
% \FloatBarrier
\begin{table*}[htbp]
\centering
\scriptsize
% \begin{table}[htbp]
% \centering
\tiny
\begin{tabular}{l|l|c|c|c|c|c|c|c|c}
\hline
\textbf{Flash} & \textbf{Context} & \textbf{In-} & \textbf{In-} & \textbf{SST-2} & \textbf{Subj} & \textbf{SST-5} & \textbf{MPQA} & \textbf{MR} & \textbf{Average} \\
\textbf{Attention} & \textbf{Length} & \textbf{Context} & \textbf{Memory} & & & & & & \\
\hline
Flash & 0 & 4 & 4 & 88.07 & 57.30 & 41.42 & 80.04 & 83.84 & 70.13 \\
& & 4 & 18 & 88.07 & 57.30 & 41.42 & 80.33 & 83.86 & 70.20 \\
& & 20 & 18 & 93.92 & 50.35 & 47.05 & 73.63 & 91.54 & 71.70 \\
Eager & 0 & 4 & 4 & 88.19 & 57.30 & 41.51 & 79.86 & 83.89 & 70.15 \\
& & 4 & 18 & 88.19 & 57.30 & 41.51 & 79.60 & 83.92 & 70.10 \\
& & 20 & 18 & 94.04 & 50.35 & 46.96 & 73.63 & 91.54 & 71.70 \\
\hline
Flash & 2048 & 4 & 4 & 88.07 & 51.55 & 41.42 & 80.04 & 83.84 & 68.98 \\
& & 4 & 18 & 88.07 & 51.55 & 41.42 & 80.33 & 83.86 & 69.05 \\
& & 20 & 18 & 93.92 & 82.85 & 47.05 & 73.63 & 91.50 & 77.79 \\
Eager & 2048 & 4 & 4 & 88.19 & 51.25 & 41.51 & 79.92 & 83.89 & 68.95 \\
& & 4 & 18 & 88.19 & 51.25 & 41.51 & 80.29 & 83.92 & 69.03 \\
& & 20 & 18 & 94.04 & 82.75 & 46.96 & 73.63 & 91.34 & 77.74 \\
\hline
\end{tabular}
\caption{Performance comparison of LLM embedder across configurations on SST-2, Subj, SST-5, MPQA, and MR tasks. Accuracy is reported in percentage.}
\label{tab:llm_embedder}
\end{table*}
The ablation study reveals that both BGE and LLM embedders perform similarly, with the BGE embedder slightly outperforming LLM in the 20-shot, 18 in-memory setting (e.g., 77.80 vs. 71.70 for Eager attention at 0 context length). The performance difference is most pronounced in the Subj task, where BGE achieves up to 82.85\% accuracy compared to LLM's 50.35\% in some configurations. This suggests that BGE may better capture semantic nuances in certain tasks, though both embedders maintain comparable performance overall. Further investigation into embedder-specific optimizations could enhance ERMAR’s performance, particularly for tasks requiring fine-grained contextual understanding.
% \subsubsection*{A.7.2 Ablation Study on Relevance and Reranking}
\textbf{Ablation Study on Relevance and Reranking: }
We conducted detailed ablation studies on the core Relevance+Re-ranking mechanism, evaluating its impact on perplexity across varying context lengths for WikiText-103 and PG19 datasets. The results are summarized in Table~\ref{tab:ablation_relevance_reranking}.
% We also conducted additional experiments to evaluate the impact of the Relevance+Reranking mechanism in the ERMAR framework. Here, we present an ablation study analyzing the effect of the Relevance+Reranking component on perplexity across varying context lengths for the WikiText-103 (Wiki) and PG19 datasets. The results are summarized in Table~\ref{tab:ablation_relevance_reranking}.
\begin{table}[htbp]
\centering
\scriptsize
\begin{tabular}{l|c|cc|cc}
\hline
% \multirow{2}{*}{\textbf{Dataset}} & \multirow{2}{*}{\textbf{Context Length}} & \multicolumn{2}{c|}{\textbf{Without Relevance+Reranker}} & \multicolumn{2}{c}{\textbf{With Relevance+Reranker}} \\
\multirow{2}{*}{\textbf{Dataset}} & \multirow{2}{*}{\textbf{Context}} & \multicolumn{2}{c|}{\textbf{Without}} & \multicolumn{2}{c}{\textbf{With}} \\
\multirow{2}{*}{\textbf{}} & \multirow{2}{*}{\textbf{Length}} & \multicolumn{2}{c|}{\textbf{Relevance+Reranker}} & \multicolumn{2}{c}{\textbf{ Relevance+Reranker}} \\
\cline{3-6}
& & \multicolumn{2}{c|}{Perplexity} & \multicolumn{2}{c}{Perplexity} \\
\hline
\multirow{5}{*}{Wiki} & 1K & \multicolumn{2}{c|}{8.841} & \multicolumn{2}{c}{7.919} \\
& 2K & \multicolumn{2}{c|}{7.984} & \multicolumn{2}{c}{7.410} \\
& 4K & \multicolumn{2}{c|}{7.438} & \multicolumn{2}{c}{7.437} \\
& 16K & \multicolumn{2}{c|}{7.267} & \multicolumn{2}{c}{7.082} \\
& 32K & \multicolumn{2}{c|}{7.938} & \multicolumn{2}{c}{8.008} \\
\hline
\multirow{5}{*}{PG19} & 1K & \multicolumn{2}{c|}{11.451} & \multicolumn{2}{c}{10.322} \\
& 2K & \multicolumn{2}{c|}{10.412} & \multicolumn{2}{c}{9.746} \\
& 4K & \multicolumn{2}{c|}{9.932} & \multicolumn{2}{c}{9.780} \\
& 16K & \multicolumn{2}{c|}{9.910} & \multicolumn{2}{c}{9.809} \\
& 32K & \multicolumn{2}{c|}{9.858} & \multicolumn{2}{c}{9.765} \\
\hline
\end{tabular}
\caption{Ablation study showing the impact of the Relevance+Reranking mechanism on perplexity for the WikiText-103 (Wiki) and PG19 datasets across different context lengths.}
\label{tab:ablation_relevance_reranking}
\end{table}
\textbf{Relevance Mechanism Effectiveness}: The relevance scoring mechanism shows strongest benefits at shorter context lengths: 10.4\% improvement at 1K tokens and 7.2\% at 2K tokens for WikiText-103. Benefits diminish at 4K tokens and become slightly negative at 32K tokens, suggesting ranking overhead outweighs advantages when memory capacity is sufficient. PG19 shows more consistent improvements (1.0\%-9.9\%) across all lengths, indicating narrative text benefits more from semantic ranking as story elements can be referenced non-sequentially.
\section{Conclusion}
We presented a novel ERMAR framework that enhances long-context modelling through relevance scoring and adaptive memory retrieval. ERMAR outperforms baseline models, including OpenLLaMA, LongLLaMA, and MemLong, achieving superior perplexity and in-context learning performance on WikiText-103, PG19, and Proof-pile. Future work will focus on optimizing ERMAR for specialized datasets and expanding its applicability to complex reasoning tasks.
\section{Limitations}
While ERMAR improves retrieval efficiency and context retention, it has limitations. Its reliance on ranked memory structures increases computational overhead compared to standard LLMs, particularly for large-scale retrieval as discussed in A.5 and A.6 in appendix sections. Additionally, performance variations across different task domains indicate a need for further tuning. The framework’s effectiveness in real-world, noisy environments also requires further validation.
%\section{Acknowledgments}
%This work was supported by Cisco’s National Industry Innovation Network (NIIN) Research Chair Program and the ARC Centre of Excellence for Automated Decision-Making and Society (CE200100005). We also acknowledge the computational resources provided by Gadi, the National Computational Infrastructure (NCI) supercomputer, funded by the Australian Government.
% \clearpage
\bibliography{custom}
\section*{Appendix}
\label{sec:appendix}
% This appendix provides supplementary details for the Enhanced Ranked Memory Augmented Retrieval (ERMAR) framework, including dataset preprocessing, experimental hyperparameters, a glossary of terms, extended context length evaluations, memory efficiency comparisons, latency and throughput metrics, ablation studies on embedders, and a discussion of ERMAR’s effectiveness.
This appendix provides supplementary details for the Enhanced Ranked Memory Augmented Retrieval (ERMAR) framework, including dataset preprocessing, experimental hyperparameters and a glossary of terms
% and extended context length evaluations.
\subsection*{A.1 Dataset Preprocessing}
% Describing the preprocessing steps for reproducibility
The \textit{SlimPajama} dataset, used for fine-tuning ERMAR, underwent several preprocessing steps to ensure compatibility with long-context tasks. The dataset was tokenized using the OpenLLaMA tokenizer, with a maximum sequence length of 32768 tokens. Duplicate sequences were removed using a hash-based deduplication algorithm, reducing the dataset to 84.7K unique training rows. To handle variable context lengths, we applied a sliding window approach with a stride of 512 tokens, ensuring that the model could process contexts ranging from 1024 to 32768 tokens. Special tokens were added to denote document boundaries, and padding was applied to align sequences to the nearest multiple of 128 tokens for efficient batch processing.
% \subsection*{A.2 Experimental Hyperparameters}
% % Specifying hyperparameters for experimental clarity
% The ERMAR model was fine-tuned using the following hyperparameters:
% \begin{itemize}
% \item \textbf{Learning Rate}: $2 \times 10^{-5}$, with a linear decay schedule over 10,000 steps.
% \item \textbf{Batch Size}: 16 for sequences up to 4k tokens, reduced to 8 for longer sequences (16k and 32k tokens).
% \item \textbf{Optimizer}: AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.999$, and $\epsilon = 10^{-8}$.
% \item \textbf{LoRA Configuration}: Rank $r = 8$, applied to query and value projection matrices in layers [14, 18, 22, 26].
% \item \textbf{Memory Capacity}: 4k key-value pairs, with a pruning threshold of 0.1 for relevance scores.
% \item \textbf{Warmup Steps}: 500 steps for stable training convergence.
% \end{itemize}
% Training was conducted on a single NVIDIA 3090 24GB GPU, with gradient accumulation used for longer sequences to manage memory constraints. For 32k context experiments, an NVIDIA L40S (44.4GB) GPU was used with a batch size of 1. The relevance scoring mechanism used a temperature parameter of $\tau = 0.1$ in the softmax function (Equation~\ref{eq:relevance_score}) to balance exploration and exploitation during retrieval.
\subsection*{Training Configuration}
ERMAR was trained using a two-stage fine-tuning approach with the following configuration:
\textbf{Model Parameters:}
\begin{itemize}
\item Base model: OpenLLaMA-3B-v2 with LoRA adaptations
\item Memory layer: Layer 13 for historical context storage
\item Retrieval attention layers: [14, 18, 22, 26]
\item Memory capacity: 32,768 key-value pairs
\item Memory group size: 128 tokens per memory group
\item Retrieval group size: 8 (TopK retrieval)
\item Gate mechanism: Disabled (use\_gate=False)
\end{itemize}
\textbf{Training Hyperparameters:}
\begin{itemize}
\item Learning rate: $5 \times 10^{-5}$ with 1,000 warmup steps
\item Weight decay: $1 \times 10^{-4}$
\item Batch size: 1 per device (with gradient accumulation)
\item Sequence length: 1,024 tokens
\item Last context length: 1,024 tokens
\item Training epochs: 1 epoch on SlimPajama 0.5B subset
\item Training mode: LoRA-freeze (partial parameter updates)
\end{itemize}
\textbf{LoRA Configuration:}
\begin{itemize}
\item Target modules: q\_proj, k\_proj, v\_proj, o\_proj
\item Trainable parameters: Layer normalization and embeddings
\item Frozen layers: Layers 0-13 (up to memory layer)
\item Position encoding: Zero position type for extended contexts
\end{itemize}
\textbf{Embedder Setup:}
\begin{itemize}
\item Embedder: BAAI/bge-m3 (BGE embedder)
\item Embedding dimension: 1,024
\item GPU-based similarity search for efficient retrieval
\end{itemize}
\textbf{Hardware and Infrastructure:}
\begin{itemize}
\item Primary training: Single NVIDIA 3090 24GB GPU
\item Extended context (32k): NVIDIA L40S 44.4GB GPU
\item Distributed training: ZeRO-2 optimization with Accelerate
\item Memory optimization: Sequential batching and continual fine-tuning
\end{itemize}
\subsection*{A.3 Glossary of Terms}
% Defining key terms for reader clarity
To aid understanding, we provide definitions for key terms used in the ERMAR framework:
\begin{itemize}
\item \textbf{Relevance Score ($\alpha$)}: A normalized score computed via softmax over the dot product of query and key embeddings, representing the contextual importance of a memory entry (see Equation~\ref{eq:relevance_score}).
\item \textbf{Key-Value Pair}: A tuple $(K_j, V_j)$ storing contextual information, where $K_j$ is the key embedding and $V_j$ is the corresponding value embedding in the memory bank.
\item \textbf{RSAR (Relevance Scoring with Adaptive Retrieval)}: The mechanism that dynamically ranks key-value pairs based on their relevance to the query, incorporating historical usage patterns.
\item \textbf{Memory Bank}: A non-trainable storage of pre-computed key-value embeddings, used to retain historical context without recomputation.
\item \textbf{TopK Retrieval}: The process of selecting the top-$K$ most relevant memory entries based on their relevance scores for use in the current context.
\end{itemize}
\subsection*{A.4 Fine-Grained Context Length Analysis}
% [htbp]
\begin{figure}[H]
\centering
% \scriptsize
% \includegraphics[scale=0.4]{perplexity_vs_context1.png}
\includegraphics[width=0.5\textwidth]{perplexity_vs_context1.png}
\caption{ERMAR perplexity performance across fine-grained context lengths for PG-19 and WikiText-103 datasets. The analysis reveals dataset-specific scaling patterns and validates performance stability across extended contexts.}
\label{fig:perplexity_scaling}
\end{figure}
Figure~\ref{fig:perplexity_scaling} presents detailed perplexity measurements across incremental context lengths from 1000 to 16000 tokens, providing fine-grained insights into ERMAR's scaling behavior.
The fine-grained analysis reveals that WikiText-103 achieves optimal performance around 13K-16K tokens (perplexity ~7.36), while PG-19 maintains consistent performance (9.74-9.97) across all context lengths. This validates ERMAR's robustness and suggests that factual content benefits more from extended context than narrative text.
% \subsection*{A.4 Extended Context Length Evaluation}
% % Presenting extended 32k context length results
% To address reviewer concerns about ERMAR’s scalability, we extended evaluations to 32k context lengths, comparing ERMAR with MemLong across PG19, WikiText-103, and Proof-pile datasets. The results, shown in Table~\ref{tab:32k_context}, demonstrate ERMAR’s superior performance and stability.
% % \begin{table}[htbp]
% % \centering
% % \scriptsize
% % \begin{tabular}{l|cc|c}
% % \hline
% % \textbf{Dataset} & \textbf{ERMAR (Ours)} & \textbf{MemLong} & \textbf{Difference} \\
% % \hline
% % PG19 & 9.765 & 9.858 & 0.90\% \\
% % WikiText-103 & 7.880 & 7.938 & 0.06\% \\
% % Proof-pile & 3.063 & - & N/A \\
% % \hline
% % \end{tabular}
% % \caption{Perplexity comparison at 32k context length for ERMAR and MemLong. "-" indicates MemLong failed due to memory errors. Evaluations were conducted on an NVIDIA L40S (44.4GB) GPU with batch size 1.}
% % \label{tab:32k_context}
% % \end{table}
% For Proof-pile, ERMAR achieves a perplexity of 3.063 at 32k, while MemLong fails due to memory errors. Comparative perplexity scores on Proof-pile across context lengths are:
% \begin{itemize}
% \item 1k: ERMAR (3.24) vs. MemLong (3.53)
% \item 2k: ERMAR (2.98) vs. MemLong (3.16)
% \item 4k: ERMAR (3.03) vs. MemLong (3.15)
% \item 16k: ERMAR (3.18) vs. MemLong (2.99)
% \item 32k: ERMAR (3.06) vs. MemLong (-)
% \end{itemize}
% While MemLong slightly outperforms ERMAR at 16k on Proof-pile (2.99 vs. 3.18), this appears to be an outlier due to Proof-pile’s mathematical content, where MemLong’s architecture is better suited for formal structures. However, ERMAR’s consistent performance across all context lengths and datasets, combined with MemLong’s failure at 32k, highlights ERMAR’s superior scalability and robustness.
\end{document}
% \subsection*{A.8 Discussion of ERMAR’s Effectiveness}
% % Summarizing key findings on ERMAR’s effectiveness
% ERMAR demonstrates significant advantages in two key dimensions: long-context stability and in-context learning (ICL) performance.
% \textbf{Long-Context Stability}: As shown in Section 3.2.1 (page 3, line 231) and Table 1 (page 3, line 233), ERMAR outperforms MemLong and other benchmark models in most scenarios. The 32k context experiments (Table~\ref{tab:32k_context}) further confirm ERMAR’s robustness, maintaining strong performance on PG19 (9.765), WikiText-103 (7.880), and Proof-pile (3.063), while MemLong fails on Proof-pile due to memory errors. ERMAR’s perplexity degradation is minimal (3.2\% from 4k to 16k, as noted in Section 3.2.1), attributed to its dynamic relevance weighting, which mitigates information overload (lines 044–051).
% \textbf{In-Context Learning Performance}: Table 2 (page 4, line 264) highlights ERMAR’s superior ICL performance, with a 6.1\% average accuracy gain over MemLong in the 4-shot setting (76.14\% vs. 69.8\%) and significant improvements in complex tasks like MPQA (85.2\% vs. 70.2\%, 21\% improvement) and Subj (82.8\% vs. 65.8\%, 26\% improvement). In the 20-shot setting, ERMAR achieves 80.54\% average accuracy compared to MemLong’s 73.4\%, demonstrating precise memory retrieval and robustness to context variations.
% \textbf{Trade-offs and Future Work}: While ERMAR excels in general language tasks, its slight underperformance on Proof-pile at 16k (3.18 vs. 2.99) suggests that its dynamic memory management may require further optimization for specialized mathematical content. This context-dependent variance highlights a trade-off between general language competence and mathematical performance, warranting further investigation into task-specific memory strategies.
\end{document}
\end{document}
\begin{thebibliography}{29}
\providecommand{\natexlab}[1]{#1}
\bibitem[{Abdin et~al.(2024)Abdin, Aneja, Awadalla, Awadallah, Awan, Bach, Bahree, Bakhtiari, Bao, Behl et~al.}]{abdin2024phi}
Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar~Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, et~al. 2024.
\newblock Phi-3 technical report: A highly capable language model locally on your phone.
\newblock \emph{arXiv preprint arXiv:2404.14219}.
\bibitem[{Azerbayev et~al.(2023)Azerbayev, Ayers, and Piotrowski}]{azerbayev2023proofpile}
Zhangir Azerbayev, Edward Ayers, and Bartosz Piotrowski. 2023.
\newblock Proofpile: A pre-training dataset of mathematical texts.
\bibitem[{Beltagy et~al.(2020)Beltagy, Peters, and Cohan}]{beltagy2020longformer}
Iz~Beltagy, Matthew~E Peters, and Arman Cohan. 2020.
\newblock Longformer: The long-document transformer.
\newblock \emph{arXiv preprint arXiv:2004.05150}.
\bibitem[{Cao et~al.(2007)Cao, Qin, Liu, Tsai, and Li}]{cao2007learning}
Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Hang Li. 2007.
\newblock Learning to rank: from pairwise approach to listwise approach.
\newblock In \emph{Proceedings of the 24th international conference on Machine learning}, pages 129--136.
\bibitem[{Chen et~al.(2023)Chen, Qian, Tang, Lai, Liu, Han, and Jia}]{chen2023longlora}
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. 2023.
\newblock Longlora: Efficient fine-tuning of long-context large language models.
\newblock \emph{arXiv preprint arXiv:2309.12307}.
\bibitem[{Fu et~al.(2024)Fu, Panda, Niu, Yue, Hajishirzi, Kim, and Peng}]{fu2024data}
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. 2024.
\newblock Data engineering for scaling language models to 128k context.
\newblock \emph{arXiv preprint arXiv:2402.10171}.
\bibitem[{Graves et~al.(2014)Graves, Wayne, and Danihelka}]{graves2014neural}
Alex Graves, Greg Wayne, and Ivo Danihelka. 2014.
\newblock Neural turing machines.
\newblock \emph{arXiv preprint arXiv:1410.5401}.
\bibitem[{Graves et~al.(2016)Graves, Wayne, Reynolds, Harley, Danihelka, Grabska-Barwi{\'n}ska, Colmenarejo, Grefenstette, Ramalho, Agapiou et~al.}]{graves2016hybrid}
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwi{\'n}ska, Sergio~G{\'o}mez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et~al. 2016.
\newblock Hybrid computing using a neural network with dynamic external memory.
\newblock \emph{Nature}, 538(7626):471--476.
\bibitem[{Guu et~al.(2020)Guu, Lee, Tung, Pasupat, and Chang}]{guu2020retrieval}
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. 2020.
\newblock Retrieval augmented language model pre-training.
\newblock In \emph{International conference on machine learning}, pages 3929--3938. PMLR.
\bibitem[{Hu et~al.(2021)Hu, Shen, Wallis, Allen-Zhu, Li, Wang, Wang, and Chen}]{hu2021lora}
Edward~J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu~Wang, and Weizhu Chen. 2021.
\newblock Lora: Low-rank adaptation of large language models.
\newblock \emph{arXiv preprint arXiv:2106.09685}.
\bibitem[{Izacard and Grave(2020)}]{izacard2020leveraging}
Gautier Izacard and Edouard Grave. 2020.
\newblock Leveraging passage retrieval with generative models for open domain question answering.
\newblock \emph{arXiv preprint arXiv:2007.01282}.
\bibitem[{Karpukhin et~al.(2020)Karpukhin, O{\u{g}}uz, Min, Lewis, Wu, Edunov, Chen, and Yih}]{karpukhin2020dense}
Vladimir Karpukhin, Barlas O{\u{g}}uz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020.
\newblock Dense passage retrieval for open-domain question answering.
\newblock \emph{arXiv preprint arXiv:2004.04906}.
\bibitem[{Khattab and Zaharia(2020)}]{khattab2020colbert}
Omar Khattab and Matei Zaharia. 2020.
\newblock Colbert: Efficient and effective passage search via contextualized late interaction over bert.
\newblock In \emph{Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval}, pages 39--48.
\bibitem[{Lewis et~al.(2020)Lewis, Perez, Piktus, Petroni, Karpukhin, Goyal, K{\"u}ttler, Lewis, Yih, Rockt{\"a}schel et~al.}]{lewis2020retrieval}
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich K{\"u}ttler, Mike Lewis, Wen-tau Yih, Tim Rockt{\"a}schel, et~al. 2020.
\newblock Retrieval-augmented generation for knowledge-intensive nlp tasks.
\newblock \emph{Advances in neural information processing systems}, 33:9459--9474.
\bibitem[{Liu et~al.(2024)Liu, Tang, Li, Chen, and Zhang}]{liu2024memlong}
Weijie Liu, Zecheng Tang, Juntao Li, Kehai Chen, and Min Zhang. 2024.
\newblock Memlong: Memory-augmented retrieval for long text modeling.
\newblock \emph{arXiv preprint arXiv:2408.16967}.
\bibitem[{Merity et~al.(2016)Merity, Xiong, Bradbury, and Socher}]{merity2016pointer}
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016.
\newblock Pointer sentinel mixture models.
\newblock \emph{arXiv preprint arXiv:1609.07843}.
\bibitem[{Nogueira and Cho(2019)}]{nogueira2019passage}
Rodrigo Nogueira and Kyunghyun Cho. 2019.
\newblock Passage re-ranking with bert.
\newblock \emph{arXiv preprint arXiv:1901.04085}.
\bibitem[{Peng et~al.(2023)Peng, Quesnelle, Fan, and Shippole}]{peng2023yarn}
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. 2023.
\newblock Yarn: Efficient context window extension of large language models.
\newblock \emph{arXiv preprint arXiv:2309.00071}.
\bibitem[{Press et~al.(2022)Press, Smith, and Lewis}]{press2022trainshorttestlong}
Ofir Press, Noah~A. Smith, and Mike Lewis. 2022.
\newblock \href {https://arxiv.org/abs/2108.12409} {Train short, test long: Attention with linear biases enables input length extrapolation}.
\newblock \emph{Preprint}, arXiv:2108.12409.
\bibitem[{Rae et~al.(2019)Rae, Potapenko, Jayakumar, and Lillicrap}]{rae2019compressive}
Jack~W Rae, Anna Potapenko, Siddhant~M Jayakumar, and Timothy~P Lillicrap. 2019.
\newblock Compressive transformers for long-range sequence modelling.
\newblock \emph{arXiv preprint arXiv:1911.05507}.
\bibitem[{Su et~al.(2024)Su, Ahmed, Lu, Pan, Bo, and Liu}]{su2024roformer}
Jianlin Su, Murtadha Ahmed, Yu~Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. 2024.
\newblock Roformer: Enhanced transformer with rotary position embedding.
\newblock \emph{Neurocomputing}, 568:127063.
\bibitem[{Touvron et~al.(2023{\natexlab{a}})Touvron, Lavril, Izacard, Martinet, Lachaux, Lacroix, Rozi{\`e}re, Goyal, Hambro, Azhar et~al.}]{touvron2023llama}
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth{\'e}e Lacroix, Baptiste Rozi{\`e}re, Naman Goyal, Eric Hambro, Faisal Azhar, et~al. 2023{\natexlab{a}}.
\newblock Llama: Open and efficient foundation language models.
\newblock \emph{arXiv preprint arXiv:2302.13971}.
\bibitem[{Touvron et~al.(2023{\natexlab{b}})Touvron, Martin, Stone, Albert, Almahairi, Babaei, Bashlykov, Batra, Bhargava, Bhosale et~al.}]{touvron2023llama2}
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et~al. 2023{\natexlab{b}}.
\newblock Llama 2: Open foundation and fine-tuned chat models.
\newblock \emph{arXiv preprint arXiv:2307.09288}.
\bibitem[{Tworkowski et~al.(2024)Tworkowski, Staniszewski, Pacek, Wu, Michalewski, and Mi{\l}o{\'s}}]{tworkowski2024focused}
Szymon Tworkowski, Konrad Staniszewski, Miko{\l}aj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Mi{\l}o{\'s}. 2024.
\newblock Focused transformer: Contrastive training for context scaling.
\newblock \emph{Advances in Neural Information Processing Systems}, 36.
\bibitem[{Vaswani(2017)}]{vaswani2017attention}
A~Vaswani. 2017.
\newblock Attention is all you need.
\newblock \emph{Advances in Neural Information Processing Systems}.
\bibitem[{Wu et~al.(2022)Wu, Rabe, Hutchins, and Szegedy}]{wu2022memorizing}
Yuhuai Wu, Markus~N Rabe, DeLesley Hutchins, and Christian Szegedy. 2022.
\newblock Memorizing transformers.
\newblock \emph{arXiv preprint arXiv:2203.08913}.
\bibitem[{Xiao et~al.(2022)Xiao, Liu, Shao, and Cao}]{xiao2022retromae}
Shitao Xiao, Zheng Liu, Yingxia Shao, and Zhao Cao. 2022.
\newblock Retromae: Pre-training retrieval-oriented language models via masked auto-encoder.
\newblock \emph{arXiv preprint arXiv:2205.12035}.
\bibitem[{Yen et~al.(2024)Yen, Gao, and Chen}]{yen2024long}
Howard Yen, Tianyu Gao, and Danqi Chen. 2024.
\newblock Long-context language modeling with parallel context encoding.
\newblock \emph{arXiv preprint arXiv:2402.16617}.
\bibitem[{Zaheer et~al.(2020)Zaheer, Guruganesh, Dubey, Ainslie, Alberti, Ontanon, Pham, Ravula, Wang, Yang et~al.}]{zaheer2020big}
Manzil Zaheer, Guru Guruganesh, Kumar~Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li~Yang, et~al. 2020.
\newblock Big bird: Transformers for longer sequences.
\newblock \emph{Advances in neural information processing systems}, 33:17283--17297.
\end{thebibliography}