Review

TL; DR

Summary

• 연구진: 알리바바, 중국과학기술대학, 칭화대학교, 난양이공대학교

• github: https://github.com/ydyjya/SafetyHeadAttribution

Main Idea

standard attention mechanism 와 safety capability 간의 관련성을 찾아, safety에 관한 interpretability를 탐구하자 !

Background & Motivation

• LLM의 safety

  

  • harmful query에 대해 답변을 거절하도록 alignment 되어 있음 (그림 왼쪽)
    • e.g. ‘I cannot’ or ‘As a responsible AI assistant’ 등의 rejection token 사용
  • but, 특정 token의 확률분포를 조정하면 (Jailbreak Attack) safety에 취약해져서 harmful query에도 답변하게 됨 (그림 오른쪽)
    • ‘I cannot’, ‘As a responsible AI assistant’ 등의 rejection token을 낮추거나
    • “Sure”, “Here is…” 등의 affirmative tokens을 높히는 것

• 기존 LLM safety 관련 논문은 주로 features, neurons, layers, parameters 관점에서 수행됨
  • e.g. 어떤 neuron이 safety를 담당하는지

⇒ multi-head attention (MHA) 관점에서 safety를 분석해보자!
즉, safety에 가장 영향이 큰 head(=safety parameters)를 정량적으로 찾아보자!

Contributions (What they’ve revealed)

• ⚙️ Settings

  • head ablation에 용이하도록 수정된 modified multi-head attention 사용

    

    • annotations

      • $W_q, W_k, W_v$: Query, Key, and Value matrices

      • $h_i$: i번째 attention head
    • 직관적으로 head ablation == 해당 head의 output을 0으로 두는 것이지만, 논문에서는 두가지 방법을 사용해 head ablation 수행
      • Undifferentiated Attention: Q 혹은 K(또는 둘 다)에 아주 작은 계수 ε를 곱해, 해당 head의 attention weight가 모든 토큰에 대해 거의 균일(평균)하게 되도록 만듦

        

        ⇒ 즉, 특정 head에서 어떤 token이 중요한지 판단하는 token selection 기능을 저하시킴

        

        이 그림에서, 모든 token이 다 비슷한 색이 됨 !

      • Scaling Contribution: attention weights는 그대로 두고, head output V를 ε로 줄여서, 이 head의 contribution만 매우 약하게 만드는 방식

        

        ⇒ 즉, 특정 head의 output을 줄여버림

  • backbone LLM: Llama-2-7b-chat, Vicuna-7b-v1.5

  • dataset: Advbench, Jailbreakbench, Malicious Instruct

  • metric: attack success rate (ASR)

    

    : 높을수록 unsafety

  • LLM input setting
    • template: template에 질문을 넣어 답변을 생성하는 방식

      

    • direct: template 없이 질문만 넣기

  • decoding setting
    • vanilla: original decoding
    • greedy: greedy decoding
    • top5: top-5 sampling decoding

• 1. Safety interpretability 연구를 위해, 처음으로 Safety-specific attention head의 존재를 밝힘

  : 두가지 방법 (undifferentiated attention, scaling contribution)으로 safety-specific attention head (즉 highest ship score를 가지는 head)를 제거해 본 결과

  

  • Llama2-7b-chat: layer 2 head 26 제거 했더니 ASR이 평균 0.04에서 0.64↑ (16배)

  • Vicuna-7b-1.5: 일부 head 제거했더니 ASR이 평균 0.27에서 0.55↑ (2배)

  • Undifferentiated Attention 방식이 Scaling Contribution보다 safety를 훨씬 크게 약화시킴.

• 2. attention head의 safety impact를 보이기 위해 SHIP (Safety Head ImPortant) Score & SAHARA (Safety Attention Head Attribution Algorithm)제안

  • SHIP score란?

    

    • 각 attention head가 safety에 얼마나 중요한지를 나타내는 점수
    • annotations

      • qH: 거절해야 하는 harmful query

      • $θ_{h_i^l}$ : target attention head

      • $θ_O$: original model parameters

      • $D_{KL}$: Kullback-Leibler divergence

      • $θ_O\backslash θ_{h_i^l}$ : original model에서 target attention head를 제거한 모델

    ⇒ safety-critical attention head를 찾기 위해 generalized된 버전의 SHIP score

    즉, dataset 레벨로 attention head ablation 수행

    

    • how to make?
      1. harmful query dataset $Q_H$의 top layer activation a를 쌓아 matrix M 구성
      2. Singular Value Decomposition $SVD(M) = UΣV^T$을 통해 left singular matrix $U_θ$ 얻기
      3. attention head $h_i^l$이 ablation된 model에서 $U_A$ 얻기

      ⇒ 2와 3 사이의 principal angle이 클수록 safety에 관여가 큼에 착안함

    • annotations

      • $σ_r$: r-th singular value

  • SAHARA란?
    • safety degration을 유도하는 head group을 찾기 위한 heuristic algorithm
      • 실험을 통해 head group size < 5 로 설정

        

        ⇒ 이전 연구보다 더 fine-grained & efficient

      • 알고리즘 요약: 각 group을 제거해보고 ASR 성능 확인해봄

        

  • 결과

    

    • Undifferentiated Attention > Scaling Contribution
    • SAHARA를 사용하여 찾은 attention head를 제거했더니 ASR 향상됨

    

    • safety 기능은 모델 전체에 골고루 있는 게 아니라 특정 layer·head에 집중됨

• 3. LLM safety 관점에서 standard multi-head attention mechanism의 중요성을 분석함으로서, LLM risk에 대한 우려를 완화하고, transperency에 기여

  • vicuna와 llama에서의 safety head는 얼마나 겹칠까?

    

    ⇒ Safety head들은 pre-training 단계에서 이미 상당량 형성되어 있고, 이후 chat-tuning(Vicuna-style instruction-tuning)에서도 그대로 유지되는 것

  • helpful-harmness trade-off는?

    

    • safety head ablation한 Llama2-7b-chat에서 zero-shot 실험 결과,

      ASR은 크게 증가(안전성 붕괴)하지만, 일반적인 pruning(SparseGPT, Wanda)와 비교했을 때 zero-shot 성능 저하는 비슷하거나 더 적음.

    ⇒ 즉, safety head는 “주로 safety용”으로 기능하고, 일반 언어능력과의 중첩(superposition)이 상대적으로 적은 것으로 보임.