Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism

• URL: http://arxiv.org/abs/2504.18574v2
• Date: Wed, 11 Jun 2025 11:06:08 GMT
• Title: Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism
• Authors: Aviv Bick, Eric Xing, Albert Gu,
• Abstract summary: State-space models (SSMs) offer efficient alternatives to Transformers for long sequences.<n>In this work, we examine how in-context retrieval operates in Transformer- and SSM-based language models.
• Score: 15.626801223435173
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: State-space models (SSMs) offer efficient alternatives to Transformers for long sequences, but their fixed-size recurrent state limits capability on algorithmic tasks, such as retrieving past context. In this work, we examine how in-context retrieval operates in Transformer- and SSM-based language models and find that both rely on a similar Gather-and-Aggregate (G&A) mechanism: a Gather Head extracts relevant information pieces from context, which an Aggregate Head integrates into a single representation. In both architectures, G&A concentrates in a few heads, forming critical bottlenecks even for simple retrieval. For example, we show that disabling a single Gather or Aggregate Head in a pruned Llama-3.1-8B impairs retrieving the correct answer letter in MMLU, reducing its accuracy from 66% to 25% (random guessing). Moreover, this retrieval bottleneck can obscure limited knowledge demands of tasks as the pruned model succeeds on MMLU with functioning G&A heads yet fails on other knowledge benchmarks. The bottleneck similarly extends to tasks where SSMs typically underperform, such as GSM8K, BBH, and dialogue comprehension. We show that SSMs' retrieval challenges manifest in these heads, creating smoother attention patterns instead of the sharp token transitions effective G&A requires. Thus, the Transformer-SSM retrieval gap exists in just a few heads, rather than the entire language model. This suggests a unified explanation for Transformer vs. SSM performance gap while showing how to merge their strengths. We find that pretrained hybrid models, where SSMs are combined with a few attention layers, delegate the role of Aggregate Heads to attention. Similarly, replacing a single G&A head in a pretrained SSM with an attention variant boosts retrieval and benchmark scores.

Related papers

• ImpRAG: Retrieval-Augmented Generation with Implicit Queries [49.510101132093396]
  ImpRAG is a query-free RAG system that integrates retrieval and generation into a unified model.<n>We show that ImpRAG achieves 3.6-11.5 improvements in exact match scores on unseen tasks with diverse formats.
  arXiv  Detail & Related papers  (2025-06-02T21:38:21Z)
• UniversalRAG: Retrieval-Augmented Generation over Multiple Corpora with Diverse Modalities and Granularities [53.76854299076118]
  UniversalRAG is a novel RAG framework designed to retrieve and integrate knowledge from heterogeneous sources with diverse modalities and granularities. We propose a modality-aware routing mechanism that dynamically identifies the most appropriate modality-specific corpus and performs targeted retrieval within it. We validate UniversalRAG on 8 benchmarks spanning multiple modalities, showing its superiority over modality-specific and unified baselines.
  arXiv  Detail & Related papers  (2025-04-29T13:18:58Z)
• Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B [46.99314622487279]
  In-Context Learning (ICL) is an intriguing ability of large language models (LLMs)<n>We use causal interventions to identify information flow in Gemma-2 2B for five naturalistic ICL tasks.<n>We find that the model infers task information using a two-step strategy we call contextualize-then-aggregate.
  arXiv  Detail & Related papers  (2025-03-31T18:33:55Z)
• SAGE: A Framework of Precise Retrieval for RAG [9.889395372896153]
  Retrieval-augmented generation (RAG) has demonstrated significant proficiency in conducting question-answering tasks.<n>RAG methods segment the corpus without considering semantics, making it difficult to find relevant context.<n>We introduce a RAG framework (SAGE) to overcome these limitations.
  arXiv  Detail & Related papers  (2025-03-03T16:25:58Z)
• Harnessing Large Language Models for Knowledge Graph Question Answering via Adaptive Multi-Aspect Retrieval-Augmentation [81.18701211912779]
  We introduce an Adaptive Multi-Aspect Retrieval-augmented over KGs (Amar) framework.<n>This method retrieves knowledge including entities, relations, and subgraphs, and converts each piece of retrieved text into prompt embeddings.<n>Our method has achieved state-of-the-art performance on two common datasets.
  arXiv  Detail & Related papers  (2024-12-24T16:38:04Z)
• Inference Scaling for Bridging Retrieval and Augmented Generation [47.091086803980765]
  Retrieval-augmented generation (RAG) has emerged as a popular approach to steering the output of a large language model (LLM)<n>We show such bias can be mitigated, from inference scaling, aggregating inference calls from the permuted order of retrieved contexts.<n>We showcase the effectiveness of MOI on diverse RAG tasks, improving ROUGE-L on MS MARCO and EM on HotpotQA benchmarks by 7 points.
  arXiv  Detail & Related papers  (2024-12-14T05:06:43Z)
• Invar-RAG: Invariant LLM-aligned Retrieval for Better Generation [43.630437906898635]
  We propose a novel two-stage fine-tuning architecture called Invar-RAG. In the retrieval stage, an LLM-based retriever is constructed by integrating LoRA-based representation learning. In the generation stage, a refined fine-tuning method is employed to improve LLM accuracy in generating answers based on retrieved information.
  arXiv  Detail & Related papers  (2024-11-11T14:25:37Z)
• Birdie: Advancing State Space Models with Reward-Driven Objectives and Curricula [23.071384759427072]
  State space models (SSMs) offer advantages over Transformers but struggle with tasks requiring long-range in-context retrieval-like text copying, associative recall, and question answering over long contexts. We propose a novel training procedure, Birdie, that significantly enhances the in-context retrieval capabilities of SSMs without altering their architecture.
  arXiv  Detail & Related papers  (2024-11-01T21:01:13Z)
• SG-FSM: A Self-Guiding Zero-Shot Prompting Paradigm for Multi-Hop Question Answering Based on Finite State Machine [27.274219226254026]
  Multi-hop Question Answering (MHQA) remains challenging for many existing models. We propose the Self-Guiding prompting Finite State Machine (SG-FSM) to strengthen multi-hop reasoning abilities.
  arXiv  Detail & Related papers  (2024-10-22T13:47:38Z)
• RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation [42.82192656794179]
  Large Language Models (LLMs) exhibit remarkable capabilities but are prone to generating inaccurate or hallucinatory responses. This limitation stems from their reliance on vast pretraining datasets, making them susceptible to errors in unseen scenarios. Retrieval-Augmented Generation (RAG) addresses this by incorporating external, relevant documents into the response generation process.
  arXiv  Detail & Related papers  (2024-03-31T08:58:54Z)
• Repeat After Me: Transformers are Better than State Space Models at Copying [53.47717661441142]
  We show that while generalized state space models are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context.
  arXiv  Detail & Related papers  (2024-02-01T21:44:11Z)
• GSVA: Generalized Segmentation via Multimodal Large Language Models [72.57095903188922]
  Generalized Referring Expression (GRES) extends the scope of classic RES to refer to multiple objects in one expression or identify the empty targets absent in the image. Current solutions to GRES remain unsatisfactory since segmentation MLLMs cannot correctly handle the cases where users might reference multiple subjects in a singular prompt. We propose Generalized Vision Assistant (GSVA) to address this gap.
  arXiv  Detail & Related papers  (2023-12-15T02:54:31Z)
• Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection [74.51523859064802]
  We introduce a new framework called Self-Reflective Retrieval-Augmented Generation (Self-RAG) Self-RAG enhances an LM's quality and factuality through retrieval and self-reflection. It significantly outperforms state-of-the-art LLMs and retrieval-augmented models on a diverse set of tasks.
  arXiv  Detail & Related papers  (2023-10-17T18:18:32Z)
• USER: Unified Semantic Enhancement with Momentum Contrast for Image-Text Retrieval [115.28586222748478]
  Image-Text Retrieval (ITR) aims at searching for the target instances that are semantically relevant to the given query from the other modality. Existing approaches typically suffer from two major limitations.
  arXiv  Detail & Related papers  (2023-01-17T12:42:58Z)
• UniKGQA: Unified Retrieval and Reasoning for Solving Multi-hop Question Answering Over Knowledge Graph [89.98762327725112]
  Multi-hop Question Answering over Knowledge Graph(KGQA) aims to find the answer entities that are multiple hops away from the topic entities mentioned in a natural language question. We propose UniKGQA, a novel approach for multi-hop KGQA task, by unifying retrieval and reasoning in both model architecture and parameter learning.
  arXiv  Detail & Related papers  (2022-12-02T04:08:09Z)
• Rethinking the Two-Stage Framework for Grounded Situation Recognition [61.93345308377144]
  Grounded Situation Recognition is an essential step towards "human-like" event understanding. Existing GSR methods resort to a two-stage framework: predicting the verb in the first stage and detecting the semantic roles in the second stage. We propose a novel SituFormer for GSR which consists of a Coarse-to-Fine Verb Model (CFVM) and a Transformer-based Noun Model (TNM)
  arXiv  Detail & Related papers  (2021-12-10T08:10:56Z)
• Why Do Pretrained Language Models Help in Downstream Tasks? An Analysis of Head and Prompt Tuning [66.44344616836158]
  We propose an analysis framework that links the pretraining and downstream tasks with an underlying latent variable generative model of text. We show that 1) under certain non-degeneracy conditions on the HMM, simple classification heads can solve the downstream task, 2) prompt tuning obtains downstream guarantees with weaker non-degeneracy conditions, and 3) our recovery guarantees for the memory-augmented HMM are stronger than for the vanilla HMM.
  arXiv  Detail & Related papers  (2021-06-17T03:31:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.