기술 요약: 고도 희소성 하에서 가지치기 강건성 향상을 위한 가중치 집중 정규화

문제 제기

심층 신경망 (DNN) 은 비전 및 언어 작업에서 최첨단 성능을 달성했으나, 종종 상당한 계산 자원을 필요로 하여 자원 제약 환경에서의 배포를 제한합니다. 모델 가지치기는 덜 중요한 매개변수를 제거하여 모델 크기를 줄이는 것을 목표로 합니다. 재학습을 동반한 반복적 가지치기는 효과적이지만 계산 비용이 많이 듭니다. 반면, 단일 단계에서 재학습 없이 중복 매개변수를 제거하는**원샷 가지치기 (One-shot pruning)**는 더 효율적인 대안을 제공합니다. 그러나 표준 목적함수로 훈련된 모델은 공격적인 희소성 하에서 상당한 정확도 저하를 겪는 경우가 많습니다.

이전 연구들은 주로 두 가지 방향을 통해 이를 완화하려 시도했습니다:

1. 정규화 기반 접근법 (예: $\ell_1$, DeepHoyer): 훈련 중 가중치 분포를 형성합니다. 그러나 $\ell_1$은 모든 가중치를 균일하게 축소하는 반면, DeepHoyer 는 스케일 불변 희소성을 유도합니다. 둘 다 정보 있는 소수의 매개변수 위에 가중치 에너지를 명시적으로 집중시키지는 않습니다.
2. 가지치기 강건 최적화기 (예: SAM, CrAM, S2SAM): 일반화를 개선하기 위해 손실 지형을 평탄하게 만듭니다.

식별된 핵심 한계는 표준 희소성 유도 방법들이 희소성과 정보 보존 사이의 균형을 명시적으로 다루지 못한다는 점입니다. 가중치 중요도가 확산되어 분포되어 있을 때, 작은 크기의 가중치를 제거하더라도 "중요한" 에너지가 소수의 지배적인 매개변수 위에 충분히 집중되지 않았기 때문에 성능이 저하될 수 있습니다.

방법론: 가중치 집중 정규화기 (WCR)

저자들은 훈련 시 가중치 크기 분포를 재구성하도록 설계된 훈련 시간 정규화기인**가중치 집중 정규화기 (Weight Concentration Regularizer, WCR)**를 제안합니다. 가중치를 균일하게 축소하는 기존 방법들과 달리, WCR 은 나머지 가중치를 0 으로 유도하면서 가중치 에너지를 소수의 매개변수 집합 위에 집중시키는 것을 목표로 합니다. 이는 소수의 매개변수가 대부분의 가중치 에너지를 보유하는 "무거운 꼬리 (heavy-tailed)" 분포를 생성하여, 크기 기반 가지치기가 무시할 수 있는 매개변수를 최소한의 성능 손실로 제거할 수 있게 합니다.

수학적 공식화

WCR 은 절대 가중치의 분산 역수를 패널티로 작용합니다. 선택된 층의 가중치 $w^{(\ell)}$에 대해, 저자들은 원점에서 미분 가능성을 보장하기 위해 먼저 부드러운 절대값 변환을 적용합니다:
$$\tilde{w}^{(\ell)}_i = \sqrt{(w^{(\ell)}_i)^2 + r} \approx |w^{(\ell)}_i|, \quad r = 10^{-8}$$
스케일 차이를 처리하기 위해 층별로 분산을 계산합니다:
$$\text{Var}(\tilde{w}^{(\ell)}) = \frac{1}{n_\ell} \sum_{i=1}^{n_\ell} \left( \tilde{w}^{(\ell)}_i - \frac{1}{n_\ell} \sum_{j=1}^{n_\ell} \tilde{w}^{(\ell)}_j \right)^2$$
정규화 패널티 $\psi(w)$는 선택된 모든 층에 걸친 분산 역수의 합입니다:
$$\psi(w) = \sum_{w^{(\ell)} \in \mathcal{W}} \frac{1}{\text{Var}(\tilde{w}^{(\ell)}) + \epsilon}$$
여기서 $\epsilon = 10^{-8}$은 수치적 안정성을 보장합니다. 총 훈련 목적함수는 다음과 같습니다:
$$\mathcal{L}_{total}(w) = \mathcal{L}(w) + \lambda \psi(w)$$
여기서 $\lambda > 0$는 균형 조정 계수입니다.

이론적 분석

본 논문은 WCR 이 확률적 경사 하강법 (SGD) 의 표준 수렴 속도를 유지함을 보여주는 수렴 분석을 제공합니다. 손실 함수의 $\beta_1$-평탄성과 정규화기의 $\beta_2$-평탄성 하에서, 결합된 목적함수는 표준 수렴 - 정적점 보장 조건을 만족합니다. 정규화기가 가중치의 결정론적 함수이므로 추가적인 확률적 경사 분산을 도입하지 않습니다. 결과적으로 점근적 수렴 속도는 표준 SGD 와 일치하는 $O(1/\sqrt{T})$로 유지됩니다.

주요 기여

1. WCR 도입: 가중치 에너지를 소수의 매개변수 집합 위에 집중시키는 새로운 훈련 시간 정규화기로, 균일한 축소가 아닌 무거운 꼬리 가중치 분포를 명시적으로 장려합니다.
2. 이론적 보장: mild 한 평탄성 가정 하에서 WCR 이 SGD 의 표준 수렴 속도를 유지함을 증명합니다.
3. 실증적 검증: 다양한 설정에 대한 포괄적 평가:
   • 이미지 분류: CIFAR-10/100, SVHN 에서의 ResNet-18/50, WideResNet-28-10 및 Tiny-ImageNet 에서의 ViT-B/32.
   • 의료 분할: LGG 뇌 MRI 데이터셋에서의 ResNet-50–UNet.
   • LLM 파인튜닝: Commonsense-170K 에서 LoRA 를 사용한 Qwen-2.5-1.5B 및 LLaMA-3.2-1B.
4. 호환성: WCR 이 기존 가지치기 강건 최적화기 (SAM, CrAM, S2SAM) 와 호환되며 이를 보완하여 추가적인 향상을 이끌어낼 수 있음을 입증합니다.

실험 결과

저자들은 재학습 없이 원샷 크기 기반 가지치기 하에서 WCR 을 평가했습니다.

• 비전 작업 (CNN 및 ViT): WCR 은 특히 고도 희소성 수준 (예: 96% 가지치기) 에서 가지치기 후 정확도를 일관되게 향상시켰습니다.
  • CIFAR-10에서 ResNet-18 을 사용할 때, SGD 에 WCR 을 추가하면 96% 가지치기 시 정확도가 19.69% 에서 25.70% 로 향상되었습니다.
  • SAM과 결합했을 때 WCR 은 극적인 개선을 보였습니다: 96% 가지치기 시 정확도가 SAM 단독 (48.36%) 에서 92.08%(SAM + WCR) 로 상승했습니다.
  • S2SAM (72.51% $\to$ 93.44%) 및 SVHN(9.74% $\to$ 95.99%) 에서도 유사한 경향이 관찰되었습니다.
  • ViT-B/32의 경우, WCR 은 CIFAR-100 에서 80% 가지치기 시 정확도를 SAM 기준 64.11% 에서 84.95% 로 향상시켰습니다.
• 의료 분할: LGG 뇌 MRI 데이터셋에서 WCR 은 85% 및 88% 가지치기 하에서 F1 점수를 크게 향상시키고 하우스도르프 거리를 감소시켰습니다. 특히, 88% 가지치기 시 SGD 기준선은 유효한 분할을 생성하지 못했으나, SGD + WCR 은 0.5280의 F1 점수를 달성했습니다. 정성적 결과는 WCR 이 기준선보다 병변 경계를 더 잘 보존함을 보여주었습니다.
• LLM 파인튜닝: LoRA 기반 파인튜닝에서 WCR 은 $\ell_1$ 및 DeepHoyer 정규화기보다 우수한 성능을 발휘했습니다. Qwen-2.5-1.5B 에서 WCR 은 비교된 정규화기 중 30% 가지치기 시 가장 높은 평균 정확도 (69.54%) 를 달성하여, 매개변수 효율적 파인튜닝 시나리오에서의 강건성을 입증했습니다.
• 일반화: 주 Hessian 고유값 분석은 WCR 이 손실 표면의 곡률을 낮춘다는 것을 나타내며, 이는 평탄한 최소점으로 최적화를 안내하여 일반화 개선과 상관관계가 있음을 시사합니다.

중요성 및 주장

본 논문은 훈련 중 가중치 크기 분포를 형성하는 것이 고도 희소성 하에서 강건성을 개선하는 효과적인 접근법이라고 주장합니다. WCR 의 중요성은 다음과 같은 능력에 있습니다:

• 큰 크기 (정보 있는) 와 작은 크기 (무시할 수 있는) 매개변수 간의 더 명확한 분리를 창출합니다.
• 크기 기반 가지치기가 기능적 영향이 최소화되면서 매개변수의 큰 비율을 제거할 수 있게 합니다.
• 독립적인 정규화기로 작동하거나 SAM 과 같은 손실 지형 평탄화 방법을 보완하는 구성 요소로 작동할 수 있습니다.

저자들은 WCR 이 아키텍처 변경이나 가지치기 후 재학습 없이 CNN, Transformer, 대규모 언어 모델 전반에 적용 가능하며, 공격적인 원샷 가지치기와 관련된 정확도 저하에 대한 실용적인 훈련 시간 솔루션을 제공한다고 결론지었습니다.