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

TL;DR

This paper introduces logit flow for interpretability and proposes back attention to improve multi-hop reasoning in large language models, significantly enhancing their reasoning accuracy.

Contribution

We develop logit flow for analyzing reasoning processes and propose back attention, a novel mechanism that boosts multi-hop reasoning performance in LLMs.

Findings

Logit flow reveals four stages in single-hop reasoning.

Back attention improves reasoning accuracy across multiple datasets.

A 1-layer transformer with back attention matches 2-layer performance.

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.