HiEdit: Lifelong Model Editing with Hierarchical Reinforcement Learning

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 은 배포된 후 시간이 지남에 따라 지식이 노후화되거나 부정확해질 수 있습니다. 이를 해결하기 위해 모델을 처음부터 재학습 (Retraining) 하는 것은 비용이 너무 많이 들기 때문에, 생애주기 모델 편집 (Lifelong Model Editing, LME) 이 중요한 과제로 대두되었습니다. LME 는 기존 지식을 손상시키지 않으면서 새로운 지식을 순차적으로 주입하는 것을 목표로 합니다.

그러나 기존 접근법들은 다음과 같은 치명적인 한계를 가지고 있습니다:

• 정적 및 밀집된 레이어 편집: 모든 편집 작업에 대해 모델의 특정 레이어 집합 (예: 모든 MLP 레이어) 을 무조건적으로 수정합니다.
• 지식별 레이어 특이성 무시: 연구에 따르면 서로 다른 지식은 모델 내의 서로 다른 레이어에 저장됩니다. 모든 편집에 동일한 레이어를 강제로 수정하면 불필요한 파라미터 교란이 발생하여, 새로운 지식의 통합을 방해하고 기존 지식 (일반적 지식 및 이전에 편집된 지식) 의 파괴적 망각 (Catastrophic Forgetting) 을 유발합니다.
• 성능 저하: 긴 순차 편집 시나리오 (수만 건의 편집) 에서 기존 방법들은 일반 성능과 편집된 지식의 유지율이 급격히 떨어집니다.

2. 방법론 (Methodology)

저자들은 이러한 문제를 해결하기 위해 HiEdit을 제안했습니다. HiEdit 은 계층적 강화 학습 (Hierarchical Reinforcement Learning, HRL) 을 활용하여 각 편집 인스턴스 (Knowledge Instance) 에 대해 가장 관련성이 높은 레이어를 적응적으로 (Adaptively) 선택하고 국소적으로 업데이트하는 프레임워크입니다.

핵심 구성 요소

1. 계층적 의사결정 구조 (Hierarchical Decision-Making):

   • LME 작업을 두 단계의 하위 작업으로 분해합니다:
     • 고수준 (High-level): "어디를 편집할 것인가?" (Layer Selection). 어떤 레이어를 편집 대상으로 선택할지 결정합니다.
     • 저수준 (Low-level): "어떻게 편집할 것인가?" (Parameter Updating). 선택된 레이어의 파라미터를 어떻게 업데이트할지 결정합니다.
   • 이를 통해 평탄하고 복잡한 행동 공간 (Action Space) 을 관리 가능한 구조화된 프레임워크로 변환합니다.

2. 하이퍼네트워크 (Hypernetworks) 설계:

   • 고수준 하이퍼네트워크 (Manager, $\pi_\phi$): 편집할 지식 $(x_t, y_t)$ 와 현재 모델 파라미터를 입력받아, 중요도 분포를 기반으로 레이어 선택 마스크 (Layer Selection Mask) 를 생성합니다. Top-K 알고리즘을 통해 가장 중요한 $K$ 개의 레이어만 활성화합니다.
   • 저수준 하이퍼네트워크 (Worker, $H_\theta$): 고수준에서 선택된 레이어에 대해 실제 파라미터 업데이트 ($\tilde{\nabla}W$) 를 생성합니다. 기존 MEND 또는 MALMEN 아키텍처를 기반으로 합니다.

3. 내재적 보상 메커니즘 (Intrinsic Reward Mechanism):

   • 희소성 (Sparsity) 을 장려하기 위해 부분 레이어 업데이트 vs 전체 레이어 업데이트의 상대적 이점을 측정하는 내재적 보상을 도입합니다.
   • 고수준 보상 ($r_{high}$) 은 전체 레이어를 업데이트했을 때의 손실과 선택된 부분 레이어를 업데이트했을 때의 손실 차이를 기반으로 계산됩니다. 이는 불필요한 파라미터 수정을 최소화하면서 편집 목표를 달성하도록 모델을 유도합니다.

4. 학습 과정:

   • 고수준과 저수준 정책을 연결하여 엔드 - 투 - 엔드 (End-to-End) 학습이 가능하도록 합니다.
   • 강화 학습 (RL) 프레임워크를 사용하여 편집 궤적 (Trajectory) 전반에 걸친 장기적 의존성을 포착하고, 누적 보상을 최대화하도록 정책을 최적화합니다.

3. 주요 기여 (Key Contributions)

1. HiEdit 프레임워크 제안: LME 의 복잡한 행동 공간을 계층적 구조로 재구성하여, 각 지식 인스턴스에 맞춰 적응적이고 국소적인 업데이트를 수행하는 새로운 접근법을 제시했습니다.
2. 내재적 보상 메커니즘: 부분 레이어 업데이트의 상대적 이점을 기반으로 한 보상 함수를 설계하여, 불필요한 파라미터 수정을 줄이고 편집 효율성을 극대화했습니다.
3. 엄격한 실험 설정 및 성과: 기존 연구에서 다루지 않았던 수만 건의 긴 순차 편집 (Long-range Sequential Editing) 시나리오를 최초로 탐구했습니다.
   • 경쟁 모델인 RLEdit 대비 평균 8.48% 의 성능 향상.
   • 편집 시 파라미터 교란을 절반 (약 50%) 으로 줄이면서도 더 높은 정확도를 달성했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 모델: ZsRE, CounterFact 데이터셋을 사용하여 Llama-3-8B 와 Gemma-2-9B 모델에서 실험 수행.
• 성능 비교:
  • 효율성 (Efficacy), 일반화 (Generalization), 특이성 (Specificity), 편집 유지율 (Edited Retention) 모든 지표에서 기존 최첨단 방법 (RLEdit, AlphaEdit 등) 을 압도적으로 상회했습니다.
  • 특히 20,000 번의 순차 편집 후에도 HiEdit 은 일반 능력 (GLUE 벤치마크) 과 이전에 편집된 지식을 거의 완벽하게 유지하는 반면, 기존 방법들은 5,000 번 편집 이후부터 성능이 급격히 하락하거나 모델 붕괴가 발생했습니다.
• 해석 가능성 (Interpretability):
  • HiEdit 은 지식의 종류 (예: 음악 vs 질병) 에 따라 서로 다른 레이어를 선택하는 것을 확인했습니다. 이는 모델이 지식의 특성에 따라 최적의 편집 위치를 동적으로 찾음을 시사합니다.
• 계산 비용: 고수준 하이퍼네트워크를 추가하더라도 전체 파라미터 증가분은 LLM 의 약 1.96% 에 불과하여 확장성 (Scalability) 이 뛰어납니다. 또한, 편집할 레이어 수를 줄여 편집 시간을 단축했습니다.

5. 의의 및 중요성 (Significance)

• 효율적인 생애주기 학습: LLM 의 지속적인 업데이트가 필요한 현실적인 요구사항에 부합하며, 재학습 없이도 수만 번의 지식 변경을 안정적으로 처리할 수 있는 가능성을 입증했습니다.
• 파괴적 망각 해결: 기존 방법들의 가장 큰 약점인 '과도한 파라미터 수정'으로 인한 기존 지식 손실 문제를, 적응적 레이어 선택을 통해 효과적으로 해결했습니다.
• 새로운 패러다임 제시: 모델 편집을 단순한 파라미터 보정이 아닌, 계층적 강화 학습을 통한 지능적인 의사결정 과정으로 재정의함으로써, 향후 모델 편집 및 지속적 학습 연구에 중요한 방향성을 제시했습니다.

결론적으로, HiEdit 은 LLM 의 생애주기 동안 발생하는 지식 업데이트 문제를 해결하기 위해, 어디를 (Which layers) 그리고 어떻게 (How) 편집할지를 학습 가능한 방식으로 통합한 혁신적인 솔루션입니다.