기술 요약: Cobras (리만 활성화 스티어링을 위한 조건부 최적 브릿지)

문제 정의

활성화 스티어링(Activation steering)은 대규모 언어 모델(LLM)의 추론 시점에 제어하기 위한 미세 조정(fine-tuning)의 경량화된 대안으로 부상하였으며, 고차원 속성이 활성화 공간 내의 선형 방향에 대응한다는 선형 표현 가설(linear representation hypothesis) 하에 작동합니다. 그러나 기존 방식들은 다음과 같은 두 가지 주요 한계점을 가집니다:

1. 휴리스틱 목적 함수: 대부분의 방법은 로그 밀도 비율($\nabla \log \frac{p_+}{p_-}$) 목적 함수를 원칙적인 최적화 문제로부터 유도하지 않고 휴리스틱하게 암묵적으로 최적화합니다.
2. 쿼리 독립성(Query-Independence): 현재의 스티어링 신호는 일반적으로 쿼리에 독립적입니다(고정된 벡터 또는 변환). 이러한 경직성은 분포 외(Out-of-Distribution, OOD) 입력에서 심각한 성능 저하를 초래할 수 있습니다. 즉, 인-디스트리뷰션(in-distribution) 프롬프트에는 유익한 방향이 OOD 쿼리에 대해서는 횡설수설하거나 잘못된 출력을 생성하게 만들 수 있습니다.

방법론: Cobras

저자들은 Cobras(Conditional Optimal Bridge for Riemannian Activation Steering)를 제안하며, 이는 활성화 스티어링을 잔차 스트림(residual-stream) 초구체(hypersphere) 상의 슈뢰딩거 브릿지(Schrödinger Bridge, SB) 문제로 재정의합니다.

1. 리만 기하학 정식화

이 방법은 트랜스포머 레이어의 잔차 스트림 활성화가 노름(norm)에 강력하게 집중되어, 실질적으로 초구체 $S = \{x \in \mathbb{R}^d : \|x\| = R\}$ 상에 존재한다는 관찰에 근거합니다. Cobras는 다음을 활용하여 리만 다양체 내에서 완전히 작동합니다:

• 접사 투영기(Tangent Projectors): $\Pi_x = I - R^{-2}xx^\top$.
• 측지 거리(Geodesic Distance): $d_S(x, y) = R \arccos(R^{-2}x^\top y)$.
• 지수/로그 맵(Exponential/Logarithm Maps): 다양체와 접공간(tangent spaces) 사이를 이동하기 위함.

2. 슈뢰딩거 브릿지 목적 함수

단순히 벡터를 더하는 대신, Cobras는 원형(sphere) 상의 참조 브라운 운동(Brownian motion)으로부터의 쿨백-라이블러(Kullback-Leibler, KL) 발산을 최소화하는 "노력"을 사용하여, 원치 않는 활성화 분포($p_-$)를 원하는 활성화 분포($p_+$)로 운송하는 최적의 확률 과정을 탐색합니다.

• 최적화: 이 문제는 마진 $P_0 = p_-$ 및 $P_T = p_+$를 만족하면서 $KL(P \| Q)$를 최소화하는 문제로 구성됩니다.
• 해법: 이는 싱크혼(Sinkhorn) 알고리즘을 통해 해결 가능한 엔트로피 최적 운송(entropic optimal transport) 문제로 귀결되며, 쌍대 포텐셜(dual potentials) $\{\phi_j\}$와 $\{\psi_i\}$를 산출합니다.

3. 스티어링 방향의 유도

해법은 확률 미분 방정식(SDE)을 유도합니다. 저자들은 이 SDE에 **확률 흐름 ODE(Probability Flow ODE)**를 적용하여 결정론적 드리프트(drift) 항을 도출합니다:
$$v(x) \propto \nabla_S \log \frac{\hat{\psi}(x, t)}{\hat{\phi}(x, t)}$$
여기서 $\hat{\psi}$와 $\hat{\phi}$는 현재 쿼리 활성화 $x$에서 평가된 슈뢰딩거 포텐셜입니다.

• 원칙적 유도: 논문은 (Proposition 1) 운송 비용이 소멸하는 극한에서, 이 드리프트가 널리 사용되는 휴리스틱인 $\nabla_S \log \frac{p_+}{p_-}$로 수렴함을 증명합니다. 이는 로그 밀도 비율 목적 함수에 대한 최초의 엄격한 최적화 기반 유도를 제공합니다.
• 쿼리 적응성(Query-Adaptivity): 결정적으로, 포텐셜이 현재 쿼리 활성화에서 평가되기 때문에, 결과적인 스티어링 방향은 본질적으로 문맥 의존적입니다.

4. 알고리즘 구현

Cobras 알고리즘은 두 단계로 작동합니다:

• 오프라인(Offline): 훈련 샘플($D_+, D_-$)을 초구체 상에 투영하고, 포텐셜을 계산하기 위해 싱크혼 반복을 통해 정적 SB 문제를 해결합니다.
• 온라인(Online, 추론):
  1. 비매개변수적 추정(Non-parametric Estimation): 커널 밀도 추정(KDE)을 사용하여 임의의 쿼리 지점으로 포텐셜을 확장합니다.
  2. ODE 적분: 구체 상에서의 측지 오일러 단계(geodesic Euler steps)를 사용하여 확률 흐름 ODE를 적분함으로써 쿼리 활성화를 스티어링합니다.
  3. 게이팅 메커니즘(Gating Mechanisms):
     • 기권 게이트(Abstain Gate): 쿼리의 훈련 분포와의 거리를 기준 반경과 비교합니다. 쿼리가 멀리 있는 경우(OOD), 성능 저하를 방지하기 위해 스티어링을 약화시킵니다.
     • 강도 게이트(Strength Gate): 쿼리의 평균 대조 방향(mean contrastive direction)과의 정렬 상태에 따라 스티어링을 조절하여, 이미 정렬된 입력에 대한 개입을 줄입니다.

주요 기여

1. 원칙적 유도: 잘 정의된 리만 다양체 상의 슈뢰딩거 브릿지 최적화 문제로부터 로그 밀도 비율 스티어링 목적 함수를 도출하는 최초의 유도를 제공합니다.
2. 쿼리 적응형 스티어링: 스티어링 방향이 특정 쿼리에 따라 동적으로 계산되도록 하여, 고정된 벡터 방식의 "일률적(one-size-fits-all)" 실패 모드를 방지하는 방법을 도입합니다.
3. 실험적 우월성: 4개의 모델(Falcon-7B, Mistral-7B, LLaMA3.1-8B, Qwen2.5-7B)과 3개의 정렬 축(도움됨, 진실성, 독성 제거)에 대해 일관된 우수한 성능을 입증했습니다.

실험 결과

저자들은 TruthfulQA (진실성), UltraFeedback (도움됨), RealToxicityPrompts (독성 제거)에 대해, 그리고 OOD 벤치마크(MMLU, GSM8K)와 함께 Cobras를 평가했습니다.

• 인-디스트리뷰션 성능: Cobras는 모든 모델과 작업에서 일관되게 가장 높은 주요 지표를 달성했습니다. 예를 들어, TruthfulQA에서 Mistral-7B의 경우 True×Info 점수를 +29.5 포인트, LLaMA3.1-8B의 경우 +25.1 포인트 개선하였으며, 이는 원본 모델 및 ODESteer, CAA, SphericalSteer과 같은 베이스라인보다 뛰어난 성능입니다.
• OOD 보존: 인-디스트리뷰션 지표를 높이는 대신 OOD 정확도가 급격히 떨어지는(예: GSM8K에서의 심각한 하락) 많은 베이스라인과 달리, Cobras는 OOD 성능을 스티어링되지 않은 베이스라인에 가깝거나 그 이상으로 유지했습니다. 이는 훈련 지원 범위(training support)에서 멀리 떨어진 입력에 대해 스티어링을 억제하는 기권 게이트 덕분입니다.
• 다양성: Cobras는 생성 품질을 저하시키는 경우가 많은 여러 베이스라인(예: CAA, SphericalSteer)보다 어휘적 다양성(Dist-n 지표)과 퍼플렉시티(perplexity)를 더 잘 보존했습니다.

의의 및 주장

이 논문은 Cobras가 휴리스틱한 활성화 스티어링과 엄격한 최적 운송 이론 사이의 간극을 메운다고 주장합니다. 주요 의의는 다음과 같습니다:

• 이론적 토대: 인기 있는 로그 밀도 비율 휴리스틱이 원칙적인 슈뢰딩거 브릿지 정식화의 극한 사례임을 확립했습니다.
• 강건성: 고정된 방향 스티어링에 내재된 OOD 저하 문제를, 개입을 본질적으로 쿼리 적응형으로 만듦으로써 해결했습니다.
• 일반성: 다양한 스티어링 패러다임(CAA, ITI, Spherical Steering 포함)을 단일 기하학적 최적화 프레임워크 아래 통합하는 동시에, 우수한 경험적 성능을 제공합니다.

저자들은 짝의 측지 거리 및 싱크혼 반복에 따른 $O(N_+ \cdot N_-)$의 계산 비용과 고정된 구형 기하학 가정에 대한 의존성을 한계점으로 언급했습니다. 향후 연구로는 계층적 운송 솔버(hierarchical transport solvers) 및 더 일반적인 리만 다양체에 대한 탐구를 제안했습니다.