기술 요약: 대규모 언어 모델의 지속적 학습을 위한 SAE-FD

문제 제기

대규모 언어 모델 (LLM) 은 여러 작업에 순차적으로 적응할 때 **파괴적 망각 (catastrophic forgetting)**이라는 도전에 직면합니다. 기존 지속적 학습 (CL) 방법들은 주로 **밀집 표현 공간 (weight, gradient, 또는 output spaces)**에서 작동합니다. 이러한 접근법의 근본적인 한계는 **특성 중첩 (feature superposition)**으로, 밀집 활성화 공간 내에서 여러 개념이 겹치는 차원에 인코딩되는 현상입니다. 이러한 얽힘은 이전 학습 지식을 선택적으로 보호하면서도 새로운 작업 학습을 부수적으로 제한하지 않도록 하는 것을 어렵게 만들어, 비최적의 가소성 - 안정성 트레이드오프를 초래합니다.

방법론: SAE-FD

저자들은 사전 훈련된 희소 오토인코더 (SAE) 의 희소 특성 공간에 모델 표현을 고정하는 **SAE-FD(Sparse Autoencoder Feature Distillation)**를 제안합니다. 밀집 활성화를 희소하고 과완전한 기저로 분해함으로써, SAE-FD 는 중첩된 표현을 분리하여 더 표적화된 정규화를 가능하게 합니다.

이 방법은 3 단계 파이프라인으로 작동합니다:

1. SAE 훈련 (1 단계)

지속적 학습 시퀀스 시작 전에, 기본 모델의 다양한 텍스트 활성화 데이터로 **게이트드 SAE(Gated SAE)**를 훈련합니다.

• 입력: 분류, QA, 지식 벤치마크 등 다양한 소스에서 얻은 마지막 레이어 MLP 활성화 ($h \in \mathbb{R}^d$).
• 아키텍처: 게이트드 SAE 는 인코딩을 게이트 경로 ($g$) 와 크기 경로 ($m$) 로 분리하여 ReLU 전 특성 ($f_{pre}$) 과 ReLU 후 특성 ($f$) 을 생성합니다.
• 목적: 재구성 손실 ($\|\hat{h} - h\|_2^2$) 과 L1 희소성 페널티 ($\lambda_{\ell1}\|f\|_1$) 를 최소화합니다.
• 결과: 밀집 활성화를 희소 특성 ($D \gg d$) 으로 분해하여 개념을 효과적으로 분리하는 고정된 SAE.

2. 작업 훈련 및 앵커 캡처 (2 단계)

시퀀스 내의 각 작업 $t$에 대해:

• 파인튜닝: 표준 언어 모델링 목적 함수를 사용하여 LoRA 로 모델을 파인튜닝합니다.
• 앵커 캡처: 작업 훈련 후, 모델은 소수의 **앵커 샘플 ($N_a$)**에 대해 순전파를 수행합니다. 토큰별 MLP 출력 ($H_i$) 을 추출하여 저장합니다.
• 버퍼: 이러한 앵커는 작업 간에 누적됩니다. 이후 훈련 중에는 버퍼에서 균일하게 배치 (batch) 를 샘플링하여 증류 손실을 계산함으로써 이전 모든 작업에 걸쳐 비용을 분산합니다.

3. 희소 인식 특성 증류 (3 단계)

새로운 작업 ($t+1$) 훈련 중, 모델은 고정된 SAE를 통해 인코딩된 현재 표현과 저장된 앵커를 비교하여 정규화됩니다. 증류 손실 ($L_{FD}$) 은 두 가지 상호 보완적인 항으로 구성됩니다:

• 코사인 방향 손실 ($L_{cos}$): 현재 ReLU 전 특성 벡터와 앵커 특성 벡터 간의 각도 편차를 패널티합니다. 이는 표현의 전역 방향을 보존합니다.
• 활성 크기 손실 ($L_{mag}$): 앵커에서 활성 (ReLU 후 $> 0$) 이었던 특성에만 평균 제곱 오차 (MSE) 를 적용합니다. 이는 이전 작업에 기여하는 특성의 크기를 제한하면서, 비활성 특성은 새로운 작업을 위해 모집될 수 있도록 자유롭게 둡니다.

적응형 $\lambda$ 메커니즘

작업 학습과 정규화 사이의 균형을 맞추기 위해, 저자들은 총 손실 $L = L_{task} + \lambda \cdot L_{FD}$의 가중 계수 $\lambda$에 대한 적응형 메커니즘을 도입합니다.

• 고정된 $\lambda$ 대신, 시스템은 총 손실에 대한 증류 손실의 고정 기여 비율 $\rho$를 목표로 합니다.
• $\lambda$는 $\rho = \frac{\lambda L_{FD}}{\lambda L_{FD} + L_{task}}$를 만족하도록 지수 이동 평균 (EMA) 을 통해 업데이트됩니다.
• 동작: $\lambda$는 새로운 작업이 시작될 때 (작업 손실 $L_{task}$가 크고 망각 위험이 높을 때) 높게 시작되어, 모델이 수렴함에 따라 완화되므로 파인튜닝 중 과도한 제약을 방지합니다.

주요 기여

1. SAE-FD 프레임워크: 밀집 공간에 내재된 특성 중첩 문제를 해결하기 위해, SAE 의 희소이고 분리된 특성 공간에서 정규화를 수행하는 최초의 지속적 학습 방법.
2. 희소 인식 증류 손실: 코사인 방향 보존과 활성 특성에 대한 선택적 크기 매칭을 결합한 새로운 손실 함수로, SAE 표현의 희소성을 활용하여 새로운 작업과의 간섭을 최소화합니다.
3. 적응형 정규화: 작업 손실과 증류 손실의 상대적 크기에 따라 정규화 강도를 동적으로 조정하는 적응형 $\lambda$ 메커니즘으로, 고위험 기간 동안 강력한 보호를 제공하면서도 수렴을 방해하지 않습니다.

실험 결과

저자들은 세 가지 모델 아키텍처 (LLaMA-2-7B, Vicuna-7B, Mistral-7B) 에 걸쳐 두 가지 벤치마크에서 SAE-FD 를 평가했습니다.

• TRACE 벤치마크 (8 개의 이질적 작업):
  • LLaMA-2-7B 에서 SAE-FD 는 **52.70% 평균 정확도 (AA)**를 달성했으며, **−0.46% 역방향 전이 (BWT)**만 발생했습니다.
  • 이는 가장 강력한 공개된 정규화 기반 베이스라인 (GORP) 보다 +2.3% AA로 우월하며, 망각을 34% 감소시켰습니다 (−0.46 대 −0.70 BWT).
  • BWT 가 −3% 보다 나쁜 다른 LoRA 기반 방법들 (EWC, O-LoRA, TreeLoRA) 보다 훨씬 뛰어난 성능을 보였습니다.
• 일반화:
  • SAE-FD 는 세 모델 패밀리 전반에 걸쳐 일관되게 최고의 AA 와 BWT 를 달성했습니다.
  • Vicuna-7B 에서 O-LoRA 보다 +10.1% AA로 우월했으며, 전체 파인튜닝 베이스라인을 포함한 모든 방법 중 최고의 BWT(−0.69%) 를 기록했습니다.
• Ablation 연구:
  • 손실 구성 요소: 코사인 방향과 활성 크기 손실 모두 필수적이며, 어느 하나만 사용하면 AA 가 낮아지거나 망각이 증가했습니다.
  • 적응형 $\lambda$: 적응형 메커니즘을 고정된 $\lambda=1.0$으로 대체하면 BWT 가 −0.46% 에서 −2.38% 로 저하되었습니다.
  • SAE 대 원시 MSE: 밀집 활성화에서의 원시 MSE 증류보다 SAE 특성 공간에서의 증류가 더 우수하여, 분리된 공간의 이점을 확인했습니다.
• 효율성: 계산 오버헤드는 SeqLoRA 의 약 1.85 배이며, 이전 작업 수에 관계없이 단계당 일정한 비용이 발생합니다. 저장 공간은 앵커 활성화를 위해 작업당 약 0.8 GB 가 필요합니다.

중요성과 주장

이 논문은 SAE-FD 가 **기계적 해석 가능성 (SAE 를 통해)**과 지속적 학습을 성공적으로 연결한다고 주장합니다. 신경망을 이해하기 위해 개발된 희소하고 분리된 표현을 활용함으로써, 이 방법은 이를 제어하기 위한 효과적인 귀납적 편향을 제공합니다.

저자들은 희소 특성 공간에서 작동함으로써 지식 보존을 더 표적화할 수 있으며, 밀집 공간 정규화를 괴롭히는 "특성 중첩" 문제를 해결한다고 주장합니다. 이 방법은 우수한 가소성 - 안정성 트레이드오프를 달성하여, LLM 이 이전에 학습한 작업에서 높은 성능을 유지하면서 처음부터 다시 훈련하지 않고도 진화하는 작업에 적응할 수 있게 합니다.

저자가 지적한 한계:

• 모델별 SAE 에 대한 일회성 훈련 비용이 필요합니다.
• 앵커 활성화에 대한 저장 요구 사항은 작업 수에 비례하여 증가합니다 (압축 가능함).
• TRACE 벤치마크의 경우 고정된 작업 순서로 실험이 수행되었으나, 다중 모델 일관성은 견고성을 시사합니다.
• 실험 설정의 차이 (예: 전체 파인튜닝 대 LoRA, 에포크 수) 로 인해 일부 베이스라인과의 비교는 신중하게 해석해야 합니다.