[ROME] Locating and Editing Factual Associations in GPT

[bigshanedogg]

Problem statement

• transformer 구조에서 factoid query에 대한 실패 (i.e. hallucination)는 공통적으로 관측되고 있음.
  • 모델의 어느 뉴런이 factual prediction에 관여하는지 식별하는 방법론 제안
  • feed-forward wieghts를 업데이트하여 특정 factual associations을 업데이트하는 ROME (Rank-One Model Editing)을 제안

Glossary

• Knowledge Tuple: knowledge editing 분야에서 (s, r, o)로 이루어진 statement를 의미
  • e.g. (The Space Needle, is located in, the city of Seatle)
  • c.f. knowledge editing 분야에서 prompt는 보통 (s, r)을 의미한다.
• TE (Total Effect): 정답 토큰에 대한 clean run과 corrupted run 사이의 prob 차이
• IE (Indirect Effect): 정답 토큰에 대한 corrupted-with-restoration run과 corrupted run 사이의 prob 차이, 즉, corrupted run에서 특정 weights에 의한 영향도
  • AIE (Average Indirect Effect): 모든 샘플들에 대한 IE의 평균

Baseline

• KE (Knowledge Editor)
• MEND (Model Editor Networks with Gradient Decomposition)

Data details

name	abbr	type	format	source	size	description	remark	related tasks
zero-shot Relation Extraction	zsRE	(s,r,o)	-	-	-	-	-	KnowledgeEditing
CounterFACT	-	(s,r,o)	-	-	-	-	-	KnowledgeEditing

Approach

• Causal Traces: the causal effects of hidden state activations

  • 방법:
  • 
  • clean run: 정상적인 forward로 각 layer의 모듈들 (attention & mlp)의 clean hidden states를 저장, 이 때 output token은 정답을 생성
  • corrupted run: subject tokens의 embedding에 noise를 더해서 corrupted hidden states로 바꿈, 이 때 output token은 오답을 생성
  • corrupted-with-restoration run: 일부 corrupted hidden states를 clean hidden states로 변경, 만약 clean hidden states로 변경했을 때 정답을 생성한다면, 해당 부분이 factual association에 영향을 주는 부분이라고 가정할 수 있음.
  • Findings:
  • (trivial) A. object token의 생성은 마지막 레이어의 마지막 시퀀스 부분에서 강한 인과성이 있다.
  • (notable) B. 중간 레이어부가 last subject token에 관여하며, 이 부분에도 강한 인과성이 있다.
  • MLP와 attention 중 상대적으로 인과성이 강한 모듈은 A는 attention module, B는 MLP module이다. 즉, subject에 대한 fact recall에 관여하는 key–value mapping이 일어나는 부분을 midlayer MLP라고 가정할 수 있다.
    • B의 경우, Attention 모듈의 AIE는 1.6%, MLP 모듈의 AIE는 6.6%
  • 실험 인사이트:
  • 위의 분석 결과와 아래의 관련 연구를 통해, '어떤 사실이라도 중간 레이어의 MLP 모듈 중 하나에 동등하게 저장되어 있을 것이다'라고 추측하고, middle layer의 MLP 모듈 weights의 수정을 통해 자의적인 사실을 저장하도록 변경할 수 있는지 실험해본다.
    • Transformer feed-forward layers are key-value memories (EMNLP 2021): MLP layers store knowledge, first layer form a key, while second layer retrieves an associated value
    • A mathematical framework for transformer circuits: an information-copying role for self-attention
    • Of non-linearity and commutativity in BERT (IJCNN 2021): transformer layer order can be exchanged with minimal change in behavior

• Rank-One Model Editing: 사실 연관관계 저장을 이해하기 위한 weights intervention

  • 전제:
  • MLP layer는 WK ≈ V 을 최적화하는 것을 통해 k-v mapping을 학습한다.
  • 아래 수식을 통해 새로운 k-v pair (k*, v*)를 k-v memory에 삽입할 수 있다.
    • 
    • 대략적으로 이해하자면, 우선 C는 본래의 key인 k의 무중심공분산에 대한 추정 상수인데, 새로 삽입하게 될 k*를 분포를 원점을 기준으로 중심화된 분포가 아닌 아닌 원래의 데이터인 k와 유사하게 만들어주기 위함, 즉, 원점을 기준으로 상대적 위치가 아닌 k-v mapping 상에서 절대적 위치로 변경하여 주는 것 (이라고 이해했다...)
    • 잔여 오류 Λ를 더하는 이유는 W의 특정 열에 (k, v) 관련된 업데이트를 한 이후에도 전체 행렬인 W가 Wk ∧ ≈V를 더 잘 만족시키도록 하기 위함, 즉, 기존의 다른 fact mapping에 주는 영향도를 최소화하기 위함 (이라고 이해했다...)
  • 결론적으로, 적절한 (k, v)를 선택하기만 하면, weights에 직접적으로 사실 관련 정보를 삽입할 수 있다.
  • 방법:
  • (step1) Choosing k∗ to Select the Subject:
    • k를 선택하기 전에 먼저 k를 추출하기 위한 input token을 먼저 선택해야 함.
    • 앞선 분석결과를 토대로 중간부 레이어들을 l, last subject token의 인덱스를 i라 할 때, `(l, i)`의 hidden states를 모은다. 이 때, 모으는 모듈의 위치는 MLP에서 non-linearity function (i.e. activation function) 직 후의 값이다.
    • hidden states는 prompt prefix에도 영향을 받기 때문에, 2~10 길이의 랜덤한 텍스트를 붙이고, 50개의 랜덤한 프롬프트에 대한 값을 평균한다.
  • (step2) Choosing v∗ to Recall the Fact:
    • 
    • (이해 중...)
  • (step3) Inserting the Fact:
    • 앞선 계산을 이용해 새로운 knowledge triplet인 (s, r, o*)에 해당하는 (k*, v*)을 W의 weights에 직접 변경한다.

Evaluation

• 위 방법을 이용해 삽입한 지식이 유의미하고, 일반화된 사실 관계로 모델 내에서 작동할 수 있는가 를 알아본다.

• zsRE

  • 각 샘플은 1개의 factual statement와 its paraphrase, 그리고 1개의 unrelated factual statement로 이루어진다.
  • factual statement와 its paraphrase를 통해 방법론이 얼마나 일반화된 사실 관계로 작용하는지 (Efficacy), unrelated factual statement를 통해 다른 지식에 영향을 주진 않았는지 (Specificity) 확인할 수 있다.
  • 
  • baselines:
  • FT: finetuning
  • FT+L: fnetuning with norm constraint on weight changes
  • previous approaches such as KE, MEND: 2가지 모델 모두 weight changes를 예측하기 위한 별도의 모델을 필요로하는 hypernetwork 구조임
  • efficacy는 매우 강한 편이고, paraphrase 점수는 조금 낮지만, zsRE 데이터 분포에 대한 학습이 필요한 KE와 MEND를 고려하면 competitive하다.
  • specificity는 높은 편이지만, 사실 zsRE가 굉장히 넓은 범위의 사실 관계에 대한 데이터셋이고 거기서 샘플링한 unrelated sample이기 때문에 이 메트릭의 의미는 크게 높지 않은 것 같다.
  • 
  • 레이어의 위치, 토큰 인덱스 i에 대한 결과 비교. 종합적인 스코어를 두고 봤을 때, 결과적으로도 중간부 layer의 last subject token의 weight를 타겟으로 하는 것이 효과가 있다.

• CounterFACT

  • metrics:

Limitations

• 한 번에 1개의 사실 관계만을 수정할 수 있기 때문에 대규모 모델 학습이라는 관점에서 실용적이지 않음.
• (s, r, o)의 구조가 다른 2개의 statements는 동일한 사실을 가리키고 있어도 별도로 2번의 edits을 적용해야 한다.
  • e.g. The iconic landmark in Seattle is the Space Needle & The Space Needle is the iconic landmark in Seattle
• 단순한 사실 관계 외의 logical, spatial, or numerical knowledge에 대해서는 다루지 못하고 있다.