Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

기술적 요약: 합리주의자로부터 배우기: 해석 가능한 중간 추론의 증류

문제 제기
심층 신경망 (DNN) 은 고위험 분야에서 탁월한 성능을 달성했지만, 여전히 불투명한 '블랙박스'로 남아 신뢰 문제를 야기합니다. 사후 설명 도구 (예: LIME, SHAP) 가 존재하지만, 종종 충실한 설명을 제공하지 못합니다. 추론 추출 (RE) 은 생성기가 특징 (추론) 을 선택하고 예측기가 오직 해당 특징만을 기반으로 결정을 내리는 '선택 - 예측' 아키텍처를 제공하여 충실성을 보장합니다. 그러나 RE 모델의 학습은 '닭이 먼저냐 달걀이 먼저냐'라는 딜레마에 직면합니다. 생성기는 중요한 특징을 식별하기 위해 예측기가 필요하고, 예측기는 생성기가 깨끗한 입력을 제공해야 하기 때문입니다. 이 문제는 경량화되거나 능력이 낮은 신경망 (학생) 을 사용할 때 더욱 악화되는데, 외부 지도 없이 효과적인 추론을 학습하고 높은 예측 성능을 달성하는 데 어려움을 겪기 때문입니다.

방법론: REKD
저자들은 REKD(지식 증류를 통한 추론 추출) 를 제안합니다. 이는 학생 RE 모델이 자체 최적화 외에도 '합리주의자'인 교사 RE 모델로부터 학습하는 프레임워크입니다. 이 방법은 신경 모델에 구애받지 않아 어떤 블랙박스 아키텍처든 백본으로 사용할 수 있습니다.

1. 추론 추출 (RE) 구성 요소:

   • 이산적인 특징 선택을 미분 가능하게 만들기 위해 Straight-Through Gumbel-Softmax 추정기를 사용합니다.
   • 생성기는 특징 선택을 위한 정규화되지 않은 로그트 (logits) 를 생성하며, 온도 매개변수 $\tau$를 사용하여 Gumbel-Softmax 를 통해 샘플링됩니다.
   • 이진 마스크는 argmax 를 통해 유도되며 (역전파 중에는 항등 함수로 근사됨), 선택되지 않은 특징을 0 으로 만듭니다.
   • 예측기는 마스크된 추론만을 기반으로 분류합니다.
   • 목적 함수는 작업 손실 (교차 엔트로피) 과 선택 손실 (목표 추론 비율에서의 편차를 패널티) 을 결합합니다.

2. 지식 증류 (KD) 구성 요소:

   • 학생은 교사의 추론(특징 선택 분포) 과 예측(클래스 확률) 에서 학습합니다.
   • 생성기 증류: 모든 특징에 걸쳐 학생과 교사의 Gumbel-Softmax 분포 간의 Kullback-Leibler (KL) 발산을 최소화합니다.
   • 예측기 증류: 학생과 교사의 온도 스케일링 출력 분포 간의 KL 발산을 최소화합니다.

3. 동기화된 온도 어닐링:

   • 주요 혁신은 KD 온도와 Gumbel-Softmax 어닐링 스케줄러의 동기화입니다.
   • 온도 $\tau$는 높은 초기 값 ($\tau_0$) 에서 낮은 최종 값 ($\tau_K$) 으로 지수적으로 감소합니다.
   • 이는 암묵적인 커리큘럼 학습 전략을 생성합니다:
     • 초기 단계 (높은 $\tau$): 학생은 교사의 넓은 분포에서 '부드러운' 지식을 흡수하여 탐색을 안내합니다.
     • 후기 단계 (낮은 $\tau$): 분포가 날카로워짐에 따라 학생은 교사의 '단단한' 고신뢰도 특징 선택을 모방하도록 전환되어 이산화 과정에서 정밀성을 강제합니다.
   • 총 손실은 RE 목적 함수와 KD 목적 함수의 가중 합입니다.

주요 기여

1. 딜레마의 식별: 이 논문은 RE 의 '닭이 먼저냐 달걀이 먼저냐' 학습 어려움을 공식적으로 식별하고 분석하며, 특히 경량 학생 모델을 어떻게 방해하는지 강조합니다.
2. REKD 프레임워크: Gumbel-Softmax 어닐링의 내재적 커리큘럼을 활용하여 해석 가능한 특징 선택 지식을 교사로부터 학생으로 전달하는 증류 프레임워크를 제안합니다.
3. 실증적 검증: BERT 및 Vision Transformer (ViT) 변형을 사용하여 언어 (IMDB) 및 비전 (CIFAR 10/100) 작업 전반에서 접근법의 타당성을 입증하여 학생 모델의 성능 향상을 보여줍니다.

실험 결과
IMDB(감성 분석), CIFAR 10, CIFAR 100(이미지 분류) 에서 BERT 및 ViT 아키텍처를 사용하여 실험을 수행했습니다.

• 성능 향상: REKD 는 표준 RE 학습에 비해 학생 RE 모델의 예측 정확도를 크게 향상시킵니다.
  • CIFAR 10: ViT Small 은 0.889(RE) 에서 0.968(REKD) 로 향상되었고, ViT Tiny 는 0.797 에서 0.936 으로 향상되었습니다.
  • CIFAR 100: ViT Small 은 0.779 에서 0.845 로 향상되었고, ViT Tiny 는 0.645 에서 0.777 로 향상되었습니다.
  • IMDB: BERT Small 은 0.881 에서 0.906 으로 향상되었고, BERT Mini 는 0.863 에서 0.892 로 향상되었습니다.
• 교사 모델 초월: 일부 경우 (예: CIFAR 10 및 CIFAR 100 의 ViT Small) 에는 학생 REKD 모델이 교사 모델 자체보다 높은 평균 정확도를 달성하여, 증류 과정이 학생이 교사의 추론을 활용하는 능력을 효과적으로 정제함을 시사합니다.
• 강건성: 이 방법은 서로 다른 무작위 시드 간 성능 변동을 줄입니다.
• 애블레이션: 저자들은 성능 향상이 RE 와 KD 의 결합에서 비롯되며, 교사의 추론이 학생이 우연한 상관관계를 학습하는 것을 방지하는 강력한 정규화 신호를 제공한다고 지적합니다.

의의 및 주장
이 논문은 REKD 가 특히 리소스 제약 시나리오 (예: 모바일 장치) 에서 해석 가능 모델의 학습 불안정성에 대한 실용적인 해결책을 제공한다고 주장합니다. 문제를 '합리주의자'로부터 학습하는 것으로 프레임함으로써, 이 방법은 초보자가 전문가가 제공하는 검증되고 해석 가능한 지식으로부터 학습하는 인간의 학습 과정과 부합합니다.

저자들은 평가 지표가 인간 주석과의 정렬이 아닌 추론 제약 하의 예측 성능에 초점을 맞춘다고 강조합니다. 그들은 데이터에 편향이 포함되어 있더라도 데이터에 대한 충실성이 인간이 납득할 수 있는 것보다 추론의 질에 대한 더 객관적인 척도라고 주장하며, 이는 모델의 실제 추론을 드러내기 때문입니다. 이 방법은 더 큰 모델의 '추론'을 작은 모델이 물려받을 수 있게 함으로써, 대규모 모델의 예측 능력을 희생하지 않으면서도 고위험 분야에서 DNN 을 더 신뢰할 수 있게 만드는 한 걸음으로 제시됩니다.