기술 요약: UniSteer

문제 정의

활성화 기반 제어는 추론 중 내부 표현에 개입하여 파라미터 업데이트 없이 대규모 언어 모델 (LLM) 을 조종하는 경량화되고 모듈식 패러다임으로 부상했습니다. 그러나 기존 방법들은 유연성과 구성 가능성 측면에서 상당한 한계에 직면해 있습니다:

• 고정된 방향: 대비적 활성화 추가 (CAA) 및 표현 공학 (RepE) 과 같은 방법들은 대비적 예시에서 도출된 사전 정의된 고정 조종 방향에 의존합니다.
• 작업별 모듈: 학습된 개입 방법들 (예: LoReFT, ODESteer) 은 각 목표 행동마다 별도의 모듈을 학습합니다.
• 구성 실패: 이러한 접근 방식들은 독립적으로 학습된 방향들이 고차원 활성화 공간에서 서로 간섭하기 때문에 세밀한 개념이나 다중 제약 조건 (예: "유용하고, 해롭지 않으며, 간결하게 응답하라") 을 처리하는 데 어려움을 겪습니다.

핵심적인 과제는 각 목표마다 별도의 모델을 적합시키지 않고도 자연어를 사용하여 LLM 을 임의의 행동, 개념, 다중 제약 조건 요구사항으로 조종할 수 있는 통합 메커니즘을 개발하는 것입니다.

방법론: UniSteer

저자들은 잔류 스트림 활성화에 대한 조건부 분포를 학습하는 텍스트 유도 활성화 흐름 매칭 모델인 UniSteer를 제안합니다. 고정된 벡터를 학습하는 대신, UniSteer 는 활성화 공간 내에서 보편적인 조건부 속도장을 학습합니다.

1. 텍스트 조건부 활성화 흐름 매칭

UniSteer 는 $p_\theta(a^{(\ell)}_i | c, \ell, i)$라는 조건부 분포를 모델링합니다. 여기서 $a$는 층 $\ell$과 토큰 $i$에서의 잔류 스트림 활성화이며, $c$는 자연어 조건 (예: "사악하게 행동하라" 또는 "간결하고 해롭지 않게 행동하라") 입니다.

• 학습 목적: 모델은 흐름 매칭 손실을 최소화함으로써 조건부 벡터장 $v_\theta$를 학습합니다. 사전 노이즈 상태 $a_0 \sim \mathcal{N}(0, I)$와 목표 활성화 $a_1$ 사이를 선형 확률 경로 $a_t = (1-t)a_0 + ta_1$를 사용하여 보간합니다. 목표 속도는 $u_t = a_1 - a_0$입니다. 모델은 텍스트 $c$에 조건부인 이 속도를 예측하도록 학습됩니다.
• 아키텍처: 흐름 모델은 고정된 텍스트 인코더 (Qwen3-0.6B) 와 층 및 토큰 위치에 대한 학습된 임베딩에 크로스 어텐션을 수행하는 DiT 스타일 트랜스포머로 구현됩니다.

2. 흐름 역변환을 통한 추론 시 활성화 조종

고정된 LLM 을 조종하기 위해 UniSteer 는 처음부터 샘플링하는 대신 흐름 역변환을 수행합니다:

1. 부분 역변환: 생성 중 관찰된 소스 활성화 $a_{src}$는 소스 조건 $c_{src}$에 조건부인 흐름을 따라 역방향으로 이동하여 중간 노이즈 잠재 상태 $a_\tau$ (여기서 $\tau = 1 - \lambda$, $\lambda$는 편집 강도) 로 변환됩니다.
2. 전진 재생성: 잠재 상태는 목표 조건 $c_{tgt}$에 조건부인 흐름을 따라 전진하여 편집된 활성화 $\hat{a}_{edit}$를 생성합니다.
3. 주입: 편집된 활성화는 고정된 LLM 의 잔류 스트림에 다시 주입되어 후속 생성에 영향을 미칩니다.

3. 활성화 공간 분류

UniSteer 는 별도의 헤드가 없이도 분류기로 기능합니다. 텍스트 레이블의 후보 집합이 주어지면, 모델은 각 조건 하에서 입력 활성화를 재구성합니다. 가장 낮은 조건부 재구성 에너지 ($\|a - \tilde{a}(c_j)\|^2_2$) 를 생성하는 레이블이 예측값으로 선택됩니다.

주요 기여

1. 통합 공식화: 본 논문은 활성화 조종을 텍스트 조건부 활성화 수송으로 공식화하여, 이질적인 제어 목표 (행동 특성, 세밀한 개념, 다중 제약 조건) 를 위한 단일 속도장을 학습하는 조건부 흐름 매칭 모델을 도입했습니다.
2. 편집을 위한 흐름 역변환: 저자들은 자연어 조건에 기반하여 활성화를 부분적으로 역변환하고 재생성함으로써 단일 모델이 단일 행동 및 구성 조종 모두를 처리할 수 있게 하는 흐름 역변환 메커니즘을 제안합니다.
3. 편집 및 분류를 위한 통합 인터페이스: 동일한 조건부 활성화 모델이 흐름 역변환을 통한 생성 조종과 재구성 에너지를 통한 내부 상태 분류라는 이중 목적을 수행하여, 작업별 분류기나 개입 모듈의 필요성을 제거합니다.

실험 결과

저자들은 다섯 가지 설정에서 세 가지 목표 LLM (Llama-3.2-1B, Qwen2.5-1.5B, Qwen2.5-7B) 에 대해 UniSteer 를 평가했습니다:

• 행동 제어 (페르소나): UniSteer 는 "사악함"과 "아첨"과 같은 특성에서 모든 모델에 걸쳐 가장 높은 목표 특성 점수를 달성하여 CAA 및 RepE 와 같은 베이스라인을 능가했습니다.
• 진실성 조종 (TruthfulQA): UniSteer 는 원래 모델 대비 Truth*Info 점수를 향상시켜, 단순한 스타일이 아닌 사실적 정확성을 위해 조종할 수 있음을 입증했습니다.
• 세밀한 개념 조종 (AxBench): UniSteer 는 단일 공유 모델을 사용함에도 불구하고 "간결하게 응답하라"와 같은 개념 조종에서 개념별 베이스라인과 경쟁하거나 우수한 성능을 달성했습니다.
• 다중 제약 조건 지시 따르기 (RECAST): UniSteer 는 동시에 5 개 또는 10 개의 제약 조건을 성공적으로 처리하여, 여러 고정 방향 간의 간섭으로 어려움을 겪는 베이스라인들을 능가했습니다.
• 활성화 공간 분류 (ToxiGen): 재구성 에너지를 사용하여 UniSteer 는 유해 콘텐츠와 비유해 콘텐츠를 분류하는 데 있어 최첨단 또는 공동 최상의 정확도와 AUC 를 달성하여 조건 의존적 활성화 분포를 포착하는 능력을 검증했습니다.
• 토큰 수준 분석: "시작하기 (start_with)" 제약 조건에 대한 분석은 UniSteer 가 위치 인식적으로 편집을 적용하여 특정 토큰 위치 (예: 응답의 시작) 에서 제약 조건 방향과 강력하게 정렬되며 전역적 교란을 적용하지 않음을 보여주었습니다.

중요성 및 주장

본 논문은 UniSteer 가 LLM 조종 및 분류를 위한 통합 인터페이스를 제공한다고 주장합니다. 고정된 방향이나 작업별 모듈에서 텍스트 조건부 속도장으로 전환함으로써, 이 방법은 기존 방법들이 결합하기 어려운 구성 제약 조건과 세밀한 개념을 자연스럽게 처리합니다.

저자들은 이 접근 방식이 다음을 가능하게 한다고 강조합니다:

• 다용도성: 단일 모델이 페르소나, 진실성, 제약 조건 등 다양한 작업을 단순히 텍스트 조건을 변경함으로써 처리할 수 있습니다.
• 구성 가능성: 다중 제약 조건 요구사항은 조건 텍스트에 직접 표현되어 별도의 조종 벡터를 결합할 때 발생하는 간섭 문제를 피합니다.
• 해석 가능성: 재구성 에너지를 통해 활성화를 분류하는 모델의 능력은 학습된 흐름이 모델 내부 상태의 의미 구조를 효과적으로 포착함을 시사합니다.

저자들은 장문 생성, 다중 턴 안정성, 복잡한 추론 작업에 대한 평가 부재와 같은 한계를 지적했습니다. 또한 안전 고려사항을 강조하며, 유용성을 개선하는 데 사용되는 동일한 메커니즘이 이론적으로 바람직하지 않은 행동을 증폭하는 데 사용될 수 있음을 인정하여 향후 안전 필터와 감사가 필요하다고 밝혔습니다.