기술 요약: 단계별 추론 에너지 (Step-Aware Reasoning Energy, SARE)

1. 문제 정의

Chain-of-Thought (CoT) 프롬프팅은 복잡한 추론 작업에서 대규모 언어 모델(LLM)의 성능을 크게 향큼 향상시켰으나, 개별 추론 단계 전반에 걸친 내부 계산 노력의 할당 방식은 여전히 불투명하다. 기존의 해석 가능성 방법론들은 두 가지 주요 한계를 가진다:

1. 표면 수준의 의존성: 많은 접근 방식이 출력 레벨의 신호(예: 토큰 로그 확률, 엔트로피)에 의존하거나 추론 궤적을 단일 스칼라 값으로 취급하여, 단계별 역동성을 포착하지 못한다.
2. 조대(Coarse)한 입도: 해석 가능성 연구는 종종 개별 토큰이나 붕괴된 집계 표현을 분석한다. 이러한 입도는 전체 추론 단계(예: 논리적 연역 또는 방정식)의 계산 역동성을 포착하기에는 불충분하며, 해당 단계 내 토큰들 사이의 관계적 구조를 소실시킨다.

결과적으로, 어떤 단계가 깊은 내부 처리를 요구하는지 아니면 사소하게 해결되는지, 그리고 이러한 내부 역동성이 추론의 성공 또는 실패와 어떻게 상관관계가 있는지에 대한 이해가 부족하다.

2. 방법론

저자들은 CoT 단계의 입도에서 계산 노력을 정량화하기 위해 설계된 기하학적 프레임워크인 **단계별 추론 에너지 (Step-Aware Reasoning Energy, SARE)**를 제안한다. 이 방법론은 기하학적 비유사성과 의미론적 상태 모델링을 통합한다.

2.1 단계별 정식화

CoT 궤적 $T$는 이산적인 텍스트 추론 단계 $s_1, \dots, s_T$로 분할된다. 특정 트랜스포머 레이어 $l$에서의 각 단계 $s_t$에 대해, 토큰 은닉 상태(hidden states)는 행렬 $H^{(l)}_t$로 표현된다. 내부의 의미론적 관계는 그람 행렬(Gram matrix) $G^{(l)}_t = H^{(l)}_t (H^{(l)}_t)^\top$를 통해 포착된다.

2.2 CKA를 통한 기하학적 비유사성

단계 내에서 소요되는 계산 노력을 측정하기 위해, 본 프레임워크는 인접한 레이어($l$ 및 $l+1$) 간의 그람 행렬에 대한 **중심 커널 정렬(Centered Kernel Alignment, CKA)**을 계산한다.

• 비유사도 점수: 레이어별 비유사도 $D^{(l)}_t$는 중심화된 그람 행렬 $\tilde{G}$에 대하여 $1 - \text{CKA}(\tilde{G}^{(l)}_t, \tilde{G}^{(l+1)}_t)$로 정의된다.
• 추론 에너지 ($E_t$): 단계의 총 에너지는 모든 레이어 전이(transition)에 대한 비유사도 점수의 합이다: $E_t = \sum_{l=1}^{L-1} D^{(l)}_t$.
• 해석: 높은 $E_t$는 레이어를 가로지르는 내부 토큰 관계 기하학의 상당한 재구성을 나타내며, 이는 진정한 계산 노력을 반영한다. 낮은 $E_t$는 사소한 전이나 일상적인 회상(recall)의 특징인 조기 구조적 안정화를 시사한다.
• 대안 대비 장점: 토큰 수준의 척도(예: 다음 토큰 분포에 대한 Jensen-Shannon Divergence)나 고윳값(eigenvalue) 비교와 달리, CKA는 인접 레이어 간의 고윳값 정렬이나 클러스터 대응을 요구하지 않고 쌍별 토큰 유사도 구조에 직접 작용한다.

2.3 의미론적 상태 모델링

본 프레임워크는 CoT 단계를 잠재적 의미 상태 간의 전이로 모델링함으로써 에너지의 맥락을 파악한다:

1. 클러스터링: 단계들은 최종 레이어의 누적 토큰 그람 행렬의 고윳값 스펙트럼으로부터 유도된 스펙트럼 임베딩(spectral embeddings)으로 표현된다. 이들은 매크로 상태(예: 문제 설정, 사실적 검색, 합성)를 식별하기 위해 K-Means를 통해 클러스터링된다.
2. 마르코프 역동성: 클러스터의 시퀀스는 1차 마르코프 체인으로 모델링되어 교차 단계 에너지 속도와 궤적 안정성을 분석한다.

3. 주요 기여

1. SARE 프레임워크: 토큰 수준의 척도에서 흔히 버려지는 인접 토큰 관계 역동성을 포착하기 위해, 토큰 은닉 상태 그람 행렬에 대한 CKA를 사용하여 단계별 추론 노력을 정량화하는 기하학 기반 방법론을 제시한다.
2. 경험적 특성화: 추론 에너지가 의미론적 단계 유형에 따라 매우 불균일하며, 궤적 수준의 지표로는 보이지 않는 위상 변화(phase-like transitions)를 보인다는 점을 규명한다.
3. 실패 탐지: 잘못된 궤적이 결정적인 추론 지점(특히 검증 및 합성 상태)에서 체계적으로 낮은 추론 에너지를 보인다는 것을 입증한다.
4. 예측 유용성: SARE 기반 피처가 오프라인 추론 실패 탐지에 있어 출력 기반 신뢰도 베이스라인(로그 확률, 엔트로피)과 대등하거나 이를 능가함을 보여준다.

4. 실험 결과

본 프레임워크는 6개의 벤치마크(GSM8K, MATH, CSQA, StrategyQA, HotpotQA, MuSiQue)와 3개의 오픈 웨이트 LLM(LLaMA-3.2-3B, Phi-4-mini, Gemma-3-4B)에 대해 평가되었다.

4.1 추론 에너지 프로필

• 비균일성: 에너지 분포는 의미론적 역할에 따라 크게 달라진다. 초기 설정 및 최종 합성 단계는 종종 에너지 극단을 형성하는 반면, 중간 경로의 사실적 검색 단계는 안정적이고 중간 정도의 에너지를 보인다.
• 모델 의존성: 에너지의 "동적 범위(dynamic range)"는 아키텍처에 따라 다르다. 예를 들어, Gemma-3-4B는 LLaMA-3.2-3B나 Phi-4-mini보다 더 넓은 에너지 스펙트럼을 활용한다.
• 의미론적 해석:
  • 수학적 추론(Phi-4-mini)에서 최종 단계는 종-종 가장 낮은 에너지를 보이는데, 이는 핵심 추론이 더 일찍 완료됨을 시사한다.
  • 다단계/상식 추론(Gemma-3, LLaMA)에서 최종 합성 클러스터는 가장 높은 에너지를 나타내는 경-향이 있다.

4.2 추론 실패 시그니처

• 오류 시 낮은 에너지: 잘못된 궤적은 특히 최종 검증 및 구성적 추론 단계에서 일관되게 낮은 추론 에너지를 보인다.
• 국소적 발산: 정답과 오답 궤적 간의 에너지 분포 차이는 균일하지 않으며, 초기 검색 단계보다는 후기 정량화 및 구성적 상태에서 시작되는 경우가 많다.

4.3 예측 성능

• 베이스라인 비교: SARE(토큰 수와 결합 시)는 대부분의 데이터셋에서 표준 신뢰도 베이스라인(Log-Prob, Entropy, Perplexity, Self-Certainty)과 대등하거나 이를 능가한다.
• StrategyQA 이상 현상: 확률적 베이스라인은 StrategyQA(True/False 작업)에서 정답 여부와 상관없이 토큰 분포가 균등했기 때문에 F1 점수 0을 기록하며 완전히 실패했다. 반면, SARE와 토큰 수는 예측력(F1 ~0.46–0.58)을 유지하여, 내부 기하학적 역동성이 표면적 확률 이상의 정보를 인코딩함을 입증했다.

5. 의의 및 주장

본 논문은 SARE를 직접적인 개입 메커니즘이라기보다 진단 및 해석 도구로 정의한다. 그 의의는 다음과 같다:

• 블랙박스 개방: 고윳값 정렬 없이도 단계별 노력을 정량화함으로써 LLM 추론의 "블랙박스" 내부를 들여-다볼 수 있는 창을 제공한다.
• 표면 신호 그 이상: 내부 기하학적 역동성이 출력 레벨의 신뢰도 점수로는 포착되지 않는 추론 정확도에 대한 예측 정보를 포함하고 있음을 보여준다.
• 향러 연구의 토대: 저자들은 이러한 구조화된 에너지 패턴이 결정적인 지점에서 지속적인 재구성을 유도하기 위한 학습 시간 신호나, 리샘플링을 위한 추론 시간 트리거와 같은 기하학 인식 추론 개선을 위한 기초가 될 수 있다고 제안한다.

인정된 한계점:

• 규모: 실험은 3–4B 파라미터 모델에 국한되었으며, 더 큰 모델로의 일반화 가능성은 미지수이다.
• 오버헤드: 모든 레이어에서 은닉 상태를 추출해야 하므로, SARE는 오프라인 분석에 가장 적합하다. 다만, 향후 효율성을 위해 레이어 서브샘플링과 같은 근사법이 제안되었다.