Back Attention: Understanding and Enhancing Multi-Hop Reasoning in Large Language Models

2025-02-15

AI Summary

This work addresses the performance degradation of large language models (LLMs) on implicit multi-hop reasoning (e.g., “Who was Mozart’s mother’s spouse?”), identifying critical failure modes in the relational attribute extraction stage. We propose Logit Flow, a novel interpretability method that quantifies layer- and position-wise contributions to final predictions via attribution paths. Furthermore, we introduce Back Attention—a first-of-its-kind mechanism enabling lower layers to dynamically attend to higher-layer hidden states across both positions and layers—thereby breaking the unidirectional hierarchical dependency of standard Transformers. We also establish the first four-stage analytical framework for single-hop reasoning, precisely localizing bottlenecks in multi-hop inference. Evaluated on five benchmark datasets, our approach achieves: (1) one-layer modified models matching the performance of original two-layer baselines; and (2) significant average accuracy gains across four mainstream LLMs, demonstrating its generalizability and effectiveness in enhancing implicit multi-hop reasoning capability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

We investigate how large language models perform latent multi-hop reasoning in prompts like"Wolfgang Amadeus Mozart's mother's spouse is". To analyze this process, we introduce logit flow, an interpretability method that traces how logits propagate across layers and positions toward the final prediction. Using logit flow, we identify four distinct stages in single-hop knowledge prediction: (A) entity subject enrichment, (B) entity attribute extraction, (C) relation subject enrichment, and (D) relation attribute extraction. Extending this analysis to multi-hop reasoning, we find that failures often stem from the relation attribute extraction stage, where conflicting logits reduce prediction accuracy. To address this, we propose back attention, a novel mechanism that enables lower layers to leverage higher-layer hidden states from different positions during attention computation. With back attention, a 1-layer transformer achieves the performance of a 2-layer transformer. Applied to four LLMs, back attention improves accuracy on five reasoning datasets, demonstrating its effectiveness in enhancing latent multi-hop reasoning ability.

Problem

• Enhance multi-hop reasoning in LLMs
• Identify stages in knowledge prediction
• Propose back attention mechanism

Innovation

• Introduces logit flow analysis
• Proposes back attention mechanism
• Enhances multi-hop reasoning accuracy