작은 AI 모델에 거대한 AI가 아는 모든 것을 가르치는 방법

목차

1. 소개
2. 어둠의 지식
3. 증류 절차
4. MNIST에 대한 실험
5. MNIST에 대한 추가 실험

1. 서론

이 글에서는 AI의 지식 증류 과정을 살펴보겠습니다. 지식 증류 과정이 일반적으로 어떻게 작동하는지, 그 중요성과 사용 이유에 대해 알아보겠습니다.

매우 큰 데이터 세트를 통해 훈련되어 데이터에서 구조를 추출한 더 큰 모델이나 모델 앙상블의 지식을 성능 저하 없이 단일 작은 모델로 압축하고 전달하려면 어떻게 해야 할까요?

하지만 왜 우리는 이것을 하고 싶어할까요? 더 큰 모델이나 앙상블 모델이 이미 테스트 데이터에서 좋은 결과를 제공하고 있는데 왜 더 작은 모델이 필요한 걸까요?

훈련 시간에 우리는 일반적으로 대규모/앙상블 모델을 훈련하는데, 주된 목표는 매우 큰 데이터 세트에서 구조를 추출하는 것이기 때문입니다. 또한 훈련 시간에 드롭아웃, 데이터 증강과 같은 많은 것을 적용하여 이러한 대규모 모델에 모든 종류의 데이터를 공급할 수도 있습니다.

하지만 예측 시점에 우리의 목표는 완전히 다릅니다. 우리는 가능한 한 빨리 결과를 얻고 싶습니다. 따라서 더 큰 모델/앙상블을 사용하는 것은 매우 비용이 많이 들고 많은 수의 사용자에게 배포하는 데 방해가 됩니다. 따라서 이제 질문은 이 더 큰 모델의 지식을 쉽게 배포할 수 있는 더 작은 모델로 압축하는 방법입니다.

구글의 Geoffrey Hinton, Oriol Vinyals, Jeff Dean은 논문을 통해 이 지식을 더 작은 모델로 옮기기 위해 증류 라는 다른 종류의 훈련을 제안했습니다. 이것은 Distill BERT 구현에서 huging face가 사용한 것과 동일한 기술입니다.

만약 우리가 이 작은 모델을 큰 모델과 같은 방식으로 일반화 하도록 훈련할 수 있다면, 이런 방식으로 훈련된 이 작은 모델은 같은 데이터에서 일반적인 방식으로 훈련된 작은 모델보다 훨씬 더 나은 성과를 낼 것입니다. 이것이 증류의 주요 원리 중 하나입니다.

2. 어둠의 지식

일반적으로, 많은 수의 클래스를 구별하도록 학습하는 모델인 머신 러닝에서 주요 학습 목표는 정답의 평균 로그 확률을 최대화하는 것입니다. 예를 들어, 이미지를 1 또는 2 또는 … 9로 분류하는 것이 목표인 MNIST 데이터 세트의 예를 들어 보겠습니다. 따라서 실제 이미지가 2인 경우 모든 모델의 목표는 P(its 2/image) (특정 이미지가 주어진 이미지에서 2일 확률로 읽을 수 있음)를 최대화하는 것입니다. 하지만 모델은 또한 모든 오답에 대한 확률을 제공하는데, 이러한 확률이 매우 작더라도 그 중 일부는 다른 것보다 훨씬 큽니다. 요점은 이러한 확률이 작더라도 오답의 상대적 확률은 모델이 어떻게 일반화할 수 있는지에 대해 많은 것을 알려준다는 것입니다. 이를 이해하기 위해 아래 예를 살펴보겠습니다.

위 그림에서 이 버전의 2는 3이 될 확률 10-6, 7이 될 확률 10-9를 받았지만 다른 버전에서는 그 반대일 수 있습니다. 이는 데이터에 대한 풍부한 유사성 구조를 정의하는 귀중한 정보입니다(즉, 어떤 2가 3처럼 보이고 어떤 2가 7처럼 보이는지 말해줍니다). 그러나 확률이 0에 매우 가깝기 때문에 전송 단계에서 교차 엔트로피 비용 함수에 미치는 영향은 거의 없습니다.

하지만 증류 과정으로 넘어가기 전에, 모델이 실제로 출력 확률을 어떻게 생성했는지에 대해 시간을 할애해 보겠습니다. 여기서 소프트맥스 활성화가 등장합니다. 모델 처리의 마지막 단계는 소프트맥스이고, 이 구성 요소가 출력 확률을 제공합니다. 소프트맥스에 대한 입력은 로짓이라고 하며, NN의 최종 계층을 숨겨진 단위 수 = 분류하려는 클래스 수가 되도록 설계합니다.

소프트맥스를 계산하는 공식은 다음과 같습니다.

위의 방정식은 각 i 에 대한 확률을 제공하고 전체 확률의 합은 1 과 같습니다. 학습 시간 동안 모든 단일 학습 예제에 대한 손실은 이러한 소프트맥스 확률을 하드 타겟(레이블)과 비교하고 역전파 계수를 사용하여 손실이 최소가 될 때까지 업데이트하여 계산됩니다.

위에서 본 것처럼 이 소프트맥스는 참 레이블에 높은 확률을 주고 틀린 레이블에 낮은 확률을 줍니다. 또한 우리는 비록 작더라도 틀린 답변의 확률은 모델이 일반화하는 데 도움이 되는 많은 정보가 숨겨져 있음을 알 수 있습니다. 우리는 이것을 다크 지식 이라고 부릅니다.

3. 증류 절차

논문에 따르면, 큰 모델의 일반화 능력을 작은 모델로 옮기는 가장 좋은 방법은 복잡한 모델에서 생성된 클래스 확률을 작은 모델을 훈련하기 위한 소프트 타겟 으로 사용하는 것입니다.

그 과정은 다음과 같습니다.

• 더 큰 모델을 학습하는 데 사용된 원래의 학습 세트를 가져온 다음, 해당 학습 데이터를 더 큰 모델에 전달하여 다양한 클래스에 대한 소프트맥스 확률을 얻습니다. 위에서 볼 수 있듯이, 참 레이블은 높은 확률을 얻고, 부정확한 레이블은 낮은 확률을 얻습니다. 하지만 우리는 이러한 낮은 확률에 많은 정보가 숨겨져 있음을 보았습니다. 따라서 이러한 확률의 중요성을 확대하기 위해 논문의 저자는 소프트맥스를 통과하기 전에 모든 로짓을 나누기 위해 Temperature(T)라는 변수를 사용했습니다. 이렇게 하면 클래스에 대한 더 부드러운 확률 분포가 생성됩니다. 아래에서 볼 수 있습니다.

temperature(T)로 소프트맥스를 적용한 결과는 소프트 타겟이라고 합니다. 이 과정은 저자들이 증류 라고 부르는 것입니다. 온도를 높여 물 속의 불순물을 제거하는 것과 유사합니다.

• 대규모 모델에서 학습된 함수에 대한 많은 정보는 소프트 타겟의 매우 작은 확률 비율에 들어 있습니다.
• 일부 용어:

1. 소프트 타겟 - 소프트맥스 동안 온도 T가 적용된 후 대형 모델의 출력
2. 소프트 예측 - 소프트맥스 동안 온도 T가 적용된 후 더 작은 모델의 출력
3. 하드 예측 - 온도 T = 1일 때 더 작은 모델의 출력(일반 소프트맥스)
4. 실제 목표 - 훈련 세트의 실제 목표

아래는 전체 교육 과정의 흐름도입니다.

따라서 작은 모델의 학습 과정에는 손실 함수가 2개 있습니다. 첫 번째 손실 함수는 소프트 예측과 소프트 타겟을 모두 취하고 교차 엔트로피 손실 함수입니다. 이는 소프트 타겟을 일치시키려고 시도하여 일반화 능력을 대형 모델에서 소형 모델로 이전하는 방식입니다. 이 손실 함수의 경우 두 소프트맥스는 모두 'T'의 온도를 사용합니다.

저자들은 또한 작은 모델을 사용하여 진정한 타겟을 일치시키는 것이 도움이 된다는 것을 발견했습니다. 이것은 두 번째 비용 함수에 통합되어 있습니다. 최종 비용은 하이퍼 매개변수 알파와 베타를 가진 이 두 비용 함수의 가중 평균입니다.

4. MNIST 실험

저자는 이 접근 방식을 테스트하기 위해 MNIST 데이터 세트를 사용했습니다. 그들은 이를 위해 두 가지 아키텍처를 사용했는데, 이는 중간 계층의 숨겨진 단위 수에서만 다릅니다. 저자는 두 경우 모두 2개의 숨겨진 계층 신경망을 사용했습니다.

1. 784 -> 800 -> 800 -> 10으로 볼 수 있는 더 작은 모델입니다(여기서 784는 이미지의 펼쳐진 차원이고, 800은 RELU 활성화를 사용한 숨겨진 유닛의 수이고, 10은 예측하는 클래스의 수입니다). 이 모델은 정규화 없이 146개의 테스트 오류를 보였습니다.
2. 784 -> 1200 -> 1200 -> 10으로 볼 수 있는 더 큰 모델(여기서 784는 이미지의 펼쳐진 차원이고, 1200은 RELU 활성화를 사용한 숨겨진 유닛의 수이고, 10은 예측하는 클래스의 수임). 이 모델은 드롭아웃, 가중치 제약, 지터링 입력 이미지를 사용하여 MNIST에서 학습되었으며, 이 넷은 67개의 테스트 오류를 달성했습니다.

더 큰 모델에서 나타난 이러한 개선 효과를 작은 모델로 옮길 수 있을까요?

저자들은 이제 큰 네트에서 얻은 소프트 타겟과 실제 타겟을 모두 사용했는데, 이때 드롭아웃이나 이미지 지터링은 없었습니다. 즉, 작은 네트가 큰 네트에서 생성된 소프트 타겟을 20의 온도에서 매칭하는 추가 작업을 추가하여 정규화되었고 그 결과는 다음과 같습니다.

784 -> 800 -> 800 -> 10을 사용하여 74개의 테스트 오류

이는 소프트 타겟이 작은 모델에 많은 지식을 전달할 수 있음을 보여주는데, 여기에는 변환된 학습 데이터에서 학습한 일반화 방법에 대한 지식도 포함됩니다. 다시 말해, 입력을 변환하여 얻은 이점은 작은 네트에 대한 입력을 변환하지 않더라도 작은 네트에 전달됩니다.

입력을 다른 변환으로 변환하면 모델이 훨씬 더 잘 일반화된다는 것은 잘 알려진 사실이며, 우리의 경우 일반화 방법에 대한 정보가 Dark knowledge에 나타나고 이것은 soft targets에 숨겨져 있습니다. 이 정보 중 어느 것도 True targets에 없습니다. 따라서 soft targets의 정보를 사용함으로써 우리의 작은 net은 훨씬 더 나은 성능을 보이고 있습니다.

소프트 타겟을 사용하는 큰 넷은 '무엇이 무엇과 비슷한가'를 학습한 유사성 메트릭을 학습했으며 이러한 지식 전달을 통해 작은 넷에 '무엇이 무엇과 비슷한가'를 알려줍니다.

위의 MNIST에 대한 모든 실험은 아래와 같이 요약됩니다.

5. MNIST에 대한 추가 실험

또한 저자는 증류를 사용하여 더 작은 모델을 훈련할 때 숫자 3의 예를 생략하려고 했습니다. 따라서 작은 모델의 관점에서 볼 때 3은 본 적이 없는 신화적인 숫자입니다. 더 작은 모델은 훈련 중에 3을 본 적이 없으므로 테스트 세트에서 3을 만나면 많은 오류가 발생할 것으로 예상합니다. 그럼에도 불구하고 증류된 모델은 테스트 세트에서 1010개의 3에 대한 133개의 오류를 포함하여 206개의 테스트 오류만 발생했습니다. 이는 증류 중에 큰 모델의 일반화 기능이 작은 모델로 이전되었으며 이로 인해 대부분의 경우 작은 모델이 3을 올바르게 예측한다는 것을 분명히 보여줍니다.

이 이야기의 교훈은 다음과 같습니다.

입력 이미지를 변환하면 일반화가 크게 향상됩니다. 타겟을 변환하면 비슷한 효과가 있고 어딘가에서 소프트 타겟을 얻을 수 있다면 더 작은 모델로 동일한 성능을 얻을 수 있으므로 훨씬 저렴합니다.

참고문헌:

1. 힌튼, 제프리, 오리올 비냔얼스, 제프 딘. "신경망에서 지식을 증류하다."
2. 인텔랩스에 의한 지식의 정제.