Concept Heterogeneity-aware Representation Steering

이 논문은 대규모 언어 모델의 내부 표현이 이질적일 수 있다는 점을 고려하여, 최적 수송 (Optimal Transport) 이론을 기반으로 클러스터 수준의 이동 벡터를 입력에 따라 가중 평균하는 'CHaRS'라는 새로운 표현 제어 방법을 제안하고, 이를 통해 기존 전역적 방향 기반 제어보다 더 효과적인 행동 제어를 가능하게 함을 보여줍니다.

핵심

1. 기존 방법: "모두에게 똑같은 나침반을 주는 것"

기존의 기술 (Global Steering) 은 언어 모델의 행동을 바꿀 때, **'한 가지 방향'**만 정해서 모든 상황에 적용했습니다.

• 비유: imagine you are a tour guide in a huge, complex library. You want to guide visitors to the "Harmless" (안전한) section.
  • 기존 방법은 "안전한 책이 있는 곳으로 가려면, 북쪽으로만 걸어라"라고 모든 사람에게 똑같은 지시를 내립니다.
  • 문제점: 도서관은 매우 복잡합니다. 어떤 구역은 북쪽으로 가면 안전하지만, 다른 구역은 북쪽으로 가면 오히려 위험한 책 (해로운 내용) 에 닿을 수도 있습니다. 또한, 방문객들이 서 있는 위치 (맥락) 가 다르면 같은 '북쪽'이 의미하는 바가 달라질 수 있습니다.
  • 결과: 모든 사람에게 똑같은 지시를 내리면, 어떤 사람은 안전해지지만 다른 사람은 여전히 위험한 곳에 머물거나, 오히려 엉뚱한 곳으로 가게 됩니다. 이를 **'균질성 (Homogeneity) 가정'**이라고 하는데, 현실은 그렇게 단순하지 않습니다.

2. 새로운 방법 (CHaRS): "맞춤형 지도와 스마트 안내"

이 논문은 "아, 도서관의 각 구역은 서로 다른 특징을 가지고 있구나!"라고 깨닫고, **최적 수송 (Optimal Transport)**이라는 수학적 개념을 이용해 문제를 해결했습니다.

• 핵심 아이디어:

  • 언어 모델의 내부 지식 공간은 하나의 큰 덩어리가 아니라, 여러 개의 **'작은 군집 (Cluster)'**으로 나뉘어 있습니다. (예: '거짓말'이라는 개념도, 농담으로 하는 거짓말, 악의적인 사기, 실수 등 여러 형태로 존재합니다.)
  • CHaRS 는 이 작은 군집들을 각각 찾아내고, '출발지 (해로운 말)'의 각 군집을 '목적지 (안전한 말)'의 해당 군집으로 가장 자연스럽게 연결하는 맞춤형 지도를 그립니다.

• 비유 (여행지 안내):

  • 이제 가이드는 "북쪽으로 가라"라고 말하지 않습니다.
  • 대신, "당신이 지금 A 구역에 있다면 동쪽으로, B 구역에 있다면 서쪽으로, C 구역에 있다면 남쪽으로 이동하세요"라고 상황에 따라 다른 지시를 내립니다.
  • 그리고 이 지시들은 부드럽게 연결되어 있어서, 방문객이 한 구역에서 다른 구역으로 넘어갈 때 갑자기 방향이 뚝 끊기지 않고 자연스럽게 이동합니다.

3. CHaRS 가 어떻게 작동하나요? (3 단계)

1. 군집화 (Clustering): 먼저, 해로운 말과 안전한 말의 데이터들을 분석하여 비슷한 것끼리 묶어 '작은 그룹'을 만듭니다. (예: '사기성 질문' 그룹, '폭력성 질문' 그룹 등)
2. 매칭 (Optimal Transport): 각 '해로운 그룹'을 가장 잘 대응할 수 있는 '안전한 그룹'을 찾아 연결합니다. 마치 퍼즐 조각을 맞춰가며 가장 자연스러운 경로를 찾는 것과 같습니다.
3. 부드러운 이동 (Steering): 사용자가 질문을 입력하면, 모델은 그 질문이 어느 '그룹'에 가깝게 속하는지 보고, 해당 그룹에 맞는 맞춤형 이동 방향을 적용합니다.

4. 왜 이것이 중요한가요? (실제 효과)

이 논문은 이 방법이 실제로 얼마나 좋은지 여러 실험으로 증명했습니다.

• 해킹 방지 (Jailbreaking): 악의적인 사용자가 모델을 속여서 유해한 내용을 출력하게 하려는 시도를 막을 때, 기존 방법보다 훨씬 효과적으로 차단했습니다. (기존 방법은 뚫리기 쉬웠지만, CHaRS 는 상황마다 다른 방어 태세를 취해서 뚫기 어렵게 만들었습니다.)
• 독성 제거 (Toxicity Mitigation): 모델이 모욕적이거나 독이 되는 말을 할 때, 이를 부드럽게 고쳐주면서도 모델의 일반적인 대화 능력은 떨어뜨리지 않았습니다.
• 이미지 스타일 조절: 텍스트뿐만 아니라, 이미지 생성 모델에서도 '사이버펑크'나 '스케치' 같은 스타일을 더 정확하게 적용할 수 있었습니다.

5. 요약: 한 마디로 정리하면?

기존의 방법은 **"모두에게 똑같은 나침반을 주고 북쪽으로 가라"**고 했다면,
이 새로운 방법 (CHaRS) 은 **"당신이 어디에 있느냐에 따라, 가장 안전하고 자연스러운 길을 찾아주는 스마트한 GPS"**를 제공한다는 것입니다.

이는 인공지능이 더 유연하고, 상황에 민감하며, 안전하게 행동하도록 돕는 중요한 기술적 발전입니다.

기술 요약

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 의 행동을 제어하기 위해 내부 활성화 (internal activations) 에 개입하는 표현 제어 (Representation Steering) 기법이 주목받고 있습니다. 기존 방법론 (예: Difference-in-Means, ActAdd 등) 은 주로 대조적인 데이터셋 (예: 유해 vs 무해) 간의 평균 차이를 계산하여 단일 전역 방향 벡터 (single global steering vector) 를 생성하고, 이를 모든 입력에 동일하게 적용합니다.

그러나 이러한 접근법은 다음과 같은 근본적인 한계를 가집니다:

• 동질성 가정의 오류: 기존 방법은 타겟 개념이 임베딩 공간에서 균일하게 분포되어 있다고 가정합니다.
• 실제 구조의 복잡성: 실제 LLM 의 표현 공간은 고차원적이며, 맥락이나 잠재 하위 개념에 따라 클러스터링된 비균질 (heterogeneous) 구조를 보입니다.
• 취약한 제어: 단일 전역 방향 벡터는 이러한 복잡한 구조를 무시하기 때문에, 맥락에 따라 일관되지 않은 제어 결과를 초래하거나 원치 않는 부작용을 일으킬 수 있습니다.

2. 방법론 (Methodology)

저자들은 표현 제어를 최적 수송 (Optimal Transport, OT) 의 관점에서 재해석하고, 이를 가우시안 혼합 모델 (Gaussian Mixture Models, GMMs) 을 기반으로 확장하여 CHaRS를 제안했습니다.

핵심 아이디어

1. 표현의 GMM 모델링:

   • 기존 방법은 두 개념을 단일 가우시안 분포로 간주하여 평균 차이 (전역 이동) 로 접근했습니다.
   • CHaRS 는 소스 (Source) 와 타겟 (Target) 표현을 가우시안 혼합 모델 (GMM) 로 모델링합니다. 즉, 각 개념이 여러 개의 클러스터 (잠재 하위 의미 영역) 로 구성되어 있다고 가정합니다.

2. 이산 최적 수송 (Discrete OT) 문제로서의 정립:

   • 두 GMM 간의 정렬을 Mixture Wasserstein Distance를 사용하여 정의합니다.
   • 이는 두 분포 간의 전역적인 이동이 아니라, 소스 클러스터와 타겟 클러스터 간의 이산적인 매칭 (discrete matching) 문제로 변환됩니다.
   • 엔트로피 정규화 (Entropy Regularization) 가 적용된 Sinkhorn 알고리즘을 사용하여 클러스터 간의 최적 결합 (coupling) 을 효율적으로 계산합니다.

3. 입력 의존적 매핑 (Input-Dependent Steering Map):

   • 계산된 결합 계획 (transport plan) 을 바탕으로 바리센트릭 투영 (Barycentric Projection) 을 수행합니다.
   • 결과적으로, 특정 입력 $x$에 대해 해당 입력이 속한 클러스터와 매칭된 타겟 클러스터 간의 이동 벡터를 커널 가중치 (kernel-weighted) 로 부드럽게 조합한 국소적 이동 벡터를 생성합니다.
   • 수식적으로, 최종 이동 벡터 $\hat{v}(x)$는 다음과 같이 표현됩니다:
     $$\hat{v}(x) = \sum_{i,j} P^*_{ij} \cdot k(x, a_i) \cdot (b_j - a_i)$$
     (여기서 $P^*_{ij}$는 클러스터 매칭 확률, $k(x, a_i)$는 입력과 클러스터 중심 간의 거리 기반 커널 가중치, $b_j - a_i$는 클러스터 간 이동 벡터입니다.)

4. 주성분 임계값 (Principal Component Thresholding, PCT):

   • 생성된 이동 벡터들의 공분산 행렬은 본질적으로 저랭크 (low-rank) 성질을 가집니다.
   • 이를 활용하여 CHaRS-PCT를 제안합니다. 주성분 분석 (PCA) 을 통해 주요 의미 축만 선별하여 노이즈를 제거하고, 더 적은 수의 방향 벡터로 효율적인 제어를 가능하게 합니다.

3. 주요 기여 (Key Contributions)

1. 단일 가우시안에서 GMM 로의 일반화: 표현 제어를 제한적인 단일 가우시안 가정에서 다중 모드 (multimodal) 가우시안 혼합 모델로 확장하여, LLM 의 비균질한 개념 구조를 정량적으로 모델링했습니다.
2. CHaRS 프레임워크 개발: 클러스터 수준의 수송 계획을 기반으로 한 입력 적응형 (input-adaptive) 제어 메커니즘을 제안했습니다. 이는 표현 매니폴드 전체에 걸쳐 부드럽게 변하는 방향 벡터를 생성하여 맥락에 민감한 제어를 가능하게 합니다.
3. CHaRS-PCT 및 저랭크 구조 활용: 이동 벡터의 공분산 구조가 저랭크임을 이론적으로 증명하고, 주성분 임계값을 적용하여 계산 효율성을 높이고 노이즈를 줄인 변형 알고리즘을 제시했습니다.
4. 광범위한 실험적 검증: 다양한 모델 (3B~32B 파라미터) 과 작업 (자일브레이킹, 독성 완화, 이미지 스타일 제어) 에서 기존 방법론 (ActAdd, DirAbl, Linear-Act) 을 능가하는 성능을 입증했습니다.

4. 실험 결과 (Results)

저자들은 Gemma2, Llama3, Qwen2.5 등 다양한 오픈 가중치 LLM 과 FLUX.1(이미지 생성 모델) 을 사용하여 실험을 수행했습니다.

• 자일브레이킹 (Jailbreaking):

  • 유해한 출력을 유도하는 공격 성공률 (ASR) 에서 CHaRS 와 CHaRS-PCT 가 ActAdd 및 DirAbl 기반 베이스라인보다 일관되게 높은 성능을 보였습니다 (최대 7% 향상).
  • 동시에 일반 언어 모델링 능력 (tinyBenchmarks) 을 유지하며, 생성 품질 저하를 최소화했습니다.

• 독성 완화 (Toxicity Mitigation):

  • 순차적 (sequential) 인 레이어별 제어 설정에서 기존 방법인 Linear-Act 보다 우수한 독성 감소 효과를 보였습니다.
  • 특히 CHaRS-PCT 는 레이어 간 노이즈 누적을 줄여 더 안정적인 성능을 발휘했습니다.
  • perplexity 나 하류 작업 성능을 저하시키지 않으면서 독성 생성을 효과적으로 억제했습니다.

• 이미지 스타일 제어 (Image Style Control):

  • 텍스트 - 이미지 생성 모델 (FLUX.1) 에서 '사이버펑크'나 '스케치'와 같은 특정 스타일을 유도하는 데 성공했습니다.
  • Linear-Act 대비 더 낮은 제어 강도 (steering strength) 로 목표 스타일을 효과적으로 구현하면서도, 원본 프롬프트의 콘텐츠 보존도 (CLIPScore) 를 높게 유지하여 스타일 유도와 콘텐츠 보존 간의 최적의 균형 (Pareto front) 을 달성했습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 표현 공간이 단순한 선형 구조가 아니라 클러스터링된 비균질한 구조임을 인식하고, 이를 해결하기 위해 최적 수송 (Optimal Transport) 이론을 도입했다는 점에서 의미가 큽니다.

• 이론적 기여: 표현 제어를 단순한 선형 이동이 아닌, 의미적 클러스터 간의 정렬 문제로 재정의함으로써 더 정교한 비선형 제어의 이론적 기반을 마련했습니다.
• 실용적 기여: CHaRS 는 모델의 내부 구조를 더 잘 이해하고 활용하여, 더 강력하면서도 안전한 행동 제어를 가능하게 합니다. 이는 AI 안전 (Safety), 편향 제거, 그리고 창의적 콘텐츠 생성 등 다양한 분야에서 모델의 행동을 정밀하게 조절할 수 있는 새로운 패러다임을 제시합니다.
• 한계 및 향후 과제: 현재는 등방성 (isotropic) 공분산과 k-means 클러스터링을 사용하여 안정성을 확보했으나, 향후 이방성 (anisotropic) 혼합 모델과 특징 가중치 메커니즘을 도입하여 표현 공간의 방향적 뉘앙스를 더 정밀하게 포착할 수 있을 것으로 기대됩니다.

요약하자면, CHaRS는 LLM 의 복잡한 내부 표현 구조를 고려하여, 맥락에 따라 유연하게 적응하는 차세대 표현 제어 기술로, 기존 방법론의 한계를 극복하고 더 효과적이고 안전한 AI 제어 체계를 구축하는 데 중요한 기여를 했습니다.