기술 요약: 딥 저랭크 잔류 증류를 통한 사전 학습 가중치 잠금

문제 정의

오픈 가중치 대규모 언어 모델 (LLM) 의 확산은 강력한 AI 시스템에 대한 접근을 민주화했지만, 동시에 모델 제공자가 이러한 가중치가 어떻게 활용되는지에 대해 갖는 통제력을 약화시켰습니다. 법적 라이선스는 무단 적응을 제한하려 시도하지만, 기술적 집행 수단이 부재합니다. 모델 가중치와 아키텍처에 대한 완전한 접근 권한을 가진 적응형 공격자는 종종 단순한 구조적 방어를 역추적하거나 최적화를 통해 이를 우회할 수 있습니다.

이 논문이 다루는 핵심 과학적 과제는 모델 잠금입니다. 이는 사전 학습된 모델을 수정하여, 이를 적응시키는 것 (예: 파인튜닝을 통한) 이 원래 모델을 적응시키는 것보다 본질적으로 더 비싸게 만들되, 모델의 추론 시 성능과 일반적 능력을 엄격히 유지하는 능력입니다. 저자들은 이를 케르크호프스 원칙에 기반한 적응형 위협 모델 하에 형식화했습니다. 여기서 방어는 공격자가 잠금 메커니즘의 설계를 알고 있더라도 여전히 효과적이어야 합니다. 목표는 "잠금 계수" $\kappa > 1$을 달성하는 것으로, 하류 작업에서 목표 성능 수준에 도달하는 데 드는 비용 (시간, FLOPs, 또는 GPU 시간) 이 잠긴 모델의 경우 원래 모델에 비해 적어도 $\kappa$배 이상 증폭되도록 하는 것입니다.

기존 접근법의 한계

이 논문은 이전 방어 전략들의 근본적인 결함을 다음과 같이 지적합니다:

1. 대칭성 기반 방어: 가중치 대칭성 (예: 스케일링 또는 가역 행렬 삽입) 을 이용하여 최적화를 방해하는 방법들은 자명하게 역추적 가능합니다. 공격자는 가중치 노름을 검사하거나 특이값 분해 (SVD) 를 계산하여 가중치를 재균형화하고 원래 최적화 지형을 복원할 수 있습니다.
2. 최적화 기반 방어: 불량 조건 손실 지형 (예: 헤시안 트레이스 최대화 또는 장벽 모듈 삽입) 을 최적화하여 방어를 내재화하려는 접근법은 "자기 파괴적" 역설에 직면합니다. 방어자는 최적화를 사용하여 최적화에 저항하는 해를 찾아야 하지만, 방어자가 활용하는 동일한 기울기 정보가 공격자에게도 존재합니다. 단순 MLP 에 대한 실험 결과, 경사 하강법은 사전 학습된 초기화 상태에서 이러한 "잠긴" 고비용 영역에 도달하는 데 종종 실패함이 확인되었습니다.

방법론: DLR-Lock

저자들은 자동 미분의 추론–훈련 비대칭성을 활용하는 방어 기법인 DLR-Lock을 제안합니다. 추론은 순방향 통과만 필요로 하고 중간 활성화 값을 폐기하는 반면, 훈련은 역전파를 위해 모든 중간 활성화 값을 저장해야 합니다. DLR-Lock 은 트랜스포머의 표준 SwiGLU MLP 모듈 각각을 **딥 저랭크 잔류 네트워크 (DLR-Net)**로 대체합니다.

아키텍처

표준 MLP 잔류 블록 $x' = z + \text{FFN}(\text{RMSNorm}(z))$은 깊이 $L$의 저랭크 잔류 네트워크로 대체됩니다:
$$h_0 = z, \quad h_{i+1} = h_i + \alpha_i U_i \sigma(V_i \text{RMSNorm}(h_i)), \quad x' = h_L$$
여기서:

• $V_i \in \mathbb{R}^{r \times d}$와 $U_i \in \mathbb{R}^{d \times r}$은 병목 랭크 $r \ll d$를 갖는 저랭크 투영 행렬입니다.
• $\alpha_i$는 학습 가능한 ReZero 스칼라입니다.
• 깊이 $L$은 전체 파라미터 수를 원래 MLP 와 일치시키거나 (또는 약간 초과하도록) 선택되며, 파라미터를 몇 개의 깊은 레이어가 아닌 많은 얕은 레이어에 분배합니다.

세 가지 잠금 메커니즘

이 아키텍처 대체는 적응에 대한 세 가지 구체적인 장벽을 도입합니다:

1. 활성화 메모리 증가: 역전파 중 활성화 값을 저장하는 데 필요한 메모리는 깊이 $L$에 비례하여 선형적으로 증가합니다. 순진한 파인튜닝을 사용하는 공격자의 경우, 활성화 메모리 오버헤드는 원래 모델 대비 $\Theta(d/r)$로 스케일링되어 더 작은 배치 크기나 더 비싼 하드웨어를 강요할 수 있습니다.
2. 최적화 불안정성: 깊은 잔류 구조는 원래 모델의 훈련 레시피와 아키텍처적 불일치를 생성합니다. 원래의 얕은 MLP 들을 위해 조정된 표준 파인튜닝 하이퍼파라미터와 옵티마이저는 깊은 DLR-Net 에 적용될 때 불안정해지거나 비효율적이 됩니다.
3. 불균형한 역방향 통과 오버헤드: 복잡한 메모리 접근 패턴과 활성화 값 재계산 (그래디언트 체크포인트를 사용하더라도) 의 필요성으로 인해 역방향 통과는 순방향 통과보다 훨씬 더 많은 오버헤드를 초래합니다. 특히, 순방향 통가를 가속화하는 컴파일러 최적화 (예: CUDA 그래프 재생) 는 역방향 통가에 적용되지 않으므로, 기울기를 계산해야 하는 공격자에게 비용 격차를 더욱 확대합니다.

증류 절차

잠긴 모델이 원래 능력을 유지하도록 (유용성 보존) 하기 위해, 저자들은 2 단계 증류 파이프라인을 사용합니다:

1. 모듈별 증류: 각 DLR-Net 은 해당 사전 학습된 MLP 레이어의 입력 - 출력 매핑을 복제하도록 독립적으로 훈련되며, 상대 평균 제곱 오차 (MSE) 를 최소화합니다.
2. 로그이트 증류: 모든 레이어가 대체된 후, 누적된 오차는 원래 모델을 교사 (teacher) 로 사용하여 DLR-Net 파라미터 (어텐션과 임베딩은 고정된 상태) 를 파인튜닝함으로써 보정됩니다. 목적은 교사와 학생 출력 분포 간의 상위 $k$개 KL 발산을 최소화하는 것입니다.

실험 결과

이 방법은 Nemotron-CC 데이터셋을 사용하여 Qwen3-0.6B 및 Qwen3-8B 모델에서 검증되었습니다.

• 품질 보존: 잠긴 모델은 원래 모델의 능력을 성공적으로 유지했습니다. 2 단계 증류 후 자연어 코퍼스 (Nemotron-CC, WikiText-103) 에 대한 퍼플렉시티와 하류 벤치마크 (MMLU, ARC, HellaSwag 등) 의 성능은 사전 학습 베이스라인과 몇 점 이내로 회복되었습니다.
• 추론 오버헤드: 이론적 FLOPs 는 유사하게 유지되지만, DLR-Net 의 순차적 특성으로 인해 월클럭 오버헤드가 발생합니다. 그러나 컴파일러 최적화 (예: torch.compile with reduce-overhead) 로 인해 이는 완화됩니다. 대형 모델과 배치 크기의 경우 오버헤드 비율이 현저히 감소합니다 (예: 8B 모델의 경우 eager 모드에서 $144\times$에서 컴파일 시 $11\times$로 감소), 합법적 사용자에게는 관리 가능한 수준으로 유지됩니다.
• 파인튜닝에 대한 저항성:
  • 수렴: WikiText-103 에서 파인튜닝 시, 그래디언트 체크포인트가 활성화된 상태에서도 잠긴 모델은 베이스라인보다 훨씬 느리게 수렴하여 더 높은 손실 수준에서 정체되었습니다.
  • 불안정성: 더 높은 학습률 ($10^{-4}$) 에서 잠긴 모델은 완전히 발산한 반면, 베이스라인은 정상적으로 훈련되어 이 방어책이 유효한 학습률 범위를 좁힌다는 것을 나타냅니다.
  • 비용 증폭: 잠긴 모델의 경우 역방향 통과가 총 훈련 시간의 89% 를 차지했습니다 (베이스라인은 83%). 목표 손실에 도달하는 데 걸리는 월클럭 시간은 훨씬 더 컸습니다. 저자들은 역방향 통과 오버헤드가 순방향 통과보다 불균형적으로 높음을 지적하며, 이 방어책이 훈련 특화 비용에 의존함을 확인했습니다.

중요성과 주장

이 논문은 방어 메커니즘의 비밀이 아닌 구조적 속성에 의존하며 옵티마이저 무관인 새로운 모델 잠금 방어 축을 확립했다고 주장합니다. 주요 기여점은 다음과 같습니다:

• 형식화: 적응형 위협 모델 하에서 적응 비용 증폭으로서의 모델 잠금을 정의합니다.
• 새로운 메커니즘: 자동 미분의 추론 - 훈련 비대칭성을 활용하여 평가는 저렴하지만 미분은 비싼 방어를 생성할 수 있음을 입증합니다.
• 강건성: 그래디언트 체크포인트, 부분 가중치 업데이트, 또는 LoRA 를 사용하는 완전한 전략 지식을 가진 적응형 공격자에도 이 방법이 견딜 수 있음을 보여줍니다.
• 실용성: 잠금 메커니즘이 모델의 유용성을 희생하지 않으며, 합법적 사용자는 거의 원래 성능으로 추론할 수 있는 반면 무단 적응에는 측정 가능한 비용 장벽을 부과함을 검증합니다.

저자들은 커스텀 커널이나 컴파일을 통해 순방향 통과 오버헤드를 완화할 수 있지만, 역전파 비용의 근본적인 비대칭성은 파인튜닝에 대한 견고한 장벽으로 남는다고 결론지었습니다.