지식 증류 (knowledge distillation)
정의
지식 증류knowledge distillation란 크고 무거운 모델이 학습한 지식을 작고 가벼운 모델에게 옮겨, 작은 모델이 큰 모델에 버금가는 성능을 내도록 하는 방법론이다. 이때 크고 무거운 모델을 교사teacher, 작고 가벼운 모델을 학생student이라 한다. 모델 압축model compression의 한 갈래이다.
설명1 2
성능이 좋은 신경망은 대체로 크다. 매개변수가 수억 개에 달하는 모델이나 여러 모델의 예측을 평균 내는 앙상블ensemble은 뛰어난 성능을 보이지만, 계산량과 저장 공간을 많이 요구하기 때문에 휴대폰이나 임베디드 기기처럼 자원이 제한된 환경에 배포하기 어렵다. 학습에 쓰는 모델과 실제 배포에 쓰는 모델의 요구 조건이 서로 다른 것이다. 지식 증류는 이 간극을 메우기 위해, 지식을 잃지 않으면서 모델을 배포하기 좋은 형태로 바꾼다.
무거운 모델의 지식을 작은 모델로 옮긴다고 할 때, 먼저 떠오르는 방법은 교사 모델이 맞힌 정답, 곧 정답 레이블을 그대로 학생 모델의 학습에 쓰는 것이다. 교사 모델을 $\operatorname{Teacher}(x)$, 학생 모델을 $\operatorname{Student}(x; \theta)$라 할 때, 입력 $x$에 대한 두 모델의 출력 차이를 데이터 전체에 걸쳐 최소화하는 것이다.
$$ \mathcal{L}(\theta) = \mathbb{E}_{x}\left[ \bigl\lVert \operatorname{Teacher}(x) - \operatorname{Student}(x; \theta) \bigr\rVert^{2} \right] $$
그러나 교사의 출력을 이처럼 하나로 정해진 예측으로만 본다면, 결국 정답 레이블만으로 작은 모델을 직접 학습시키는 것과 크게 다르지 않다. 지식증류의 핵심은 학습된 모델이 갖는 지식을 매개변수의 구체적인 값과 동일시하지 않는 데 있다. 지식을 좀 더 추상적으로 바라보면, 그것은 입력 벡터를 출력 벡터로 옮기는 사상mapping이라고 할 수 있다. 여러 클래스를 분류하는 모델은 정답 클래스의 확률을 최대화하도록 학습되지만, 그 부산물로 오답 클래스들에도 저마다 다른 확률을 부여한다.
예컨대 BMW 자동차 사진을 분류할 때, 그것을 쓰레기차로 오인할 확률은 매우 작더라도 당근으로 오인할 확률보다는 훨씬 크다. 이처럼 오답들 사이의 상대적인 확률은 교사 모델이 데이터를 어떻게 일반화generalize하는지에 관한 풍부한 정보를 담고 있다. 지식 증류는 바로 이 정보를 학생 모델에게 전달한다. 정답 하나만 알려주는 대신, 교사 모델이 각 클래스에 부여한 확률 분포 전체를 학습 목표로 삼는 것이다. 이렇게 클래스 전체에 걸쳐 부드럽게 퍼진 확률 분포를 연성 목표soft target라 하고, 정답만 $1$이고 나머지는 $0$인 원-핫 형태의 정답 레이블을 경성 목표hard target라 한다. 연성 목표는 경성 목표보다 학습 사례 하나당 더 많은 정보를 담고 있으며, 사례들 사이의 기울기 분산도 작다. 따라서 학생 모델은 원래 교사 모델을 학습시킬 때보다 훨씬 적은 데이터와 높은 학습률로도 학습될 수 있다.
온도 소프트맥스와 손실함수
분류를 위한 신경망은 보통 마지막에 소프트맥스 층을 두어, 각 클래스에 대해 계산된 값인 로짓logit $z_{i}$를 확률 $q_{i}$로 변환한다.
$$ q_{i} = \dfrac{\exp(z_{i})}{\sum_{j} \exp(z_{j})} $$
지식 증류에서는 여기에 온도temperature $T$를 도입하여 다음과 같이 변형된 소프트맥스 함수를 사용한다.
$$ q_{i} = \dfrac{\exp(z_{i} / T)}{\sum_{j} \exp(z_{j} / T)} $$
$T = 1$이면 통상적인 소프트맥스와 같고, $T$를 크게 할수록 클래스에 대한 확률 분포가 더 부드러워진다.
위 그림은 같은 로짓에 온도만 달리 적용한 결과이다. $T = 1$에서는 정답 클래스 하나에 확률이 거의 몰려 있 지만, $T$를 키울수록 분포가 점점 고르게 퍼지면서 오답 클래스들에 매겨진 작은 확률들의 상대적인 크기가 드러난다.
교사 모델의 로짓을 $v_{i}$, 학생 모델의 로짓을 $z_{i}$라 하고, 온도 $T$에서 각각이 만드는 연성 확률을 $p_{i}$, $q_{i}$라 하자. 학생 모델은 다음 두 손실함수의 가중평균 으로 학습된다.
$$ \mathcal{L} = \alpha T^{2}\mathcal{L}_{\text{soft}} + (1 - \alpha)\mathcal{L}_{\text{hard}} $$
- $\mathcal{L}_{\text{soft}}$는 온도 $T$에서 교사의 연성 목표 $p_{i}$와 학생의 연성 예측 $q_{i}$ 사이의 크로스 엔트로피이다.
- $\mathcal{L}_{\text{hard}}$는 온도 $T = 1$에서 정답 레이블과 학생 예측 사이의 크로스 엔트로피이다.
- $\alpha$는 두 항의 비중을 조절하는 가중치이다.
온도 $T$는 연성 목표가 얼마나 부드러운지를 조절한다. MNIST처럼 교사 모델이 거의 언제나 정답을 매우 높은 확신으로 맞히는 과제에서는, 학습된 함수에 관한 정보 대부분이 아주 작은 확률들의 비율 속에 숨어 있다. (위 그림 참고) 예를 들어 어떤 숫자 $2$는 $3$일 확률이 $10^{-6}$, $7$일 확률이 $10^{-9}$로 매겨질 수 있는데, 이 값들은 어떤 $2$가 $3$을 닮았고 어떤 $2$가 $7$을 닮았는지를 알려주는 귀중한 정보이다. 하지만 확률이 $0$에 너무 가까워 통상적인 크로스 엔트로피에는 거의 영향을 주지 못한다. 온도를 높이면 이 작은 확률들이 도드라져, 그 안에 담긴 유사도 구조가 학습에 반영된다. 위 그림에서 $T$를 $1$에서 $3$으로 올리자, $3$일 확률이 $10^{-6}$에서 $0.1$로, $7$일 확률이 $10^{-9}$에서 $0.01$로 커져, 두 클래스 사이의 상대적인 크기가 훨씬 잘 드러난다. 그러나 $T = 10$과 같이 온도를 너무 높게 설정하면, 정답 클래스와 오답 클래스의 구분이 흐려진다.
한편 연성 목표가 만드는 기울기의 크기는 $1/T^{2}$에 비례하여 작아진다. 그래서 경성 목표와 함께 쓸 때는 연성 항에 $T^{2}$을 곱해 두 항의 상대적인 기여가 온도에 상관없이 대체로 유지되도록 한다. 이것이 손실함수에서 $\mathcal{L}_{\text{soft}}$ 앞에 $T^{2}$이 붙는 이유이다.
일반화
지식증류는 사실 학생모델의 로짓을 교사모델의 로짓에 맞추는 것의 일반화라고 볼 수 있다. 학생 모델이 예측한 확률을 $q_{i} = e^{z_{i}/T} / \sum_{j} e^{z_{j}/T}$, 교사 모델이 예측한 확률을 $p_{i} = e^{v_{i}/T} / \sum_{j} e^{v_{j}/T}$라 하자. 그러면 $i$번째 클래스에 대한 학생의 로짓 $z_{i}$에 기여하는 크로스 엔트로피 기울기는 다음과 같다.
$$ \begin{align*} \dfrac{\partial C}{\partial z_{i}} &= \dfrac{\partial}{\partial z_{i}} \left( -\sum_{k} p_{k} \log q_{k} \right) \\ &= \dfrac{\partial}{\partial z_{i}} \left[ -\sum_{k} p_{k} \left( \dfrac{z_{k}}{T} - \log \sum_{j} e^{z_{j}/T} \right) \right] \\ &= \dfrac{\partial}{\partial z_{i}} \left( -\dfrac{1}{T}\sum_{k} p_{k} z_{k} + \log \sum_{j} e^{z_{j}/T} \right) \\ &= -\dfrac{p_{i}}{T} + \dfrac{1}{T} \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} \\ &= \dfrac{1}{T}\left( \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} - p_{i} \right) \\ &= \dfrac{1}{T}\left( \dfrac{e^{z_{i}/T}}{\sum_{j} e^{z_{j}/T}} - \dfrac{e^{v_{i}/T}}{\sum_{j} e^{v_{j}/T}} \right) \end{align*} $$
온도 $T$가 로짓의 크기에 비해 충분히 크면, 지수함수의 매클로린 급수 전개에 의해, $e^{x/T} \approx 1 + x/T$로 근사할 수 있으므로 다음을 얻는다.
$$ \dfrac{\partial C}{\partial z_{i}} \approx \dfrac{1}{T}\left( \dfrac{1 + z_{i}/T}{N + \sum_{j} z_{j}/T} - \dfrac{1 + v_{i}/T}{N + \sum_{j} v_{j}/T} \right) $$
여기서 $N$은 클래스의 개수이다. 모든 클래스에 대한 로짓의 평균을 $0$이라 가정하면 $(\sum_{j} z_{j} = \sum_{j} v_{j} = 0)$, 위 식은 다음으로 간단해진다.
$$ \dfrac{\partial C}{\partial z_{i}} \approx \dfrac{1}{N T^{2}}(z_{i} - v_{i}) $$
곧 온도가 높은 극한에서 지식 증류는 $\frac{1}{2}(z_{i} - v_{i})^{2}$을 최소화하는 것, 즉 학생과 교사의 로짓 차이를 줄이는 것과 같아진다. 반대로 온도가 낮으면 평균보다 훨씬 작은(매우 음수인) 로짓에는 거의 신경 쓰지 않게 된다. 이러한 로짓은 교사 모델의 학습 과정에서 거의 제약을 받지 않아 잡음일 수 있으므로, 이를 무시하는 편이 도움이 되기도 한다. 학생 모델이 교사의 지식을 모두 담기에 너무 작을 때는 중간 정도의 온도가 가장 잘 작동하는 것으로 알려져 있다.