DualOptim+: Bridging Shared and Decoupled Optimizer States for Better Machine Unlearning in Large Language Models

DualOptim+ 기술 요약

문제 정의
대규모 언어 모델 (LLM) 에서의 기계적 망각 (Machine Unlearning, MU) 은 특정 학습 데이터 (망각 세트) 의 영향을 제거하면서도 나머지 데이터 (유지 세트) 에 대한 모델의 일반적 유용성을 보존하는 것을 목표로 합니다. 기존 접근법들은 이러한 상충되는 목표들을 균형 있게 달성하는 데 상당한 어려움을 겪고 있습니다.

• 연합 업데이트 (Joint Updating): 대부분의 방법은 망각 손실과 유지 손실의 합을 동시에 최소화합니다. 이는 과도한 망각으로 인해 모델 유용성이 크게 저하되는 결과를 초래합니다.
• 교대 업데이트 (Alternate Updating): 망각과 유지 목표를 교대로 업데이트하는 것은 성능을 개선하지만, 불안정성과 하이퍼파라미터에 대한 민감성으로 고통받습니다.
• 분리된 옵티마이저 (DualOptim): 이전 접근법인 DualOptim 은 각 목표에 대한 옵티마이저 상태 (예: 모멘텀 항) 를 분리하여 그래디언트 얽힘을 방지합니다. 컴퓨터 비전에서는 효과적이었으나, DualOptim 은 LLM 에서는 미미한 개선만 가져옵니다. 저자들은 LLM 망각에서 망각과 유지를 위한 그래디언트가 항상 직접적으로 상충되는 것은 아니며, 특히 학습 후기 단계에서는 공통된 표현을 공유하는 경우가 많음을 관찰했습니다. 완전히 분리된 상태는 이러한 공유 특징을 활용하지 못하며, 완전히 공유된 상태는 상충을 처리하지 못합니다.

방법론: DualOptim+
이 논문은 완전히 공유된 상태와 완전히 분리된 옵티마이저 상태 사이의 간극을 메우도록 설계된 플러그 앤 플레이 최적화 프레임워크인 **DualOptim+**를 제안합니다. 핵심 혁신은 옵티마이저 상태를 두 가지 구성 요소로 분해하는 것입니다:

1. 공유 베이스 상태 ($B$): 이 상태는 망각과 유지 목표 모두에서 공유되는 공통 표현을 포착합니다. 두 목표의 그래디언트 ($\nabla_\theta L_f$ 및 $\nabla_\theta L_r$) 를 사용하여 연합적으로 업데이트됩니다.
2. 분리된 델타 상태 ($\Delta_f, \Delta_r$): 이 상태들은 목표별 잔차를 보존합니다. 특정 목표의 그래디언트와 편향 보정된 베이스 상태 간의 차이를 기반으로 업데이트됩니다.
   • $\Delta_f \leftarrow \beta \Delta_f + (1-\beta)(\nabla_\theta L_f - \hat{B})$
   • $\Delta_r \leftarrow \beta \Delta_r + (1-\beta)(\nabla_\theta L_r - \hat{B})$

파라미터 업데이트:
파라미터 업데이트에 사용되는 최종 모멘텀 항은 편향 보정된 베이스 상태와 각각의 편향 보정된 델타 상태 ($\hat{B} + \hat{\Delta}_f$ 또는 $\hat{B} + \hat{\Delta}_r$) 를 결합하여 재구성됩니다.

핵심 이론적 통찰:
이론적 및 수치적 분석은 DualOptim+ 가 적응형 중간체로 기능함을 보여줍니다.

• 그래디언트가 양의 상관관계를 가질 때 (공유 특징이 지배적), 델타 상태는 0 으로 수렴하며, 이 방법은 교대 (Alternate) 방식 (공유 베이스 활용) 처럼 동작합니다.
• 그래디언트가 음의 상관관계를 가질 때 (상충되는 목표), 베이스 상태는 0 으로 수렴하며, 이 방법은 DualOptim 방식 (분리된 상태 활용) 처럼 동작합니다.
  이러한 적응성은 옵티마이저가 망각과 유지 그래디언트 간의 방향적 상충에 기반하여 동적으로 조정되도록 합니다.

DualOptim+ 8bit
추가적인 옵티마이저 상태 저장으로 인한 메모리 오버헤드 (표준 Adam 옵티마이저의 3 배) 를 해결하기 위해, 저자들은 DualOptim+ 8bit를 도입합니다. 이 변형은 옵티마이저 상태를 8 비트 정밀도로 양자화하여 성능을 저하시키지 않으면서 메모리 소비를 (표준 Adam 의 약 3/4 수준으로) 크게 줄입니다.

주요 기여

1. 새로운 프레임워크: 공유 베이스 상태와 분리된 델타 상태를 도입하여 공유 및 분리된 옵티마이저 상태를 효과적으로 연결하는 DualOptim+ 제안.
2. 플러그 앤 플레이 설계: 상태를 저장하는 모든 옵티마이저 (예: AdamW, Muon) 와 다목적 최적화에 적용 가능한 프레임워크.
3. 실증적 검증: 가짜 데이터셋 (TOFU), 실제 시나리오 (Llama-3 망각), 안전 정렬, 다중 작업 학습 작업에 걸친 광범위한 실험.
4. 효율성: 메모리 오버헤드를 완화하는 8 비트 양자화 변형 도입.

실험 결과
이 논문은 다양한 LLM(Phi-1.5, Llama-2, Llama-3) 과 작업에 걸쳐 DualOptim+ 를 연합, 교대, DualOptim 기준선과 비교 평가합니다:

• 가짜 시나리오 (TOFU): DualOptim+ 는 망각 효율성 (FE) 과 모델 유용성 (MU) 간의 절충으로 정의된 전체 성능 (OVR) 에서 일관되게 최상의 성과를 거둡니다. 특히 공유 특징을 더 잘 활용함으로써 표적 망각에서 DualOptim 을 능가합니다.
• 실제 시나리오: Llama-3 에서 실제 개인을 망각하는 작업에서 DualOptim+ 는 기준선보다 하류 작업 (ARC-c, MMLU 등) 에서 우수한 유용성을 유지하는 반면, DualOptim 은 때로 더 간단한 교대 접근법보다 성능이 떨어지기도 합니다.
• 안전 정렬: 안전하지 않은 지식을 망각하되 지시 따르기 능력을 유지하는 작업에서 DualOptim+ 는 XSTest 벤치마크에서 가장 높은 안전 - 유용성 절충과 가장 낮은 과도 거절률을 달성합니다.
• 다중 작업 학습: 그래디언트 상충이 높은 망각과 달리, 상충이 덜 두드러지는 다중 작업 학습에서 DualOptim+ 는 상당한 성능 향상을 계속 제공하는 반면, DualOptim 은 부정적인 이득을 보입니다.
• 메모리 효율성: 8 비트 변형은 전체 정밀도 방법의 성능 이점을 유지하면서 표준 기준선과 거의 동등한 메모리 사용량을 달성합니다.

의의 및 주장
저자들은 DualOptim+ 가 기존 방법들보다 지식 제거와 유용성 보존 간의 더 안정적이고 효과적인 절충을 제공한다고 주장합니다. 이 방법을 단순한 망각 도구가 아닌, LLM 정렬 및 다목적 학습과 같은 여러 상충되는 학습 목표를 포함하는 더 넓은 시나리오에 적용 가능한 일반화 가능한 최적화 프레임워크로 위치시킵니다. 논문은 공유 및 분리된 상태를 연결하는 적응적 특성이 LLM 망각에서 발견되는 동적 그래디언트 관계를 처리하는 데 필수적임을 강조합니다.