LLM의 Attention Map을 활용한 문장 Reranking 실험

동기

RAG(Retrieval-Augmented Generation)에서 검색된 문서를 LLM에 넣기 전에 관련성 높은 문장만 골라내는 것은 성능과 비용 모두에 중요한 문제입니다.

일반적으로는 별도의 Reranker 모델(Cross-Encoder 등)을 사용하지만, 한 가지 흥미로운 관점이 있습니다:

LLM이 텍스트를 처리할 때 이미 각 토큰에 대한 “관심도(Attention)“를 계산하고 있다면, 이 신호를 직접 활용할 수 있지 않을까?

이 아이디어를 실험으로 검증해보았습니다.

핵심 아이디어

Transformer 기반 LLM은 입력 토큰들 간의 Attention Score를 계산합니다. 이때 마지막으로 생성할 토큰이 어떤 입력 토큰에 주목하는지를 관찰하면, 해당 쿼리에 대해 어떤 문장이 더 관련성이 높은지를 판단할 수 있습니다.

전체 흐름

flowchart LR
    C[Context] --> LLM
    Q[Query] --> LLM
    A[Answer Prefix] --> LLM
    LLM -->|Attention| R[Reranking]

Attention 관찰 방식

일반적인 LLM 추론에서는 마지막 토큰의 Attention을 기반으로 다음 토큰을 생성합니다. 이 실험에서는 다음 토큰을 생성하는 대신, 그 Attention 분포 자체를 문장 관련성 점수로 활용합니다.

flowchart TB
    subgraph Context
        S1[Sentence 1]
        S2[Sentence 2]
        S3[Sentence 3]
    end
    Anchor[Anchor Token] -->|high| S1
    Anchor -->|low| S2
    Anchor -->|mid| S3
    S1 --> R1[Rank 1]
    S3 --> R2[Rank 2]
    S2 --> R3[Rank 3]

예를 들어, Context에 “영희는 서울대 학생입니다”, “철수는 25살입니다”, “영희는 철수 옆집에 살고 있습니다” 가 있고, Query가 “어디로 갈까요?“일 때, Answer Prefix인 “영희:"(앵커 토큰)의 Attention이 “서울대 학생” 문장에 가장 높게 나타납니다.

핵심은 answer_hint_prefix(위 예시에서 “영희:")가 앵커 토큰 역할을 한다는 점입니다. 이 토큰의 Attention이 Context 중 어디에 집중되는지를 측정하여, 현재 상황에 가장 관련 있는 문장을 찾아냅니다.

관련 논문

AttentionRAG (arXiv:2503.10720)

이 실험의 기반이 된 논문입니다. 핵심 기여:

• 쿼리를 Next-Token Prediction 형태로 변환: “Where is Daniel?” → “Daniel is in the ____”
• 앵커 토큰: 빈칸 위치의 토큰이 의미적 초점을 단일 토큰에 집중시킴
• 전체 레이어 집계: 얕은 층(구문 정보) + 깊은 층(의미 정보)을 모두 합산
• 결과: LLMLingua 대비 약 10% 성능 향상, 최대 6.3배 컨텍스트 압축

In-Context Re-ranking (ICLR 2025)

• LLM의 생성(generation) 없이 Attention 패턴 변화만으로 문서를 재순위화
• 보정(Calibration): 의미 없는 쿼리(“N/A”)로 기준선을 측정하여 위치 편향 제거
• O(1) 순방향 패스로 RankGPT 대비 60% 이상 지연시간 감소

Contrastive Retrieval Heads (arXiv:2510.02219)

• 모든 Attention Head가 동등하지 않다는 관찰
• 전체 헤드의 1% 미만만으로 최신 재순위기 성능 달성
• 유용한 헤드는 중간 레이어에 집중 분포

실험 설계

두 모델로 동일한 실험을 수행했습니다:

입력 구성

messages = [
    {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
    {"role": "assistant", "content": answer_hint_prefix},  # 예: "영희:"
]

answer_hint_prefix의 마지막 토큰이 앵커 역할을 합니다. 이 토큰이 Context의 어떤 부분에 Attention을 집중하는지를 측정합니다.

처리 과정

1. 특수 토큰 제거: 채팅 템플릿의 <end_of_turn>, <|im_end|> 등 제거
2. Attention Sink 회피: Context 영역의 토큰만 대상으로 점수 계산
3. 노이즈 토큰 제외: 마침표(.), 개행(\n\n) 등 의미 없이 높은 점수를 받는 토큰 제거
4. 문장별 평균 점수 계산: 마침표를 기준으로 문장을 구분하고 각 문장의 평균 Attention 점수를 산출

핵심 코드

def aggregate_attention_scores(inputs, layer_numbers):
    with torch.no_grad():
        outputs = model(**inputs)
        attentions = outputs.attentions

    target_index = inputs['input_ids'].shape[1] - 1  # 마지막 토큰(앵커)

    per_layer_attentions = []
    for layer_num in layer_numbers:
        attention_matrix = attentions[layer_num].squeeze(0).mean(dim=0).cpu().float().numpy()
        focused_attention = attention_matrix[target_index, :]  # 앵커가 주목하는 대상
        per_layer_attentions.append(focused_attention)

    return per_layer_attentions

실험 시나리오

다음과 같은 Context를 고정하고, 다양한 상황의 질문을 던져 Attention 기반 Reranking이 적절한 문장을 상위로 올리는지 확인했습니다:

영희는 23살 입니다.
영희는 서울대학교 학생입니다.
철수는 지방대 학생입니다.
철수는 25살 입니다.
영희가 민수와 사귀고 있습니다.
영희는 철수 옆집에 살고 있습니다.
민수는 철수와 중학교 동창입니다.
민수는 대학생입니다.
영희는 민수에게 폭행을 당했습니다.
철수는 아르바이트를 하여 최근 월급을 받았습니다.
영희는 철수의 친척입니다.
영희는 매일 운동하는 습관을 가지고 있습니다.
영희는 돈을 갚지 않는 습관이 있습니다.

시나리오 1: “영희가 울고 있습니다. 왜 울어?”

→ “영희는 민수에게 폭행을 당했습니다"가 상위에 랭크되는지 확인

시나리오 2: “택시를 탔습니다. 어디로 갈까요?”

→ “영희는 서울대학교 학생입니다"가 상위에 랭크되는지 확인

시나리오 3: “돈을 빌려달라고 합니다”

→ “영희는 돈을 갚지 않는 습관이 있습니다”, “철수는 최근 월급을 받았습니다"가 상위에 랭크되는지 확인

관찰 결과

효과적인 점

• 상황에 맞는 문장이 상위에 랭크됨: 대부분의 시나리오에서 직관적으로 관련성 높은 문장이 높은 점수를 받음
• 추가 학습 없이 작동: 별도의 Reranker 모델 학습 없이, 기존 LLM의 Attention만으로 동작
• 전체 레이어 합산이 효과적: 특정 레이어보다 전체 레이어를 합산했을 때 더 안정적인 결과

한계 및 발견

• 중간 레이어가 더 좋을 수 있음: CoRe-R 논문의 발견과 일치하게, 모든 레이어를 사용하는 것보다 중간 레이어만 사용하는 것이 더 나은 경우가 있음
• 모델 크기에 따른 성능 차이: 작은 모델에서는 Attention 신호의 품질이 떨어짐
• Attention Sink: 첫 번째 토큰이나 특수 토큰에 비정상적으로 높은 Attention이 집중되는 현상 → Context 영역만 대상으로 해야 함
• 개행 토큰의 높은 점수: \n\n 등이 의미와 무관하게 높은 점수를 받음 → 제외 필요
• 문장 구분: PoC 수준에서 마침표를 기준으로 구분했지만, 다국어 지원을 위해서는 <sep> 토큰 활용이 적절

모델별 차이

Qwen3-Reranker는 이름에서 알 수 있듯 Reranking에 특화된 모델이라, Attention 신호의 품질이 더 좋을 것으로 기대했으나, 두 모델 모두 유사한 수준의 결과를 보였습니다.

결론 및 향후 방향

LLM의 Attention Map은 별도의 학습 없이도 문서/문장 관련성을 판단하는 데 활용할 수 있는 유용한 신호를 포함하고 있습니다.

활용 가능한 방향:

• RAG 컨텍스트 압축: 긴 검색 결과에서 관련 문장만 추려서 LLM에 전달
• 롤플레잉/대화 시스템: 캐릭터 설정(Context) 중 현재 상황에 맞는 정보만 선별
• 경량 Reranker: 별도 모델 없이 추론 중 얻어지는 Attention을 활용

개선 과제:

• 최적 레이어 조합을 모델별로 찾는 방법 (CoRe-R 방식의 대조적 헤드 선별)
• ICR 논문의 보정(Calibration) 기법 적용하여 위치 편향 제거
• 문장 구분을 마침표 대신 토크나이저의 특수 토큰으로 처리

참고 자료

• AttentionRAG: Attention-Guided Context Pruning in RAG
• Attention in LLMs Yields Efficient Zero-Shot Re-Rankers (ICLR 2025)
• Contrastive Retrieval Heads Improve Attention-Based Re-Ranking
• 실험 노트북: Gemma 3
• 실험 노트북: Qwen 3