기술 요약: 매니폴드 기반 어텐션 조종 (MAGS)

문제 제기

대형 언어 모델 (LLM) 은 종종 복잡한 추론 작업에 필요한 근본적인 지식을 보유하고 있음에도 불구하고, 생성 과정에서 이러한 능력을 신뢰성 있게 표현하지 못하는 경우가 많습니다. 동일한 모델과 프롬프트를 반복적으로 샘플링하면 정답과 오답이 모두 도출될 수 있으며, 과정 수준의 주석에 따르면 오류는 일반적으로 지식의 부재가 아닌 중간 추론 단계에서 발생합니다.

기존의 추론 시간 활성화 조종 (inference-time activation steering) 방법들은 잔여 스트림이나 어텐션 출력에 고정된 사전 계산된 교정 벡터를 적용함으로써 이러한 행동을 수정하려 시도합니다. 그러나 이러한 접근법은 추론 작업과 중요한 구조적 불일치를 겪습니다:

1. 무분별한 교란: 모델의 현재 상태와 관계없이 교정을 균일하게 적용합니다. 이는 이미 올바른 단계를 손상시키면서도 아직 나타나지 않은 오류를 차단하지 못하게 할 위험이 있습니다.
2. 정적 성격: 추론 궤적이 오류 상태로 치우칠 때를 감지할 메커니즘이 부족하여, 모든 생성 단계를 교정이 필요한 것으로 동일하게 취급합니다.

이 연구의 핵심 가설은 추론 오류가 개별 어텐션 헤드의 출력 공간에서 저차원 부분 공간 내의 구조화된 편향 (drift) 으로 나타난다는 것입니다. 올바른 궤적과 잘못된 궤적은 기하학적으로 분리 가능한 영역을 차지하며, 올바른 행동에서 잘못된 행동으로의 전환은 특정 저랭크 (low-rank) 방향을 따릅니다.

방법론: 매니폴드 기반 어텐션 조종 (MAGS)

MAGS 는 "정확성 매니폴드"에서의 편향이 감지될 때만 어텐션 헤드 출력을 조종하는 적응형, 궤적 인지 추론 시간 개입입니다. 이 방법은 세 가지 단계로 작동합니다:

1. 대비 오류 매니폴드 구성

훈련 단계 동안, 시스템은 동일한 문제에 대한 대비 쌍의 추론 흔적 (하나의 정답, 하나의 오답) 을 분석하여 오류의 기하학적 구조를 학습합니다.

• 헤드별 차이 벡터: 각 어텐션 헤드 $(l, h)$에 대해 시스템은 올바른 흔적 ($\mu_c$) 과 잘못된 흔적 ($\mu_e$) 에 대한 평균 활성화를 계산합니다. 차이 벡터 $\delta = \mu_e - \mu_c$는 문제 특정 특징을 상쇄하고 오류로 인한 방향적 이동을 분리해 냅니다.
• 부분 공간 학습: 이러한 차이 벡터를 여러 문제에 걸쳐 쌓고 주성분 분석 (PCA) 을 수행함으로써, MAGS 는 저차원 기저 $B$(상위 $k$개 특이 벡터) 를 추출합니다. 이 기저는 "오류 부분 공간"을 spanning 하는데, 이는 헤드의 출력 공간에서 올바른 상태에서 잘못된 상태로의 편향이 가장 큰 분산을 보이는 방향들입니다.
• 전역 기준점: 모든 올바른 흔적에서 계산된 전역 중심점 $\mu_c$가 추론 중 기준점으로 사용됩니다.

2. 근접 기반 오류 감지

추론 중 MAGS 는 선택된 어텐션 헤드의 출력을 실시간으로 모니터링합니다.

• 근접 점수: 각 토큰 단계 $t$에서 시스템은 학습된 오류 부분 공간으로 투영된 중심화된 헤드 출력의 제곱 노름을 계산합니다:
  $$d_t = \| B(a_t - \mu_c) \|^2$$
  높은 점수는 헤드의 활성화가 오류 매니폴드 방향으로 크게 편향되었음을 나타냅니다.
• 동적 트리거링: 각 헤드의 임계값 $\tau$는 올바른 흔적에서 보정됩니다 (예: 올바른 궤적의 점수 상위 백분위수로 설정). $d_t > \tau$일 때만 교정이 적용됩니다. 이는 이미 올바른 궤적에 있는 단계가 교란되지 않도록 보장합니다.

3. 표적 교정

트리거가 발동되면 MAGS 는 해당 헤드의 출력이 레이어의 출력 투영으로 전달되기 전에 제자리 교정을 적용합니다:
$$\tilde{a}_t = a_t - \alpha B^\top B (a_t - \mu_c)$$
여기서 $\alpha \in (0, 1]$는 조종 강도 하이퍼파라미터입니다.

• 기하학적 해석: 이 연산은 편향 $(a_t - \mu_c)$을 오류 부분 공간으로 투영하고 이를 빼냄으로써, 오류 성분을 효과적으로 제거하면서 직교 여집합 (오류 매니폴드와 관련 없는 $(d_h - k)$ 차원) 에 있는 모든 정보를 보존합니다.
• 정보 보존: 이 방법은 오류 매니폴드와 무관한 정보가 보존되도록 보장하여, 정적 방법에서 관찰되는 "과도 교정"을 방지합니다.

다목적 확장 (MAGSu)

상충되는 목표 (예: 문법적 유효성 대 화학적 결합 친화도) 를 처리하기 위해 저자들은 MAGSu 를 제안했는데, 이는 각 목표에 대해 독립적인 매니폴드를 학습하고 교차 간섭을 방지하기 위해 불연속적인 헤드 집합을 통해 교정을 적용합니다.

주요 기여

1. 가설 검증: 이 논문은 추론 오류가 특정 어텐션 헤드의 저차원 부분 공간 내의 구조화된 편향으로 나타난다는 경험적 증거를 제공합니다. 진단 실험은 올바른 궤적과 잘못된 궤적이 이러한 부분 공간에 의해 매우 잘 분리되며, 오류 신호가 희소하고 특정 헤드에 집중되어 있음을 보여줍니다.
2. 적응형 개입 메커니즘: MAGS 는 헤드가 오류 부분 공간으로 편향될 때만 교정을 적용하는 동적 근접 기반 트리거를 도입합니다. 이는 현재 활성화 상태와 관계없이 고정된 벡터를 적용하는 정적 방법과 대조됩니다.
3. 경험적 우위성: 이 방법은 다양한 벤치마크에서 조종되지 않은 베이스라인 및 기존 정적 조종 접근법 (ITI, Angular Steering) 보다 일관되게 우수한 성과를 거둡니다.

실험 결과

저자들은 세 가지 모델 계열 (Llama-3.1-8B-Instruct, Gemma-4-E4b-it, GPT-OSS-20B) 을 세 가지 도메인에 걸쳐 평가했습니다:

• 수학적 추론 (MATH-500, GSM8K): MAGS 는 MATH-500 에서 가장 큰 개선을 이루었으며, 조종되지 않은 Llama 베이스라인 대비 정확도를 +5.2% 향상시켰고, 조종되지 않은 Gemma 베이스라인 대비 +3.4% 향상시켰습니다. 오류가 누적되기 쉬운 다단계 유도 작업에서 개선이 가장 두드러졌습니다.
• 코드 생성 (HumanEval, MBPP): MAGS 는 조종되지 않은 베이스라인 대비 HumanEval pass@1 을 약 4 포인트 향상시켰습니다. 주목할 점은 정적 방법 (ITI, Angular Steering) 이 문법적 일관성을 교란시켜 코드 생성 성능을 저하시킨 반면, MAGS 는 올바른 단계에서의 교정을 억제함으로써 이를 피했다는 것입니다.
• 분자 생성 (SMILES): 문법적 유효성과 높은 결합 친화도 모두를 요구하는 다목적 작업에서 MAGSu 는 두 지표를 동시에 성공적으로 향상시켰습니다. 반면, ITI 는 유효성은 향상시켰지만 결합 친화도를 크게 저하시켜 정적 조종에서 얽힌 목표가 야기하는 간섭을 보여주었습니다.
• 유창성과 오버헤드: MAGS 는 조종되지 않은 베이스라인과 비교 가능한 퍼플렉시티 수준을 유지하여 생성 분포를 왜곡하지 않음을 나타냈습니다. 계산 오버헤드는 무시할 수 있으며, 각 단계당 $O(K \cdot k \cdot d_h)$ 연산이 주를 이루며, 여기서 $K$는 모니터링되는 헤드의 수입니다.

중요성과 주장

이 논문은 MAGS 가 활성화 조종을 일정한 교정 작업이 아닌 편향 감지 문제로 재정의한다고 주장합니다. 오류 발발을 예측하는 저랭크 기저를 식별함으로써 MAGS 는 스칼라 크기 기반 방법들이 lacking 하는 해석 가능한 기하학적 구조를 제공합니다.

저자들은 다음과 같이 주장합니다:

1. 정확성 매니폴드의 보편성: 올바른 행동과 잘못된 행동을 분리하는 저차원 부분 공간의 존재는 추론, 코딩, 과학적 생성 전반에 적용 가능한 LLM 어텐션 기하학의 일반적인 특징으로 보입니다.
2. 궤적 인지의 중요성: 무분별한 교정은 유효한 추론 단계를 교란시킵니다. 편향이 감지될 때 만 개입할 수 있는 능력이 MAGS 의 우수한 성능과 하이퍼파라미터 민감성에 대한 강건성의 주요 동인입니다.
3. 기하학적 제약이 과도 교정을 방지: 활성화를 오류 부분 공간의 여집합으로 투영함으로써 MAGS 는 오류 특정 성분만 제거하면서 모델의 고유 능력을 보존합니다.

이 연구는 MAGS 가 매니폴드 피팅을 위해 대비 흔적 집합이 필요하지만, 생성 궤적의 동적 특성을 고려하지 못하는 정적 개입보다 LLM 의 추론 일관성을 개선하기 위한 견고하고 오버헤드가 낮은 솔루션을 제공한다고 결론지었습니다.