Back Attention: Understanding and Enhancing Multi-Hop Reasoning in Large Language Models
Paper Review
Abstract
이 논문은 LLM이 “Mozart의 어머니의 배우자” 같은 질의를 처리할 때, 겉으로는 중간 추론 과정을 생성하지 않아도 내부적으로 여러 단계를 거쳐 답을 낸다는 점에 주목한다. 저자들은 이런 잠재적 멀티홉 추론(latent multi-hop reasoning) 이 실제로 어떤 방식으로 이뤄지는지 분석하기 위해 logit flow라는 해석 방법을 제안한다.
논문은 먼저 단일 홉 지식 예측이 네 단계로 이루어진다고 설명한다. 그리고 2-hop 추론이 실패하는 주된 원인이 relation attribute extraction 단계에서 발생한다고 분석한다. 이를 개선하기 위해 상위 레이어의 중요한 hidden state를 하위 레이어가 다시 활용할 수 있도록 하는 back attention을 설계했고, 작은 트랜스포머와 여러 LLM에서 성능 향상을 확인했다고 보고한다.
⸻
Introduction
서론의 문제의식은 분명하다. 최근 멀티홉 추론을 향상시키는 대표적인 방법은 chain-of-thought 계열이지만, 이 방식은 추론 과정을 길게 생성해야 하므로 비용이 크다. 따라서 중요한 것은 중간 reasoning text를 길게 생성하지 않더라도, 모델 내부에서 여러 단계를 처리할 수 있는 latent multi-hop reasoning 능력이다.
저자들은 기존 연구가 주로 “The r2 of the r1 of e1 is” 같은 형식에 집중했다고 지적한다. 반면 이 논문은 “e1’s r1’s r2 is” 같은 형식을 중심으로 분석하면서, 어느 위치의 어떤 레이어가 어떤 정보를 담고 있고, 그 정보가 최종 답으로 어떻게 이어지는지를 추적한다. 또한 단순히 성능을 올리는 것이 아니라, 내부 메커니즘을 먼저 파악한 뒤 그 병목을 겨냥해 구조를 바꾸는 방향을 취한다.
이 논문의 기여는 크게 세 가지다.
- logit flow라는 해석 기법을 제안한다.
- 단일 홉과 이중 홉 추론의 내부 메커니즘을 분석해 실패 원인을 정리한다.
- 그 분석을 바탕으로 back attention을 제안하고 실제 성능 개선을 보여준다.
⸻
Experimental Settings
논문은 분석용 데이터로 TwoHop 데이터셋을 사용한다. 각 샘플은 두 개의 삼중항 <e1, r1, e2>와 <e2, r2, e3>로 구성된다. 예를 들면 <Mozart, mother, Maria Anna Mozart>와 <Maria Anna Mozart, spouse, Leopold Mozart> 같은 식이다. 이를 바탕으로 1-hop 질의 “e1’s r1 is”, 두 번째 홉 질의 “e2’s r2 is”, 최종 2-hop 질의 “e1’s r1’s r2 is”를 만든다.
데이터셋은 단순 지름길로 풀리는 사례를 줄이기 위해 필터링되었다. 특정 타입이 지나치게 많아지는 문제도 줄였고, 결과는 타입 수준 가중 평균으로 집계했다. 논문은 가장 큰 단일 카테고리가 전체의 12.7%에 불과하다고 설명한다.
분석도 목적에 따라 나뉜다. 단일 홉 메커니즘 분석에는 첫 번째 홉, 두 번째 홉, 전체 2-hop 질의가 모두 맞는 사례 889개를 사용한다. 반면 2-hop 실패 분석에는 e1, e2, e3가 모두 human 타입인 568개 사례를 따로 모아 정답과 오답의 내부 패턴을 비교한다.
이 설정은 메커니즘을 해석하기에는 깔끔하다. 다만 분석 대상이 비교적 통제된 지식 질의 형식에 집중되어 있다는 점은 함께 염두에 둘 필요가 있다.
⸻
Mechanism of Single-Hop Prediction
Background
이 부분에서는 decoder-only transformer가 residual stream, attention, FFN을 통해 어떻게 다음 토큰 분포를 만드는지 설명한다. 특히 FFN 출력은 개별 뉴런의 합으로, attention 출력은 여러 헤드와 위치의 가중합으로 볼 수 있다는 점을 이용해 나중에 어떤 뉴런과 어떤 위치가 정답 로짓 형성에 기여했는지를 추적한다.
즉 이 논문의 해석은 단순히 “이 레이어가 중요하다” 수준이 아니라, 어떤 위치의 어떤 모듈이 실제 정답 토큰 로짓을 올리는 데 기여했는가를 구체적으로 파악하는 데 초점이 있다.
Logit Flow
logit flow는 최종 예측에 기여하는 FFN 뉴런과 attention 뉴런을 찾고, 그것이 레이어와 토큰 위치별로 어떻게 분포하는지를 보는 해석 방법이다. 중요한 attention 뉴런을 활성화시키는 FFN 뉴런까지 함께 추적함으로써, 직접 정답 로짓을 저장하는 경로와 그것을 간접적으로 활성화하는 경로를 구분하려고 한다. 저자들은 이 방식이 activation patching보다 더 세밀한 뉴런 수준 분석을 가능하게 한다고 주장한다.
이 방법의 장점은 “어딘가에 정보가 존재한다”와 “그 정보가 실제로 최종 답을 만드는 데 사용된다”를 어느 정도 분리해서 볼 수 있다는 점이다. 해석 연구에서 자주 생기는 모호함을 줄이려는 시도라고 볼 수 있다.
Four Stages in Single-Hop Prediction
논문은 Llama2-7B와 889개의 1-hop 쿼리를 분석한 결과, 단일 홉 지식 예측이 다음 네 단계로 이루어진다고 설명한다.
- Entity Subject Enrichment
엔티티 위치에서 하위~중간 레이어 FFN 뉴런이 엔티티 자체의 특징을 풍부하게 만든다. 쉽게 말하면 “Mozart”라는 이름이 단순 문자열이 아니라, 그 인물과 관련된 속성들을 담는 방향으로 표현이 강화되는 단계다.
- Entity Attribute Extraction
attention 뉴런이 앞 단계에서 강화된 엔티티 표현을 읽어낸다. 즉 엔티티에 담긴 속성 정보가 다른 위치로 전달되기 시작한다.
- Relation Subject Enrichment
relation 위치와 마지막 위치에서 FFN 뉴런이 관계와 결합된 주체 표현을 강화한다. 예를 들어 “Mozart”와 “mother”가 결합되면서 이제 단순한 인물 정보가 아니라 “Mozart의 어머니”라는 방향으로 표현이 정리된다.
- Relation Attribute Extraction
attention과 FFN이 함께 작동하면서 관계가 반영된 최종 정답 후보를 끌어올린다. 이 단계에서 비로소 “Maria Anna Mozart” 같은 실제 답 토큰이 강하게 예측된다.
이 네 단계는 기존에 막연하게 이야기되던 “엔티티 정보가 저장되고 attention이 retrieval한다”는 설명을 더 세밀하게 나눈 것이다. 특히 엔티티 특징 형성과 관계 결합 과정을 구분해 보여준 점이 의미 있다.
⸻
Mechanism of Two-Hop Prediction
이제 논문은 더 어려운 질문으로 넘어간다. 왜 모델은 1-hop 두 개를 각각 풀 수 있으면서도, 2-hop 전체 질의에서는 자주 틀릴까. 기존 연구는 이를 “첫 번째 홉을 너무 늦게 해결한다”는 hopping too late 가설로 설명했다. 이 논문은 그 가설을 일부 지지하면서도, 실제 실패 메커니즘은 더 복잡하다고 본다.
분석 결과, 2-hop 질의에서는 single-hop과 달리 relation 위치의 중요도가 줄어들고, 모델이 최종 정답 e3 대신 중간 엔티티 e2에 끌리는 경우가 많다. 논문은 human→human→human 568개 사례에서 “e1’s r1’s r2 is” 형식을 실험했을 때, 정답 e3를 맞힌 비율이 52.3%, e2를 그대로 예측한 비율이 42.4%, “e1’s r2”에 해당하는 다른 충돌 답을 낸 경우가 5.3%라고 보고한다.
핵심은 실패가 relation attribute extraction 단계에서 집중된다는 점이다. false case에서는 r1 위치의 고층 정보가 이 단계에서 강하게 반영되는데, 이 정보는 이미 e2와 강하게 연결되어 있다. 그 결과 “e2 features + r2 → e3” 경로도 일부 활성화되지만, 정작 더 강하게 작용하는 것은 e2 자체의 로짓을 밀어 올리는 효과다.
즉 모델은 두 번째 홉을 전혀 못 하는 것이 아니다. 오히려 두 번째 홉을 시도하긴 하지만, 그 과정에서 첫 번째 홉의 중간 결과가 너무 강하게 남아 있어서 최종 정답보다 중간 엔티티가 더 크게 증폭된다. 이 설명은 기존의 단순한 “너무 늦게 hopping한다”는 설명보다 훨씬 구체적이다.
⸻
Back Attention: Letting Lower Layers Capture Higher-Layer Features
이 분석을 바탕으로 저자들은 back attention을 제안한다. 아이디어는 비교적 직관적이다. 멀티홉 추론이 실패하는 이유가 하위 레이어에서 필요한 고수준 정보를 충분히 활용하지 못하는 데 있다면, 상위 레이어의 중요한 hidden state를 다시 하위 레이어가 참조할 수 있게 만들면 도움이 될 수 있다는 것이다.
구조적으로는 먼저 일반적인 forward pass를 수행해 각 레이어와 위치의 hidden state를 얻는다. 그다음 특정 하위 레이어에서 query는 하위 레이어 입력으로부터 만들고, key와 value는 더 높은 레이어들의 hidden state로부터 만들어 attention을 수행한다. 이렇게 얻은 back attention 출력 B를 하위 레이어 입력에 더한 뒤, 이후 계산을 다시 진행한다.
이 말은 곧 하위 레이어가 원래는 접근할 수 없었던 상위 레이어의 표현을 되가져와 사용할 수 있다는 뜻이다. 논문에서는 이를 통해 멀티홉 추론에 필요한 고수준 특징을 더 이른 시점에서 활용할 수 있다고 본다.
Training from Scratch
저자들은 먼저 2-digit addition 데이터셋에서 back attention의 원리를 검증한다. 결과는 다음과 같다. 1-layer transformer는 83.8%, 1-layer transformer + back attention은 93.8%, 2-layer transformer는 92.5%를 기록했다. 즉 1-layer 모델에 back attention을 추가한 결과가 2-layer 모델과 비슷하거나 더 나은 수준을 보였다. 게다가 파라미터는 2-layer 모델의 56.7% 수준이라고 설명한다.
이 실험은 back attention이 단지 pretrained LLM용 보조 트릭이 아니라, 구조적으로 얕은 모델이 더 깊은 표현을 활용하게 만드는 장치라는 점을 보여준다.
Adding Back Attention in Pre-trained LLMs
이후 논문은 Llama, Mistral, Qwen 계열의 pretrained LLM에 back attention을 적용한다. 먼저 Llama-7B에서 산술 데이터셋으로 어느 레이어에 추가할지 탐색했는데, 기본 성능 67.1%가 6번째 레이어에 넣었을 때 93.2%까지 상승했다. 이후 더 높은 레이어에 넣으면 성능이 다시 낮아졌다고 보고한다.
본 실험에서는 layer 6을 사용해 1DC, SVAMP, MultiArith, TwoHop, StrategyQA 다섯 개 데이터셋에서 평가한다. 결과는 대부분 매우 큰 향상을 보여준다. 예를 들어 Llama3-8B는 1DC에서 72.7→97.0, MultiArith에서 21.1→88.9, TwoHop에서 11.5→47.8, StrategyQA에서 65.1→86.2로 상승했다. Qwen2.5-14B 역시 TwoHop에서 3.9→34.6, StrategyQA에서 71.2→91.4로 올랐다.
정성 분석도 제시된다. Mozart 예시에서 back attention score를 보면 “mother” 위치의 상층 표현이 마지막 위치의 하위 레이어로 복원되는 패턴이 관찰된다고 한다. 이는 back attention이 실제로 상위 레이어 특징을 하위 레이어로 끌어내리는 방식으로 작동한다는 시각적 증거로 해석된다.
또한 논문은 계산 효율 측면에서 chain-of-thought와의 비교도 제시한다. CoT는 생성 토큰 수가 길어질수록 비용이 커지지만, back attention은 토큰 수를 크게 늘리지 않고 토큰당 계산량만 어느 정도 증가시키는 방식이라 전체 비용 측면에서 유리할 수 있다고 주장한다. 다만 실제 대규모 서빙 환경에서의 상세 벤치마크까지는 제시하지 않는다.
⸻
Related Work
관련 연구는 크게 두 갈래로 정리된다. 첫째는 멀티홉 추론 성능을 높이는 chain-of-thought, self-consistency, tree/graph형 reasoning 계열이다. 이들은 외부로 추론 과정을 길게 생성하게 만드는 방향에 가깝다. 반면 이 논문은 추론 과정을 길게 출력하지 않고도 내부적으로 멀티홉 reasoning을 강화할 수 있는가에 더 관심이 있다.
둘째는 transformer 해석 연구다. FFN이 사실 지식을 저장한다는 연구, activation patching, residual stream 해석 등이 이 논문의 배경이 된다. 저자들은 자신들의 logit flow가 attention과 FFN의 역할을 더 세밀하게 분리해볼 수 있다는 점에서 기존 해석 기법과 차별화된다고 본다.
결국 이 논문은 단순한 프롬프팅 연구도 아니고, 순수한 해석 연구도 아니다. 해석을 통해 reasoning 병목을 찾고, 그 분석을 바탕으로 구조적 개선을 제안한 연구라고 보는 것이 가장 정확하다.
⸻
Conclusion
논문이 최종적으로 말하는 바는 분명하다. 단일 홉 지식 예측은 네 단계로 설명될 수 있고, 2-hop 추론 실패는 relation attribute extraction 단계에서 중간 엔티티가 최종 정답보다 더 강하게 증폭되는 현상과 밀접하다. 그리고 이를 해결하기 위해 제안한 back attention은 상위 레이어 정보를 하위 레이어로 다시 전달함으로써 멀티홉 추론 성능을 크게 향상시킨다.
즉 이 논문은 “왜 실패하는가”를 먼저 분석하고, 그 실패 원인에 정확히 맞는 구조를 설계해 실제 개선까지 연결한다는 점에서 완성도가 높다.
⸻
Limitations
저자들도 한계를 분명히 언급한다. 이 논문의 해석은 주로 single-hop과 two-hop knowledge query에 집중되어 있다. 따라서 다른 유형의 reasoning, 예를 들어 장기 계획, 복잡한 수학 증명, 코드 추론, 비정형 상식 추론까지 같은 방식으로 설명할 수 있는지는 아직 확실하지 않다.
그럼에도 여러 reasoning 데이터셋과 여러 LLM에서 back attention의 향상이 관찰되었다는 점은 어느 정도 일반성을 시사한다. 다만 이 메커니즘이 모든 reasoning 문제의 보편 원리라고 단정하기에는 아직 이르다.
