# Enhancing Long Context Performance in LLMs Through Inner Loop Query Mechanism

arXiv: https://arxiv.org/abs/2410.12859 · HTML: https://mlsanity.com/paper/2410.12859/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 26 of its 35 references that are in the map.
- First posted on arXiv 11 Oct 2024.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools](https://mlsanity.com/paper/2406.12793/) (2024)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval](https://mlsanity.com/paper/2401.18059/) (2024)
- [Beyond Goldfish Memory: Long-Term Open-Domain Conversation](https://mlsanity.com/paper/2107.07567/) (2021)
- [BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack](https://mlsanity.com/paper/2406.10149/) (2024)
- [Raise a Child in Large Language Model: Towards Effective and Generalizable Fine-tuning](https://mlsanity.com/paper/2109.05687/) (2021)
- [Long Time No See! Open-Domain Conversation with Long-Term Persona Memory](https://mlsanity.com/paper/2203.05797/) (2022)
- [Memory Augmented Large Language Models are Computationally Universal](https://mlsanity.com/paper/2301.04589/) (2023)
- [LLM-based Medical Assistant Personalization with Short- and Long-Term Memory Coordination](https://mlsanity.com/paper/2309.11696/) (2023)
- [Joint Passage Ranking for Diverse Multi-Answer Retrieval](https://mlsanity.com/paper/2104.08445/) (2021)
- [Dense Hierarchical Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2110.15439/) (2021)
- [A Theory of Consciousness from a Theoretical Computer Science Perspective: Insights from the Conscious Turing Machine](https://mlsanity.com/paper/2107.13704/) (2021)
- [Towards Tracing Factual Knowledge in Language Models Back to the Training Data](https://mlsanity.com/paper/2205.11482/) (2022)
- [Memory and Knowledge Augmented Language Models for Inferring Salience in Long-Form Stories](https://mlsanity.com/paper/2109.03754/) (2021)

## Built on by (most cited first, 4 of 4)

- [Memory in the Age of AI Agents](https://mlsanity.com/paper/2512.13564/) (2025)
- [Beyond Pipelines: A Survey of the Paradigm Shift toward Model-Native Agentic AI](https://mlsanity.com/paper/2510.16720/) (2025)
- [EvoMem: Improving Multi-Agent Planning with Dual-Evolving Memory](https://mlsanity.com/paper/2511.01912/) (2025)
- [Computability of Agentic Systems](https://mlsanity.com/paper/2602.13222/) (2026)

## Abstract (arXiv)

Transformers have a quadratic scaling of computational complexity with input size, which limits the input context window size of large language models (LLMs) in both training and inference. Meanwhile, retrieval-augmented generation (RAG) besed models can better handle longer contexts by using a retrieval system to filter out unnecessary information. However, most RAG methods only perform retrieval based on the initial query, which may not work well with complex questions that require deeper reasoning. We introduce a novel approach, Inner Loop Memory Augmented Tree Retrieval (ILM-TR), involving inner-loop queries, based not only on the query question itself but also on intermediate findings. At inference time, our model retrieves information from the RAG system, integrating data from lengthy documents at various levels of abstraction. Based on the information retrieved, the LLM generates texts stored in an area named Short-Term Memory (STM) which is then used to formulate the next query. This retrieval process is repeated until the text in STM converged. Our experiments demonstrate that retrieval with STM offers improvements over traditional retrieval-augmented LLMs, particularly in long context tests such as Multi-Needle In A Haystack (M-NIAH) and BABILong.