기술 요약: 임베딩을 잊지 마세요: 임베딩의 정밀한 편집을 통한 견고한 지식 삭제

문제 정의

언어 모델(LM)이 점점 더 많이 배포됨에 따라, 특정 지식(예: 개인 데이터, 저작권 콘텐츠 또는 안전 민감 주제)을 영구적으로 삭제하는 능력은 매우 중요합니다. 기존의 지속적인 삭제 방법들은 주로 모델의 파라미터, 특히 다층 퍼셉트론(MLP) 레이어를 업데이트하는 데 집중하며, 종종 희소 오토인코더(SAE)나 가중치 스티어링(weight steering)을 활용합니다. 그러나 이러한 방법들은 종-종 견고한 삭제를 달성하는 데 실패합니다. 대상 지식은 적대적 프롬프팅이나 재학습(개념 관련 데이터에 대한 미세 조정)을 통해 여전히 복구될 수 있는 경우가 많으며, 삭제 효능은 형태에 따라 의존적입니다(객관식에서 개방형 생성으로 일반화하는 데 실패함).

저자들은 이러한 한계가 기존 기술들이 토큰 임베딩 레이어를 간과하고 있기 때문이라고 가설을 세웁니다. MLP가 지식 회상에 핵심적인 역할을 한다는 것은 알려져 있지만, LM은 토큰 임베딩에 직접 상당한 양의 지식을 인코딩하기도 합니다. 이 레이어를 무시하면 모델이 삭제된 개념을 "재학습"할 수 있는 경로가 남게 되며, 기존 방법들이 모델의 전반적인 유용성을 저해하는 공격적이고 광범위한 파라미터 업데이트에 의존하게 만듭니다.

방법론: EMBER

본 논문은 기존의 MLP 중심 삭제 방법을 보완하기 위해 토큰 임베딩을 정밀하게 편집하도록 설계된 플러그 앤 플레이 모듈인 EMBER(EMBedding ERasure)를 소개합니다. EMBER는 다음 세 단계로 구성된 "국소화 후 편집(localize-then-edit)" 접근 방식을 따릅니다:

1. 특징 국소화를 위한 희소 행렬 분해(Sparse Matrix Factorization, Sparse MF):
   SAE를 사용하는 대신, EMBER는 임베딩 행렬 $E$에 직접 희소 행렬 분해를 적용합니다. 이는 특정 토큰 집합(개념 관련 문장 및 중립 문장에 등장하는 토큰)의 임베딩을 요인 사전(factor dictionary) $Z$(임베딩 특징)와 계수 행렬 $Y$로 분해합니다.

   • 희소성 제약(Sparsity Constraint): 특징이 국소화되고 해석 가능한 개념에 대응하도록 하기 위해, 계수 $Y$에 "승자 독식(winner-takes-all, WTA)" 제약을 적용하여 각 요인당 크기가 상위 1%인 것만 유지합니다. 이는 분해가 작고 구체적인 토큰의 하위 집합에 의해 공유되는 특징을 식별하도록 강제합니다.

2. 개념 관련 특징 식별:
   이 방법은 회복된 특징들을 점수화하여 개념 특화 특징과 일반 특징을 구분합니다. 개념 라벨이 붙은 토큰에 대한 특징의 기여도와 중립 토큰에 대한 기여도를 비교하는 질량비(mass-ratio) 통계량을 계산합니다. 높은 비율을 가진 특징은 LLM-as-a-judge를 사용하여 관련 토큰이 대상 개념과 의미적으로 일치하는지 검증함으로써 추가로 필터링됩니다.

3. 차감을 통한 정밀 삭제:
   개념 관련 특징($F_C$)이 식별되면, EMBER는 관련 토큰의 임베딩에서 해당 기여도를 뺍니다. 토큰 임베딩 $e_t$에 대한 업데이트는 다음과 같습니다:
   $$e_t \leftarrow e_t - \delta \sum_{i \in F_C} Y_{i,t} z_i$$
   여기서 $\delta$는 강도 하이퍼파라미터입니다. 이 연산은 나머지 임베딩은 온전하게 유지하면서 개념 특화 성분만을 제거합니다. 이 방법은 식별된 개념 특징에 대해 0이 아닌 계수를 가진 토큰에 대해서만 적용됩니다.

EMBER는 RMU, CRISP, SNMF와 같은 기존의 MLP 기반 삭제 방법과 결합되어 임베딩 레이어와 MLP 레이어를 모두 타겟팅하는 앙상블을 구축하도록 설계되었습니다.

주요 기여

• 임베딩 병목 현상의 식별: 본 논문은 임베딩 레이어가 견고한 개념 삭제의 결정적인 병목 구간임을 입증하는 첫 번째 실증적 연구를 제공합니다. 임베딩을 무시하면 모델이 삭제된 지식을 더 쉽게 복구할 수 있음을 보여줍니다.
• EMBER 알고리즘: 희소 행렬 분해를 사용하여 전체 모델 미세 조정 없이 토큰 임베딩에서 개념 관련 특징을 분리하고 정밀하게 제거하는 새로운 방법입니다.
• 플러그 앤 플레이 증강: EMBER는 기존 삭제 프레임워크와 호환되도록 설계되었으며, 삭제 파이프라인의 완전한 재설계 없이도 성능을 크게 향상시킵니다.

실험 결과

저자들은 Gemma-2-2B-it와 Llama-3.1-8B-Instruct를 대상으로 18가지 다양한 개념(해리 포터와 같은 허구적 작품부터 대마초와 같은 안전 주제까지)에 대해 EMBER를 평가했습니다.

• 효능 및 특이성: 기존 방법(RMU, CRISP, SNMF)에 EMBER를 증강했을 때, 삭제 효능(개념 정확도 감소)과 특이성(유사 도메인 및 MMLU에서의 성능 유지)이 일관되게 향상되었습니다.
• 재학습에 대한 견고성: 이것이 주요 이점입니다. EMBER는 기존 방법들의 재학습 능력을 극적으로 감소시켰습니다.
  • Llama-3.1에서, 기존 방법(RMU, CRISP)의 재학습 정확도는 EMBER와 결합했을 때 **70%–76%**에서 **35%**로 감소했습니다(약 50% 감소).
  • Gemma에서, 재학습 정확도는 16%–24% 감소했습니다.
  • 전체 MF 기반 파이프라인(SNMF + EMBER)은 가장 낮은 재학습 정확도(Llama에서 35%)를 달성하여 베이스라인을 크게 앞질렀습니다.
• 일관성 및 국소화: 생성 일관성에 미치는 비용은 미미했습니다. 분석 결과, EMBER는 개념당 어휘의 최대 **0.14%**만을 수정했습니다. 일관성 저하는 "개념 전용" 토큰(예: 해리 포터를 삭제할 때의 "Hogwarts")에 매우 국소적으로 나타났으며, 일반적인 토큰(예: 프린스 해리의 맥락에서의 "Harry")은 거의 영향을 받지 않았습니다.
• 일반화: 하이퍼파라미터는 객관식(MC) 작업에 대해 튜닝되었지만, 개선 사항이 개방형(OE) 질의응답으로 효과적으로 전이되었습니다. 이는 편집이 형식 특화된 지름길이 아니라 개념적 지식을 타겟팅함을 시사합니다.

의의 및 주장

본 논문은 정밀하고 타겟팅된 임베딩 행렬 편집이 견고한 개념 삭제를 달성하는 데 근본적이라고 주장합니다. 기존 연구의 MLP 중심 뷰를 넘어 임베딩 레이어로 초점을 전환함으로써, EMBER는 지식이 단순히 억제될 뿐 제거되지 않는 "얕은 언러닝(shallow unlearning)" 문제를 해결합니다.

저자들은 다음과 같이 결론짓습니다:

1. 재학습 공격을 방지하기 위해서는 임베딩 수준의 개입이 필수적입니다.
2. 기존 삭제 방법은 임베딩 수준의 편집을 증강함으로써 크게 향상될 수 있습니다.
3. 광범위한 가중치 업데이트를 적용하는 대신 특정 개념 관련 특징에 편집을 국소화함으로써, 삭제 효능과 모델 유용성(일관성) 사이의 트레이드오프를 개선할 수 있습니다.

이 연구는 모델 편집의 새로운 방향을 제시하며, 임베딩 레이어가 안전 및 컴플라이언스 애플리케이션을 위해 반드시 다루어야 할, 중요하지만 이전에 충분히 탐구되지 않은 지식 저장소임을 강조합니다.