PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zeping Yu, Yonatan Belinkov, Sophia Ananiadou

arXiv:2502.10835v1cs.CL

TL;DR

Latent multi-hop reasoning in LLMs remains difficult to explain, and explicit reasoning chains can be costly. The paper analyzes logit propagation with logit flow and proposes back attention; it reports improved accuracy across reasoning datasets and one-layer performance matching a two-layer transformer.

  • Problem

    Latent multi-hop reasoning mechanisms remain unclear, accuracy can be low even when individual hops are correct, and explicit chain-of-thought methods can require costly multiple inference steps or extensive sampling.

    Source: §1 · p. 1

  • Method

    The paper introduces logit flow to trace logits across layers and positions and proposes back attention, which lets lower layers use higher-level information.

    Source: §1 · p. 2

  • Results

    Back attention boosts accuracy across four LLMs on five reasoning datasets, and a 1-layer transformer with it matches a 2-layer transformer’s accuracy.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    The analysis identifies relation attribute extraction as a common source of two-hop failures, where conflicting logits can lower prediction accuracy.

    Source: §7 · p. 8

    The interpretability analysis focuses on single-hop and two-hop knowledge queries, and other reasoning tasks may involve mechanisms it does not capture.

    Source: §8 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames latent multi-hop reasoning as a costly and incompletely understood capability, then uses logit flow to analyze its failures and motivate back attention.

  • Latent multi-hop accuracy can remain low even when both hops are correct, while its mechanisms remain unclear and explicit chain-of-thought methods can be costly.Source: §1 · p. 1
  • Logit flow traces logits across layers and positions, identifying entity subject enrichment, entity attribute extraction, relation subject enrichment, and relation attribute extraction in single-hop prediction.Source: §1 · p. 2
  • In false two-hop cases, late relation attribute extraction captures high-layer r1 information, reinforcing e2 over e3 and reducing accuracy.Source: §1 · p. 2
  • Back attention lets a 1-layer transformer match a 2-layer transformer’s accuracy and boosts accuracy across four LLMs on five reasoning datasets.Source: §1 · p. 2
  • Back attention lets lower layers capture higher-level information and is effective both when training from scratch and adapting pretrained LLMs.Source: §1 · p. 2

2 Experimental Settings

The experiments use two-hop entity-relation triples, evaluate first-hop, second-hop, and combined prompts, and select cases according to query correctness and frequency.

  • Each TwoHop reasoning instance contains two linked triples, such as Mozart–mother–Maria Anna Mozart and Maria Anna Mozart–spouse–Leopold Mozart.Source: §2 · p. 2
  • The prompts query the first hop, second hop, or both; shortcut cases are removed, and retained instances require correct first-hop and two-hop predictions.Source: §2 · p. 2
  • After frequency filtering, Section 3 analyzes 889 all-correct cases, while Section 4 examines 568 human-entity cases containing correct and incorrect two-hop predictions.Source: §2 · p. 3

3 Mechanism of Single-Hop Prediction

The section develops logit flow as a neuron-level account of logits across layers and positions, then uses it to characterize single-hop prediction stages and examine their recurrence in second-hop cases.

3.1 Background

  • The analysis defines a neuron as the product of its coefficient score and subvalue, and describes FFN output as a weighted sum of neurons.Subvalues are columns of the second MLP, while subkeys are rows of the first MLP.Source: §3.1 · p. 3

3.2 Logit Flow: Tracing the Logits on Different Layers and Positions

  • Logit flow traces direct contributions from deep neurons storing final-prediction logits and indirect contributions from shallow neurons that amplify deeper-neuron coefficients.The analysis retains neuron contributions by position and layer to identify where important attention neurons are activated.Source: §3.2 · p. 4

3.3 Four Stages in Single-Hop Prediction

  • Attention neurons storing logits concentrate at entity, relation, and final positions, while FFN neurons with large inner products generally appear earlier.The entity-position stages align with prior layer-level conclusions about entity embeddings activating FFN features that attention layers process.Source: §3.3 · pp. 4–5
  • At entity positions, lower-layer neurons encode entity features; at relation and final positions, deeper layers combine those features with the relation to refine the answer.Subject enrichment and attribute extraction also occur at relation and final positions.Source: §3.3 · p. 5
  • Correct-versus-conflicting answer logits begin to diverge at relation and final positions after entity-attribute extraction; the entity position alone cannot distinguish the answers.The comparison averages logit differences across correct human-to-human cases.Source: §3.3 · p. 5
  • The same four stages appear in 889 second-hop cases, and activation patching also identifies important entity, relation, and final positions.Logit flow provides a more granular neuron-level perspective than the layer-level activation-patching analysis.Source: §3.3 · p. 5

4 Mechanism of Two-Hop Prediction

In 568 human-to-human-to-human queries, models often predicted an intermediate entity instead of the correct two-hop answer. Logit-flow analysis links this failure to stronger influence from the first-relation position, which can reinforce the intermediate answer.

4 Mechanism of Two-Hop Prediction

  • Among 568 cases evaluated in Llama2-7B, 52.3% correctly predicted e3, 42.4% predicted e2 (the answer to e1’s r1), and 5.3% predicted the answer to e1’s r2.e1, e2, and e3 are all human entities.Source: §4 · p. 5
  • False cases showed stronger influence from the first-relation position than correct cases, consistent with activation-patching results and analyses of two other Llama models.The authors report a similar trend in activation patching and in Llama3.1-8B and Llama3.2-3B.Source: §4 · pp. 5–6, Figure 3 · p. 5
  • At the first-relation position, high-layer states encode the first-hop entity and remain unchanged when the second relation is appended, reinforcing that intermediate entity’s probability.This reinforcement reduces the probability of selecting the correct two-hop answer.Source: §4 · p. 6
  • When first-relation features are extracted at layer 28, they activate parameters for the second hop at layers 28–31, but amplify the intermediate entity more than the final answer.This partially aligns with the “hopping too late” hypothesis, while showing that second-hop parameters still activate but contribute less than the direct influence of the intermediate entity.Source: §4 · p. 6

5 Back Attention: Letting Lower Layers Capture Higher-Layer features

Back attention lets lower layers use hidden states from higher layers and different positions. The paper reports improved accuracy across four LLMs and five reasoning datasets, and a one-layer transformer with back attention matches a two-layer transformer’s accuracy in a training-from-scratch experiment.

5 Back Attention: Letting Lower Layers Capture Higher-Layer features

  • Back attention computes queries from a lower source layer and keys and values from a higher target layer or stack of higher layers.In the illustrated one-layer setup, it runs after the original pass, adds its output to the layer-0 input, and recomputes the forward pass.Source: §5 · p. 6, Figure 4 · p. 6
  • The one-layer transformer with back attention reached 93.8% accuracy, versus 92.5% for a two-layer transformer and 83.8% for a one-layer transformer.The one-layer model with back attention used 56.7% of the two-layer model’s parameters and converged faster than the one-layer model.Source: §5 · p. 7
  • Adding back attention on the sixth layer achieved a peak test accuracy of 93.2%, followed by a steady decline at higher layers.The pretrained Llama-7B experiment froze the LLM parameters while fine-tuning back attention on each layer.Source: §5 · p. 7, Figure 5 · p. 7
  • Back attention significantly increased accuracy on all four tested LLMs across five datasets.The experiment added back attention on the sixth layer; the table reports accuracy before and after the addition.Source: Table 1 · p. 7, §5 · p. 7
  • For the Mozart example, sixth-layer back attention recovered the mother position’s layers 27–30 hidden states into the last position’s sixth layer.Darker scores indicate larger back-attention scores.Source: §5 · p. 7, Figure 6 · p. 7
  • Back attention can be added during fine-tuning, and its parameters account for 0.002% of Llama3’s 8 billion parameters.In Llama3.1-8B, applying it to a single layer increased average accuracy from 46.9% to 77.0%.Source: §5 · p. 8

6 Related Work

Related work covers explicit reasoning strategies and methods for interpreting model internals. The paper positions logit flow as an analysis of how neurons contribute to predictions, complementing prior neuron-level localization work.

6.1 Multi-Hop Reasoning in LLMs

  • Prior multi-hop reasoning work includes chain-of-thought, complexity-based prompting, self-consistency, self-play fine-tuning, increased inference sampling, and tree-based methods.These approaches seek to improve reasoning performance through explicit reasoning chains, sampling, or generated data.Source: §6.1 · p. 8

6.2 Mechanistic Interpretability

  • Mechanistic interpretability studies internal model mechanisms using tools such as logit lens, causal mediation, circuit construction, sparse autoencoders, and neuron-level analyses.Related work also examines the residual stream and represents FFN and attention outputs as weighted sums of neurons.Source: §6.2 · p. 8
  • Unlike prior neuron-level studies focused on localization, logit flow analyzes how neurons activate and contribute to final predictions.The authors present it as a way to examine influence beyond identifying important neurons.Source: §6.2 · p. 8

7 Conclusion

The paper analyzes latent multi-hop reasoning, finding that errors often emerge during relation attribute extraction. It introduces back attention, which improves reasoning accuracy across the reported settings.

7 Conclusion

  • Two-hop failures often arise during relation attribute extraction, where conflicting logits lower prediction accuracy.The analysis identifies this stage as a key factor affecting multi-hop reasoning accuracy.Source: §7 · p. 8
  • Back attention lets lower layers access higher-layer hidden states to restore important features.The paper presents this as a method for addressing the identified reasoning failures.Source: §7 · p. 8
  • A 1-layer transformer with back attention matches the accuracy of a 2-layer transformer, and the method improves accuracy across five datasets and four models.The reported improvements apply to pretrained LLMs and multi-hop reasoning.Source: §7 · p. 8
  • Logit flow reveals four stages in single-hop knowledge prediction: entity subject enrichment, entity attribute extraction, relation subject enrichment, and relation attribute extraction.The paper presents the analysis as providing insight into factors affecting reasoning accuracy.Source: §7 · p. 8

8 Limitations

The interpretability analysis centers on single-hop and two-hop knowledge queries, while back attention is evaluated in a single-layer configuration. The supplied findings also describe how logit differences vary by position and identify four stages in second-hop predictions.

8 Limitations

  • The interpretability analysis focuses on single-hop and two-hop knowledge queries, so other reasoning tasks may involve mechanisms it does not capture.The authors call for further validation on more diverse tasks and refinement of the interpretability framework.Source: §8 · p. 9
  • Back attention is applied to only one layer in this work; extending it to multiple layers is proposed as future research.The authors suggest that two or more layers could potentially yield greater improvements.Source: §8 · p. 9

A Logit Difference at Different Positions

  • At the entity position, the logit difference stays small; at relation and last positions, it rises sharply after layers 19–20.The deeper relation and last positions encode relation-specific knowledge that differentiates the correct prediction.Source: A Logit Difference at Different… · p. 12

B Results of Logit Flow on Second-Hop Queries in Llama2-7B

C Results of Activation Patching on Single-Hop Queries in Llama2-7B

Activation patching gives a less distinct view of position importance than logit flow because it combines the importance of FFN and attention modules. The supplied second-hop results report the same four prediction stages as first-hop queries.

C Results of Activation Patching on Single-Hop Queries in Llama2-7B

  • Second-hop queries show four stages similar to those found in first-hop queries.The supplied Figure 8 describes logit flow on “e2’s r2 is” -> “e3”.Source: Figure 8 · p. 12
  • Activation patching shows greater importance at entity and last positions, while relation positions appear less significant than in logit flow.The comparison is based on single-hop queries in Llama2-7B.Source: C Results of Activation Patching on… · p. 12
  • Activation patching combines FFN and attention importance, whereas logit flow separates them for a more granular, neuron-level view.The supplied activation-patching result also notes importance at high layers of the r1 position, though less clearly than in logit flow.Source: C Results of Activation Patching on… · p. 12, Figure 9 · p. 12

D Results of Logit Flow on Two-Hop Queries in Llama2-7B

In two-hop queries, relation positions are less important in logit flow than in single-hop queries. Activation patching of correct and false cases also shows a clearer influence from high layers at the r1 position in false cases.

D Results of Logit Flow on Two-Hop Queries in Llama2-7B

  • Relation positions are significantly less important in two-hop than in single-hop logit flow, while e1 features are mainly extracted into the last position.The authors suggest this may activate parameters associated with “e1’s r1”, “e1’s r2”, and “e1’s r1’s r2”.Source: D Results of Logit Flow on Two-Hop… · pp. 12–13, Figure 10 · p. 13

E Results of Activation Patching on Correct and False Two-Hop Queries in Llama2-7B

  • False two-hop cases show a clearer influence from high layers at the r1 position than correct cases, increasing the probability of e2 and reducing accuracy.The activation-patching trend is reported as similar to the logit-flow findings.Source: E Results of Activation Patching on… · p. 13

F Results of Logit Flow and Activation Patching on Correct and False Two-Hop Queries in Llama3.1-8B and Llama3.2-3B

Back attention brings a 1-layer transformer to performance similar to a 2-layer transformer, while logit flow and activation patching show greater high-layer impact at r1 in false than correct two-hop cases.

F Results of Logit Flow and Activation Patching on Correct and False Two-Hop Queries in Llama3.1-8B and Llama3.2-3B

G Loss and Accuracy of Back Attention on 1-Layer Transformer

Loading 2502.10835v1…