본문 바로가기
AI/Paper Review

[논문 리뷰] Steering Llama 2 via Contrastive Activation Addition

by hyeok1235 2025. 2. 12.

핵심 : 원하는 행동(behavior)의 프롬프트와 그 반대 행동의 프롬프트는 activation에서 차이가 난다. 그 둘을 빼서 steering vector로 사용하면 원하는 행동을 더 강화하거나 약화할 수 있다.

0. Abstract

해당 연구에서는 Contrastive Activation Addition(CAA)라는 방식을 소개한다. 긍정적인 input과 부정적인 input에 대해 residual stream activation의 차이값을 활용하여 유저의 프롬프트에 대해 language model을 조정할 수 있다. 파인튜닝이나 시스템 프롬프트 디자인보다 더 효과적임을 보인다. 또한, 추상적인 레벨의 컨셉들이 LLM에서 어떻게 표현되는지 엿볼 수 있다.

1. Introduction

LLM의 성능이 좋아질수록, “helpful, honest, and harmless”한 특징을 보장하는게 중요해졌다. RLHF, 파인튜닝, 프롬프트 엔지니어링의 방식으로 LLM을 조정하는 시도들은 있지만, 다양한 데이터셋을 모으는게 어렵고 hallucination, out-of-distribution과 같은 경우들을 모두 핸들링하는 것에 어려움이 있다.

“activation engineering” 또는 “representation engineering”으로 불리는 정렬 방식은 모델의 activation을 변형한다. 해당 연구에서는 CAA라는 방식으로, high-level에서 alignment-relevant하게 모델을 조정하는 것을 선택한다. Llama 2 Chat은 안전이라는 목표 하에 RLHF을 사용했기에, 기본 Llama 2에 activation engineering을 적용한 것과 비교를 한다.

2. Related Work

  • 특정 layer와 토큰 위치에서의 activation의 차이를 사용한다. 이 방식은 다양한 동작이나 프롬프트에 대해 일관성이 없다. → CAA는 하나의 constrast pair가 아닌 contrast pair들이 모인 데이터셋으로 학습을 하여 노이즈도 감소하고, 토큰 위치만을 신경쓴다.
  • attention head들을 탐색하면서 가장 진실에 가까운 (truthful) 집합들을 찾고, 거짓에 가까운 분포와의 차이를 구해서 사용한다. → CAA는 각 head들을 탐색하지 않고도 더 넓은 범위에서 RLHF로 학습된 모델을 alignment-relevant하게 조정할 수 있다.
  • 정직도나 감정과 같은 high-level한 컨셉들에 대해 정보를 추출하는 시도도 존재한다 → CAA는 선택지가 여러 개인 포맷을 사용하여, 토큰 하나만 달라지는 조건에서 조정이 가능하다. 또한, 정보 추출보다는 steering에 집중한다.
  • 모델을 steering하는 시도도 존재한다. → 기존연구는 attention activation을 건드리는 반면, CAA는 residual stream에서 조정한다.

3. Method

CAA의 핵심 아이디어는 steering 벡터를 만들어서 원하는 방향으로 모델의 output 분포를 바꾸는 것이다. steering 벡터를 만드는 방법은 긍정적인 내용의 프롬프트 하나와 부정적인 내용의 프롬프트 하나씩 쌍을 지은 다음, activation의 차이를 구하는 것이다. 나머지 내용은 동일하되 원하는 행동(behavior)과 정반대인(e.g. sycophancy의 유무) 프롬프트의 activation의 차이를 구하면, 원하는 행동에 대한 부분만 추출할 수 있다.

AI 모델에서 정렬이 되는지 확인하는 task는 다음과 같다:

  • Coordination with Other AIs (다른 AI와의 협력)
  • Corrigibility (수정 가능성)
  • Hallucination (환각)
  • Myopic Reward (근시안적 보상)
  • Survival Instinct (생존 본능)
  • Sycophancy (아첨)
  • Refusal (거부)

behaviorial clustering은 약 3분의 1 정도가 진행되었을 때 드러나기 시작한다. 특정 layer에서 갑자기 구분이 명확하게 되는 경향성을 보인다. 예를 들어 refusal의 경우 9번째에서 10번째 layer로 넘어갈 때 separability가 높아지는 것을 확인할 수 있다.

4. Effect on CAA on behaviors

크게 2가지 유형에 대해서 실험한다. 첫번째는 선택지가 있는 multiple-choice question 데이터셋이고, 효과를 더 검증하기 위해서 자유 형태로 답할 수 있는 open-ended generation이다. steering vector를 넣을 때 곱해주는 multiplier에 대해서도 여러 값으로 실험을 하는데, 너무 커지게 되면 생성되는 결과의 퀄리티가 떨이지게 된다는 것을 보인다. 따라서 적절한 균형을 찾아 steering을 수행해야 한다.

5. CAA and system-prompting

LLM을 조정할 때 사용되는 방식 중 하나는 system 프롬프팅이다. LLM에게 역할을 부여하는 식으로 하는 방식인 system 프롬프팅은, 어떤 식의 행동을 표현해야 할지 명시적으로 지시할 수 있다. system 프롬프팅과 CAA를 비교했을 때, 모델의 행동을 더 크게 변화시킬 수 있다는 것을 볼 수 있다. 그 이유를 CAA 방식은 multiplier를 통해 steering을 더 미세한 부분까지 조정할 수 있고, 원하는 행동(behavior)만 추출해서 강화할 수 있다는 것으로 추측한다.

6. Comparison to finetuning

모델을 파인튜닝한 것과 CAA 방식을 비교하는 섹션이다. 파인튜닝을 함으로써 높은 성능을 보이지만, out-of-distribution에서 generalization이 필요한 경우에는 CAA가 더 높은 성능을 보인다. 파인튜닝한 모델에서 CAA 방식을 사용할 수도 있는데, 앞서 언급된 7개의 task 중에서 3개는 더 나은 성능을 보인다. 나머지 task에서는 비직관적이지만, 성능이 더 낮아지는 것이 보여진다. 또한, 파인튜닝을 할 경우에는 시간도 오래 걸리고 GPU도 많이 필요로 한다는 단점이 있다.

7. Effects of CAA on general capabilities

MMLU (Massive Multitask Language Understanding) 벤치마크에 대해서도 CAA를 적용한 뒤에 성능을 확인한다. CAA 방식을 적용해도 MMLU 점수에 영향을 끼치지 않는 것을 확인하고, sycophacy 관련해서 CAA를 적용했을 때 조금 나아지는 모습도 확인할 수 있다.

8. Understanding and interpreting CAA

8.1 Similarity between steering vectors and per-token activations

먼저 추출한 steering vector와 실제로 그런 행동을 가진 프롬프트와의 유사도를 비교해본다. cosine 유사도를 계산했을 때, 실제로 그러한 행동이 존재하는 토큰과 steering vector는 유사한 값을 가지는 것을 확인할 수 있다.

8.2 Similarity between vectors generated at different layers

근접한 layer에서 뽑은 벡터들은 유사도가 높지만, 거리가 멀어질수록 유사도는 줄어든다. 하지만 뒷 부분의 layer들에서는 유사도가 줄어드는 정도가 감소하는 것을 보아, 한번 high-level 정보를 추출하면 그 위치는 수렴하는 것으로 보인다. 그런데 관련 정보가 이미 사용하고 처리되면, 뒷 부분에서는 다른 정보로 대체될 수 있다는 것을 보인다.

8.3 Comparing representations between base and chat models

Llama 2 Chat 모델과 Base 모델을 비교했을 때 유사도는 layer가 진행될 수록 감소하지만, 전반적으로 유사한 양상은 RLHF가 정보가 표시되는 것에 대한 영향은 작다는 것을 보인다.

9. Discussion

CAA 방식은 steering 벡터를 generalize하여 실용적으로도 다재다능함을 보인다.

장점은 다음과 같다.

  • 기존의 모델 성능을 감소시키지 않는다.
  • 시스템 프롬프팅이나 파인튜닝과 같은 방식과도 호환 가능한 additive 성격을 가졌다.
  • 모델 내부에서 정보들이 어떻게 저장되는지 high-level representation을 볼 수 있다.

⇒ 결론적으로, 정확하교, 효과적으로 모델의 행동을 조정할 수 있다는 점에서, AI를 더 통제 가능하고 긍정적인 방향으로 쓰일 수 있게 한다.

추후에는 더 작은 단위에서 steering vector를 적용하여, 생성되는 결과의 품질이 낮아지는 것의 커트라인을 조금 완화시킬 수 있을 것이다. 어떻게 해야 부정적인 행동을 하는지 파악하여, red-teaming에도 활용할 수 있다.

10. Limitations

  • GPT-4는 내부적으로 scoring 과정이 있기 때문에 CAA 방식에 방해가 될 수 있다.
  • 파인튜닝을 실험할 때, hyperparameter를 더 열심히 조정하면 더 좋은 결과를 얻을 수 있다.
  • 프롬프팅과 관련해서는 모든 경우를 실험할 수 없다는 한계가 존재한다.
  • layer 단위에서 더 큰 범위에서의 norm이 있는데, 그 부분은 건드리지 않았다. 따라서 결과를 조금 왜곡 시켰을 수도 있다.

 

논문 링크 : 
https://arxiv.org/abs/2312.06681
728x90
반응형