From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

기술 요약: 가중치에서 특징으로: LLM 지속 학습을 위한 SAE 유도 활성화 정규화

문제 정의
대규모 언어 모델(LLM)의 지속 학습(Continual Learning, CL)은 순차적 미세 조정이 이전 작업에 필요한 표현을 덮어쓰는 치명적 망각(catastrophic forgetting) 문제에 직면해 있다. Elastic Weight Consolidation (EWC)와 같은 정규화 기반 방법들이 표준적인 방식이지만, LLM에서는 성능이 현저히 떨어진다. 저자들은 이러한 실패가 구조적인 문제, 즉 LLM의 "다의성(polysemantic)" 특성에서 기인한다고 주장한다. 중첩 가설(superposition hypothesis)로 인해 개별 뉴런과 그에 연관된 가중치는 여러 무관한 개념을 동시에 인코딩한다. 결과적으로, 가중치 공간의 중요도 추정치(예: 피셔 정보)는 너무 거칠다. 특정 개념 A를 보존하기 위해 가중치를 보호하면 의도치 않게 개념 B까지 제약하게 되어, 비선택적인 페널티를 생성하고 가소성(plasticity)을 저해한다. 또한, 기존 방법들은 새로운 작업을 학습하는 데 필요한 가소성을 명시적으로 강제하지 못해, 모델이 기존 지식은 유지하지만 새로운 적응은 하지 못하는 "무기력함(intransigence)"을 초래하는 경우가 많다.

방법론
본 논문은 **SAE 유도 활성화 정규화(SAE-guided activation regularization)**를 제안한다. 이는 정규화 대상을 엉킨 파라미터 공간에서 사전 학습된 희소 오토인코더(Sparse Autoencoder, SAE)를 사용하는 더 선택적인 활성화 공간으로 전환하는 방법이다.

1. 좌표계 전환: 가중치 드리프트(drift)를 페널티로 부과하는 대신, 이 방법은 고정된, 사전 학습된 SAE를 고정된 단의적(monosemantic) 특징 사전(dictionary)으로 사용한다. SAE는 밀집된 모델 활성화를 해석 가능하고 덜 얽힌 희소 특징 방향으로 분해한다.
2. 제약 최적화 정식화: 학습 목적 함수는 임의적인 페널티 설계가 아닌 제약 최적화 문제로부터 도출된다. 모델은 다음 두 가지 제약 조건을 준수하면서 현재 작업 손실을 최소 최소화해야 한다.
   • 안정성(Stability): 작업과 무관한 특징(낮은 관련성)은 특정 예산($\varepsilon_s$) 이상으로 드리프트해서는 안 된다.
   • 가소성(Plasticity): 작업과 관련된 특징(높은 관련성)은 퇴보적인 해(degenerate solutions, 모델이 아무것도 배우지 못하는 상태)를 피하기 위해 충분히($\ge \varepsilon_p$) 드리프트해야 한다.
3. 손실 함수: 라그랑주 완화법(Lagrangian relaxation)과 제곱 힌지 페널티(squared-hinge penalty)를 사용하여, 저자들은 보호 손실(protect loss)(보호된 특징의 과도한 드리프트를 페널티로 부과)과 가이드 손실(guide loss)(관련된 특징의 불충분한 드리프트를 페널티로 부과)로 구성된 학습 손실을 도출한다. EWC는 표현 공간이 가중치이고 가소성 제약이 없는 이 프레임워크의 특수한 사례로 간주된다.
4. 마스크 구축: 각 새로운 작업에 대해, 현재 작업 데이터를 고정된 베이스 모델과 SAE에 통과시켜 연속적인 관련성 마스크($M_t$)를 구축한다. 이 마스크는 활성화 크기를 통해 관련 특징을 식별하고, 디코더 벡터에 대한 k-NN을 사용하여 인접한 특징으로 관련성을 전파한다. 결정적으로, 마스크가 구축되면 모든 작업 데이터는 폐기되며, 압축된 마스크만이 유지된다.
5. 학습: 학습 중에 현재 모델(LoRA 적용)과 고정된 베이스 모델은 모두 SAE에 의해 인코딩된다. 특징 드리프트는 마스크에 의해 가중치가 부여되어 보호 및 가이드 손실을 계산한다.

주요 기여

• SAE 유도 정규화: 활성화 공간에서 SAE 특징을 사용하여 정규화하는 실용적인 CL 방법이다. 이는 이전 작업 데이터, 저장된 활성화 값, 작업별 파라미터를 필요로 하지 않으며, 마스크 구축 후에는 추론 시점의 태스크 라우팅도 필요하지 않다.
• 원칙적 도출: 안정성-가소성 트레이드오프를 SAE 특징 공간에서의 제약 최적화 문제로 정식화하여, 라그랑주 완화법을 통해 제곱 힌지 학습 손실을 산출했다.
• 다의성에 대한 기계적 증거: 본 논문은 가중치 공간의 보호가 비선택적이라는 기계적 증거를 제공한다. 실험 결과, 작업 관련 표현은 SAE 특징 공간에서는 선형적으로 분리 가능(AUC 0.88)하지만, 가중치/뉴런 공간에서는 우연 수준과 구별할 수 없다(AUC ~0.50). 또한, 하나의 작업을 위해 가중치를 보호하는 것은 후속 작업의 특징을 원래 작업의 9196% 비율로 제약하는 반면, 특징 공간에서의 보호는 이러한 부수적 제약을 43~61%로 줄인다.
• 확장성: 파라미터 공간에 비해 낮은 차원의 특징 공간 덕분에, 작업당 저장 용량이 가중치 앵커(기가바이트 단위)에서 서브 메가바이트 단위의 특징 마스크로 감소한다.

결과
방법론은 Gemma-2 9B-it와 LoRA를 사용하여 TRACE-5000(8개 교차 도메인 작업) 및 MedCL(10개 생물 의학 작업)에서 평가되었다.

• 성능: TRACE에서 SAE 유도 정규화는 가장 강력한 성능(OP = 0.545)을 달 기록하며, 가중치 공간 정규화 방법(EWC: 0.447, SI: 0.513), 그래디언트 투영 방법, 그리고 리플레이 베이스라인(데이터 10% 사용 시에도)을 능가했다. MedCL에서도 이전 작업 예시를 보유하지 않는 방법들 중 선두를 차지했다(OP = 0.510).
• 안정성-가소성 트레이드오프: 가소성을 심하게 희생하여 높은 유지력(ARA = 0.987)을 달성한 EWC와 달리, 제안된 방법은 학습된 성능의 95%를 유지하면서도 높은 가소성(0.567, 보호되지 않은 미세 조정과 유사)을 유지했다.
• 효율성: 이 방법은 최소한의 계산 오버헤드(~1.88s/step vs 1.70s baseline)를 추가하며, 저장 요구 사항을 획기적으로 줄였다(EWC 앵커의 6.5 GB 대비 작업당 412 KB).

의의 및 주장
본 논문은 LLM에서 가중치 공간 정규화의 저조한 성능이 하이퍼파라미터 튜닝의 문제가 아니라 다의성의 구조적 결과라고 주장한다. 정규화 대상을 단의적 특징 기반으로 이동함으로써, 이 방법은 더 선택적인 지식 보호를 달성한다. 저자들은 정규화의 효과가 제약이 부과되는 표현 공간에 결정적으로 의존한다고 주장한다. 이 방법은 정규화가 모델의 내부 개념 구조와 일치하는 좌표계에서 작동한다면, 과거 데이터를 저장하거나 추론 시점에 모델 아키텍처를 수정하지 않고도 강력한 지속 학습 성능을 달성할 수 있음을 보여준다. 결론적으로, SAE 기반 정규화는 LLM을 위한 보다 선택적이고 확장 가능한 대안을 제공한다.