[논문 리뷰] PMET: Precise Model Editing in a Transformer

Xiaopeng Li, Shasha Li외 4명|arXiv (Cornell University)|2023. 08. 17.

Topic Modeling인용 수 6

한 줄 요약

PMET는 Transformer에서 다중 헤드 자기주의(MHSA)와 피드포워드 네트워크(FFN)의 히든 상태를 동시에 최적화하면서, FFN 가중치 업데이트에만 최적화된 FFN 히든 상태를 사용함으로써 정밀한 모델 편집 방법을 제안한다. 이 방법은 MHSA의 지식 추출 기능을 수정하지 않고 활용함으로써, counterfact 및 zsRE 벤치마크에서 최신 기술 수준의 성능을 달성하며 편집의 신뢰성과 특이성을 향상시킨다.

ABSTRACT

Model editing techniques modify a minor proportion of knowledge in Large Language Models (LLMs) at a relatively low cost, which have demonstrated notable success. Existing methods assume Transformer Layer (TL) hidden states are values of key-value memories of the Feed-Forward Network (FFN). They usually optimize the TL hidden states to memorize target knowledge and use it to update the weights of the FFN in LLMs. However, the information flow of TL hidden states comes from three parts: Multi-Head Self-Attention (MHSA), FFN, and residual connections. Existing methods neglect the fact that the TL hidden states contains information not specifically required for FFN. Consequently, the performance of model editing decreases. To achieve more precise model editing, we analyze hidden states of MHSA and FFN, finding that MHSA encodes certain general knowledge extraction patterns. This implies that MHSA weights do not require updating when new knowledge is introduced. Based on above findings, we introduce PMET, which simultaneously optimizes Transformer Component (TC, namely MHSA and FFN) hidden states, while only using the optimized TC hidden states of FFN to precisely update FFN weights. Our experiments demonstrate that PMET exhibits state-of-the-art performance on both the COUNTERFACT and zsRE datasets. Our ablation experiments substantiate the effectiveness of our enhancements, further reinforcing the finding that the MHSA encodes certain general knowledge extraction patterns and indicating its storage of a small amount of factual knowledge. Our code is available at https://github.com/xpq-tech/PMET.

연구 동기 및 목표

• MHSA와 잔차 연결에서 유도되는 관련 없는 정보를 포함한 Transformer 레이어(TL) 히든 상태에 의존하는 기존 모델 편집 방법에서의 부정확한 가중치 업데이트 문제를 해결하기 위해.
• MHSA가 사실 지식을 저장하는지, 아니면 오직 일반적인 지식 추출 패턴만 저장하는지 조사하여, 그 가중치가 업데이트되어야 할지 여부에 영향을 미칠 수 있는지를 밝히기 위해.
• MHSA와 FFN 히든 상태 최적화를 분리함으로써 모델 편집의 신뢰성과 특이성을 향상시키고, MHSA의 구조적 역할를 유지하기 위해.
• MHSA 가중치를 수정하지 않고도 편집 효율성, 일반화 능력, 특이성에서 최신 기술 수준의 성능을 달성하기 위해.
• 구성 요소 수준의 히든 상태 동역학을 분석함으로써 더 견고하고 해석 가능한 LLM 편집 프레임워크를 제공하기 위해.

제안 방법

• PMET는 Transformer 레이어 내에서 MHSA와 FFN 구성 요소의 히든 상태를 동시에 최적화하며, 이를 별개의 최적화 목표로 간주한다.
• 오직 최적화된 FFN 히든 상태만을 타겟 지식 표현으로 사용하여, 오염된 MHSA 또는 잔차 정보로부터 벗어나 정밀한 최소 제곱 기반의 FFN 가중치 업데이트를 수행한다.
• MHSA 히든 상태는 FFN 최적화의 기능 공간을 확장하기 위해 최적화되며, MHSA 가중치는 업데이트하지 않으며, MHSA가 일반적인 지식 추출 패턴을 인코딩한다는 통찰에 기반한다.
• 주어진 증분 가중치 업데이트를 핵심 레이어에 균형 있게 분배하기 위해 잔차 확산 전략(제곱근 확산)을 활용하여, 신뢰성과 모델 보존 능력 간의 균형을 확보한다.
• MHSA와 FFN 히든 상태를 동시에 최적화하면서도 잔차 연결 흐름의 무결성을 유지하는 새로운 손실 함수를 도입한다.
• MHSA 가중치 수정을 피하고, 오직 정제된 최적화된 표현을 통해 FFN 가중치만 업데이트함으로써, MHSA의 지식 추출 기능을 유지한다.

실험 결과

연구 질문

• RQ1 MHSA는 사실 지식을 저장하는가, 아니면 오직 안정적인 편집 간에도 유지되는 일반적인 지식 추출 패턴만 저장하는가?
• RQ2 MHSA 가중치를 업데이트하지 않고도 MHSA 히든 상태를 최적화함으로써, FFN 가중치 업데이트의 정밀성과 신뢰성을 향상시킬 수 있는가?
• RQ3 잔차 확산 전략의 선택(예: 제곱근 대비 균일)이 편집 신뢰성과 모델 특이성 간의 트레이드오프에 어떤 영향을 미치는가?
• RQ4 전체 TL 히든 상태를 사용하는 방법과 비교해, MHSA와 FFN 히든 상태 최적화를 분리함으로써 편집 성능이 얼마나 향상되는가?
• RQ5 혼합 또는 잔차 영향을 받는 표현 대비, 오직 최적화된 FFN 히든 상태만을 사용해 가중치를 업데이트할 경우의 영향은 무엇인가?

주요 결과

• PMET는 counterfact 및 zsRE 벤치마크에서 모두 최신 기술 수준의 성능을 달성하며, 기존 방법 대비 편집의 신뢰성과 특이성에서 뛰어난 성능을 보였다.
• 제거 실험 결과, 오직 FFN 히든 상태 최적화만 하는 것보다 MHSA와 FFN 히든 상태를 모두 최적화함으로써 더 높은 신뢰성이 확보되었으며, 기능 공간의 확장 효과를 확인했다.
• MHSA 가중치를 FFN 가중치와 함께 업데이트하면 일반화 능력은 약간 향상되지만 특이성은 크게 악화되어, MHSA가 일부 사실 지식을 저장하고 있음을 시사하며, 이는 MHSA 가중치를 업데이트해서는 안 된다는 것을 뒷받침한다.
• 균일한 확산 전략은 모델 유지 능력(유창성 및 특이성)을 향상시키지만, 효율성과 일관성은 극적으로 감소시켜 정보 손실의 위험을 드러낸다.
• 제곱근 확산 전략은 모든 핵심 지표에서 균일한 확산 전략을 능가하며, 모델 무결성 유지와 업데이트의 신뢰성 간 최적의 균형을 제공한다.
• 결과는 MHSA가 일반적인 지식 추출 패턴을 인코딩하고 있으며, 소량의 사실 지식도 저장하고 있음을 확인하며, 이는 PMET에서 MHSA 가중치를 업데이트하지 않는 것이 타당하다는 것을 정당화한다.