1. 문제 (왜 이걸 했는가)
모델을 경량화하면 항상 같은 문제가 발생한다.
“속도는 빨라졌는데 정확도가 떨어진다.”
YOLOv8m을 ModelOpt Pruning으로 압축한 결과는 다음과 같았다. 이전 글
- 25.9M → 16.5M (36% 감소)
- mAP50-95: 0.8392 → 0.7845 (-6.5%)
속도는 개선되었지만 성능 손실이 발생했다.
목표는 단순 복구가 아니라 다음과 같다.
- 압축된 모델이 원본을 넘도록 만든다
2. 접근 전략 (왜 Distillation인가)
Pruning 이후 모델은 구조적으로 정보가 손실된 상태다.
이 상태에서 성능을 회복하는 방법은 두 가지다.
- Fine-tuning
- GT label만 사용
- 표현 복구 한계 존재
- Knowledge Distillation
- Teacher 모델이 출력하는 클래스 확률 분포를 함께 학습
- 정답뿐 아니라 클래스 간 유사도 정보까지 전달
따라서 다음 전략을 선택했다.
GT + Teacher signal을 동시에 학습
3. 설계 핵심
Student는 기존 detection loss와 함께
Teacher의 확률 분포 기반 학습 신호를 동시에 학습한다.
4. KD Loss 구성
YOLOv8 detection head의 출력은 세 가지로 구성된다.
- classification (cls)
- bounding box (box)
- distribution (dfl)
이에 따라 KD loss도 세 가지 요소로 구성된다.
먼저 classification에 대해서는 Teacher의 클래스 확률 분포를 전달하기 위해 temperature scaling을 적용한 KL divergence를 사용한다.
\[L_{KD}^{cls} = T^2 \cdot KL\left( \text{softmax}\left(\frac{s}{T}\right) \,\|\, \text{softmax}\left(\frac{t}{T}\right) \right)\]이 과정에서 hard label이 아닌 클래스 간 관계까지 함께 학습된다.
bounding box는 분포를 맞추는 대신, 최종 좌표 공간에서 직접 정렬한다.
\[L_{KD}^{box} = \| b_s - b_t \|_1\]이는 위치 자체를 명확하게 맞추기 위한 선택이다.
DFL은 좌표의 불확실성을 나타내는 분포이기 때문에
분포 자체를 KL divergence로 정렬한다.
이를 통해 위치 예측의 sharpness까지 전달된다.
모든 anchor를 사용하는 것은 오히려 noise를 증가시킨다.
Teacher의 confidence가 일정 임계값(threshold) 이상인 anchor만 선택하여 유의미한 영역에 대해서만 KD를 적용한다.
\[\mathcal{M} = \{ a \mid \max(\sigma(t_{cls})) > \tau \}\]최종적으로 KD loss는 다음과 같이 구성된다.
\[L_{KD} = L_{KD}^{cls} + L_{KD}^{box} + L_{KD}^{dfl}\]5. Hyperparameter 설계 및 실험
KD 성능은 loss 구성보다 하이퍼파라미터에 더 크게 영향을 받는다.
이번 실험에서는 다음 세 가지 변수에 집중했다.
- $\alpha$: detection loss와 KD loss의 균형
- $T$: 확률 분포의 정보량 조절
- $\tau$: KD를 적용할 anchor 선택 기준
$\alpha$는 KD가 얼마나 강하게 작용할지를 결정한다.
- 너무 작으면 KD 효과가 거의 없다
- 너무 크면 detection loss를 압도하여 학습이 불안정해진다
$T$는 Teacher의 확률 분포를 얼마나 부드럽게 만들지 결정한다.
- 낮을수록 one-hot에 가까워짐
- 높을수록 클래스 간 관계 정보가 강조됨
$\tau$는 어떤 anchor를 KD에 사용할지 결정한다.
- 낮으면 background까지 포함되어 noise 증가
- 높으면 학습 데이터가 부족해짐
실험 결과, 다음 설정에서 가장 안정적인 성능을 얻었다.
\[\begin{aligned} \alpha &= 0.5 \\ T &= 3.0 \\ \tau &= 0.3 \end{aligned}\]이 설정에서는 detection loss와 KD loss가 균형을 이루면서
과도한 regularization 없이 성능이 향상되었다.
6. 결과
| 모델 | mAP50-95 | 속도 |
|---|---|---|
| Teacher | 0.8392 | 10.60ms |
| Pruned | 0.7845 | 8.05ms |
| Pruned + KD | 0.8679 | 8.05ms |
| Pruned + Feature KD | 0.8737 | 7.92ms |
Pruning으로 감소한 성능은 KD를 통해 복구되었고,
Feature KD까지 적용했을 때 Teacher를 초과하는 결과를 얻었다.
7. 왜 Teacher를 넘었는가
압축된 모델이 원본보다 더 높은 성능을 보이는 것은 직관적으로 이상해 보인다.
“정보를 줄였는데 성능이 더 좋아진다?”
이 현상은 Knowledge Distillation이 단순한 성능 복구가 아니라,
학습 신호 자체를 바꾸기 때문에 발생한다.
기존 학습은 GT label만 사용한다.
- 정답 클래스: 1
- 나머지 클래스: 0
즉, 모델은 “맞다 / 틀리다”만 학습한다.
반면 Teacher는 확률 분포를 제공한다.
- 특정 클래스는 0.9
- 유사한 클래스는 0.2
- 관련 없는 클래스는 0.01
이 분포에는 클래스 간 관계와 모호성이 포함되어 있다.
Student는 이 정보를 학습하면서 다음과 같은 변화가 발생한다.
-
Gradient가 더 풍부해진다
hard label에서는 0/1로 끊기지만,
soft label에서는 모든 클래스에 대해 gradient가 흐른다. -
Overfitting이 줄어든다
GT만 사용할 경우 특정 데이터에 과하게 맞춰지지만,
Teacher 분포는 이를 완화시키는 역할을 한다. -
표현 공간이 정렬된다
Feature KD를 통해 intermediate representation까지 맞춰지면서
backbone과 neck 수준에서 더 안정적인 feature를 학습한다.
결과적으로 Student는 단순히 Teacher를 복제하는 것이 아니라,
더 제한된 구조 안에서 더 일반화된 해로 수렴하게 된다.
이 과정에서 Teacher보다 더 좋은 성능이 실제로 나타났다.
8. Feature KD 역할
Feature KD는 FPN의 multi-scale 정보를 전달한다.
- P3: small object
- P4: medium object
- P5: large object
이를 통해 backbone과 neck 수준의 표현까지 개선된다.
9. 정리
이번 설계의 핵심은 다음 세 가지로 요약된다.
- KD loss: 성능 복구가 아니라 성능 초과를 유도
- alpha: detection loss와 KD 사이의 균형을 결정
- temperature / threshold: 정보의 질과 학습 대상 선택을 제어
- feature KD: representation 수준 정렬로 추가 성능 확보
10. 결론
Pruning은 단순히 모델을 줄이는 과정이 아니다.
정보를 압축하는 과정이다.
Knowledge Distillation은
그 압축된 모델에 다시 정보를 주입하는 과정이다.
결과적으로
- 더 작은 모델이
- 더 빠르게 동작하면서
- 더 높은 성능을 낼 수 있다
이번 실험에서는 실제로 Teacher 대비 성능 초과를 확인했다.