🤖 machine learning

Distill on a Diet: Efficient Knowledge Distillation via Learnable Data Pruning

이 논문은 영향 함수(influence functions)와 학습 가능한 베타 분포 기반 샘플링 정책을 활용하여 데이터 프루닝을 최적화함으로써, 현저히 줄어든 데이터와 연산량으로 훈련된 학생 모델이 전체 데이터셋으로 증류된 모델보다 더 뛰어난 성능을 낼 수 있도록 하는 효율적인 지식 증류 프레임워크인 IF-Beta를 제안한다.

원저자: Yifan Wu, Yiqi Wang, Xichen Ye, Wenjing Yan, Xiaoqiang Li, Cheng Jin, Xiangyu Yue, Weizhong Zhang

게시일 2026-06-25

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yifan Wu, Yiqi Wang, Xichen Ye, Wenjing Yan, Xiaoqiang Li, Cheng Jin, Xiangyu Yue, Weizhong Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

큰 그림: "식단 조절"을 통한 학생 가르치기

당신이 마스터 셰프(스승)가 되어 견습생(제자)에게 완벽한 요리를 만드는 법을 가르치고 있다고 상상해 보세요. 보통 제자를 가르칠 때, 당신은 거대한 요리책에 담긴 모든 레시피를 연습하게 만듭니다. 이는 시간이 너무 오래 걸리고, 많은 연료(컴퓨팅 파워)를 소모하며, 매우 지치는 일입니다.

**지식 증류(Knowledge Distillation, KD)**는 마스터 셰프가 제자를 지도하는 과정입니다. 목표는 제자가 마스터만큼 잘할 수 있도록 만들되, 훨씬 더 빠르고 적은 장비로도 가능하게 하는 것입니다.

하지만 이 논문은 한 가지 문제를 지적합니다. 비록 최종적인 제자는 작고 효율적이지만, 훈련 과정 자체는 제자가 요리책 전체를 읽어야 하기 때문에 여전히 거대하다는 점입니다. 저자들은 다음과 같이 질문합니다. "만약 우리가 제자에게 가장 중요한 레시피들로만 구성된 정교하게 선택된 '식단'을 제공할 수 있다면 어떨까? 그러면 제자가 똑같이 잘 배울 수 있으면서도, 더 빠르고 적은 노력으로 가능하지 않을까?"

이것이 바로 Distill on a Diet의 핵심 아이디어입니다.

기존 방식의 문제점

이 논문이 나오기 전, 사람들은 제자에게 줄 '최고의 레시피'를 고르기 위해 두 가지 주요 방법을 사용했는데, 둘 다 결함이 있었습니다.

"재학습(Re-Training)" 방식: 어떤 레시피가 어렵거나 쉬운지 알기 위해서는, 먼저 제자가 요리책 전체를 통해 요리하는 과정을 지켜봐야 합니다. 이는 시간을 아끼려는 목적에 어긋납니다! 이는 마치 어떤 레시피를 보여줄지 결정하기 위해 탐정을 고용해 한 달 동안 제자를 감시하게 하는 것과 같습니다.
"경험칙(Rule of Thumb)" 방식: 어떤 사람들은 "중간 난이도의 레시피를 골라라"와 같은 단순한 규칙을 사용합니다. 하지만 이는 경직되어 있습니다. 때로는 "중간" 레시피가 최선이 아닐 수도 있고, 아주 쉽거나 아주 어려운 레시피가 골고루 섞여 있어야 할 때도 있습니다. 이는 마치 신체의 실제 필요를 고려하지 않고 "무조건 사과 3개와 바나나 2개를 먹으라"고 강요하는 엄격한 식단 계획과 같습니다.

해결책: IF-Beta

저자들은 별도의 재학습 없이도 완벽한 "식단"을 선택할 수 있는 IF-Beta라는 새로운 시스템을 제안합니다. 이는 두 가지 스마트한 아이디어를 결합한 것입니다.

1. "영향 함수(Influence Function)" (수정구슬)

당신에게 수정구슬이 있어서 이렇게 말해준다고 상상해 보세요. "만약 요리책에서 이 특정 레시피를 제거한다면, 제자의 요리 실력이 나빠질까 아니면 좋아질까?"

논문에서는 이를 **영향 함수(Influence Function, IF)**라고 부릅니다.

기존 방식: 이 수정구슬은 대개 고장 났거나 사용하기에 너무 비용이 많이 들었습니다.
새로운 방식: 저자들은 이 수정구슬을 효율적으로 작동시키는 방법을 찾아냈습니다. 그들은 **평탄 검증 최솟값(Flat Validation Minima, FVM)**이라는 기술을 사용합니다. 이것은 스승의 직관을 "날카롭게" 만드는 과정이라고 생각하면 됩니다. 스승의 지식 지형을 더 "평탄하고" 안정적으로 만듦으로써, 수정구슬은 제자가 직접 고군분투하는 과정을 지켜보지 않고도 어떤 데이터 포인트가 진정으로 중요한지 정확하게 예측할 수 있습니다.

2. "베타 정책(Beta Policy)" (유연한 셰프)

수정구슬을 통해 어떤 레시피가 중요한지 알게 되었다면, 이제 그것들을 어떻게 고를까요?

기 기존 방식: 딱딱한 자를 사용했습니다. "상위 10%의 어려운 것을 가져와" 또는 "점수 5점에서 7점 사이의 것을 가져와"와 같은 식입니다. 이는 배가 얼마나 고픈지와 상관없이 "사과 3개와 바나나 2개를 정확히 먹어라"라고 하는 것과 같습니다.
새로운 방식: 저자들은 **베타 정책(Beta Policy)**을 사용합니다. 유연하고 늘어나는 그물망을 상상해 보세요. 딱딱한 자 대신, 이 그물은 최고의 데이터를 잡기 위해 스스로의 모양을 바꿀 수 있습니다. 상황에 따라 주로 어려운 예시를 잡도록 늘어나거나, 쉬운 예시를 잡거나, 혹은 완벽한 혼합을 잡을 수 있습니다. 이는 상황에 따라 적응하는 "학습 가능한" 식단입니다.

작동 원리 (두 단계의 댄스)

이 시스템은 완벽한 식단을 찾기 위해 두 단계의 춤처럼 작동합니다.

내부 루프 (연습 단계): 시스템은 작은 임시 데이터 식단을 통해 제자가 학습하는 과정을 빠르게 시뮬레이션합니다. 이를 매우 빠르게 만들기 위해, 전체 뇌를 새로 훈련하는 대신, 얼어붙은(frozen) 스승의 지식 위에 간단한 "선형 헤드(linear head, 마치 빠른 메모와 같은 것)"만을 훈련시킵니다.
외부 루프 (코치 단계): 시스템은 이 임-시 식단에서 제자가 얼마나 잘했는지 확인합니다. 만약 식단이 좋았다면, "코치"(베타 정책)는 더 좋은 데이터를 잡기 위해 그물의 모양을 조정합니다. 만약 식단이 나빴다면, 코치는 그물의 모양을 다시 바꿉니다.

이 과정은 매우 빠르게 일어나므로, 전체 훈련이라는 무거운 비용 없이도 최적의 데이터 부분 집합을 찾을 수 있습니다.

결과: 적은 음식, 더 나은 요리

이 논문은 유명한 이미지 데이터셋(CIFAR, ImageNet 등)을 통해 테스트를 진행했습니다. 결과는 놀라웠습니다.

전체 데이터보다 우수함: 많은 경우, 단 **50%~70%**의 데이터만 사용해 훈련한 제자(IF-Beta 식단 사용)가 100%의 데이터로 훈련한 제자보다 실제로 더 높은 성능을 보였습니다.
더 빠르고 저렴함: 더 적은 데이터를 사용했기 때문에, 훈련 시간과 컴퓨팅 파워가 현저히 줄어들었습니다.
경쟁 상대를 압도함: IF-Beta는 값비싼 재학습이나 경직된 규칙을 사용하는 다른 모든 방법들을 이겼습니다.

핵심 요약

이 논문은 데이터의 양보다 질이 중요하다고 주장합니다. 인간이 도서관의 모든 책을 읽는 대신 몇 명의 고품질 멘토와 예시를 통해 기술을 더 빨리 배울 수 있는 것처럼, 머신러닝 모델도 훈련 데이터의 "식단"을 정교하게 큐레이션하면 더 잘 배울 수 있습니다.

저자들은 이를 **"Distill on a Diet"**라고 부릅니다. 왜냐로 제자 모델에게 전체 데이터셋을 강제로 먹이는 대신, 더 작고 영양가 높으며 정교하게 선택된 양의 데이터를 제공함으로써, 제자가 전체를 다 공부했을 때보다 더 강하고 빠르게 성장할 수 있게 하기 때문입니다.

기술 요약: Distill on a Diet: 효율적인 지식 증류를 위한 학습 가능한 데이터 프루닝

문제 정의

지식 증류(Knowledge Distillation, KD)는 거대한 교사(Teacher) 모델을 압축된 학생(Student) 모델로 압축하여 추론 효율성을 높이는 표준적인 기술입니다. 그러나 학습 과정 자체는 여전히 계산 비용이 많이 드는데, 이는 일반적으로 학생들이 거대한 교사의 가이드 아래 전체 데이터셋을 대상으로 증류되기 때문입니다. 데이터 프루닝(데이터의 일부를 선택하는 기법)은 이러한 비용을 줄일 수 있는 유망한 방안을 제시하지만, 기존 방법들은 KD 맥락에서 두 가지 결정적인 장애물에 직면해 있습니다:

점수 추정의 효율성-효과성 트레이드오프: EL2N, Forgetting, 또는 AUM과 같은 고충실도 지표는 학습 역학(예: 에포크별 예측값)에 의존하는데, 이는 사전 학습된 교사만 제공된 상황에서는 사용할 수 없습니다. 이러한 점수를 계산하려면 교사를 다시 학습시키는 비싼 과정이 필요하며, 이는 프루닝을 통한 효율성 이득을 상쇄합니다. 반대로, 단순한 크로스 엔트로피(Cross-Entropy, CE) 손실은 계산 비용은 저렴하지만 신뢰도가 낮습니다. 과잉 매개변수화된 네트워크는 난도가 높거나 레이블이 잘못된 샘플에 과적합될 수 있어, 낮은 손실 값이 실제 샘-플의 난도를 반영하지 못할 수 있기 때문입니다.
경직되고 최적이 아닌 샘플링: 기존의 프루닝 전략은 고정된 휴리스틱(예: top-k 선택, 고정된 컷오프를 가진 층화 추출, 또는 슬라이딩 윈도우)에 의존하는 경우가 많습니다. 이러한 경직된 규칙은 다양한 KD 설정의 특정 데이터 분포와 교사-학생 역학에 적응하지 못하여 최적이 아닌 부분 집합을 생성하게 됩니다.

방법론: IF-Beta

저자들은 위의 과제들을 해결하기 위해 영향 함수(Influence Functions, IF)와 학습 가능한 확률적 샘플링 정책을 결합한 IF-Beta 프레임워크를 제안합니다.

1. 사후(Post-Hoc) 추정기로서의 영향 함수(IF)

재학습을 피하기 위해, 본 논문은 사전 학습된 교사로부터 직접 샘플의 중요성을 추정하기 위해 영향 함수를 활용합니다.

평탄한 검증 최소값(Flat Validation Minima, FVM): 심층 신경망에서 IF 추정을 안정화하기 위해, 저자들은 Sharpness-Aware Minimization(SAM)을 사용하여 검증 세트에서 사전 학습된 교사를 미세 조정합니다. 이는 손실 지형(loss landscape)에 국소적 평탄성을 강제함으로써, 표준 IF 방식에 비해 영향력 추정의 신뢰성을 크게 향상시킵니다.
점수 계산: 훈련 샘플 $z_{tr}$ 이 검증 세트에 미치는 영향은 미세 조정된 교사 파라미터에 대한 손실의 그래디언트와 역 헤시안(Inverse Hessian, 효율성을 위해 대각 피셔 정보 행렬으로 근사됨)을 사용하여 근사됩니다. 이를 통해 낮은 값일수록 더 어렵거나 정보량이 많은 샘플임을 나타내는 점수를 얻습니다.

2. 학습 가능한 Beta 정책

고정된 휴리스틱 선택 대신, IF-Beta는 Beta 분포로 매개변수화된 Beta 정책을 도입합니다.

적응형 샘플링: $[0, 1]$ 로 정규화된 영향 점수가 주어지면, 각 샘플의 선택 확률은 Beta 분포 $Beta(\alpha(\phi), \beta(\phi))$ 에 의해 모델링됩니다. 파라미터 $\alpha$ 와 $\beta$ 는 학습 가능한 파라미터 $\phi$ 에 의해 제어되며, 이를 통해 정책이 특정 KD 작업에 맞춰 형태(예: 어려운 샘플에 집중, 쉬운 샘플에 집중, 또는 특정 윈도우에 집중)를 적응시킬 수 있도록 합니다.
이중 레벨 최 optimization (Bilevel Optimization): 정책 파라미터 $\phi$ $ϕ$ 의 최적화는 다음과 같은 이중 레벨 문제로 공식화됩니다:
- 내부 루프(Inner Loop): 현재 정책에 의해 샘플링된 부분 집합(구체적으로는 고정된 교사 특징 위에 있는 선형 분류기)을 사용하여 KD 손실을 최소화하도록 학생 모델을 학습시킵니다.
- 외부 루프(Outer Loop): 샘플링된 부분 집합에서 학습된 학생의 검증 성능을 최대화하도록 $\phi$ 를 업데이트합니다.
효율적인 프록시(Proxy): 이중 레벨 최적화를 실행 가능하게 만들기 위해, 내부 루프는 전체 신경망을 재학습하지 않습니다. 대신, 고정된 교사 특징에 부착된 경량 선형 분류기를 최적화합니다. 이 "프록시-학생(proxy-student)" 공식은 무시할 수 있는 계산 비용으로 KD 역학을 포착합니다.

주요 기여

KD를 위한 IF: 본 논문은 평탄한 검증 최소값(FVM)을 통해 안정화된 영향 함수가 KD를 위한 효과적이고 효율적인 사후 점수 추정기 역할을 수행함을 입증하며, 궤적 기반 지표에 필요한 값비싼 교사 재학습 과정을 제거합니다.
학습 가능한 Beta 정책: 저자들은 데이터 프루닝을 위한 매개변수화된 Beta 분포를 도입하여, 이중 레벨 최적화를 통해 학습 가능한 적응형 샘플링 분포를 학습함으로써 고정된 휴리스틱 규칙(top-k 또는 고정 윈도우 등)의 경직성을 극복합니다.
효율적인 이중 레벨 최적화: 내부 루프를 고정된 특징 위의 선형 분류기로 구성함으로써, 탐색 단계 동안 과도한 전체 학생 재학습 비용 없이 효율적인 프루닝 정책 최적화를 달者합니다.

실험 결과

다양한 교사-학생 아키텍처(ResNet, WideResNet, ViT)와 함께 CIFAR-10, CIFAR-100, ImageNet에 대한 광범위한 실험 결과는 다음과 같습니다:

우수한 성능: IF-Beta는 다양한 프루닝 비율에 걸쳐 재학습이 없는 방식(Medium-Difficulty) 및 재학습 기반 방식(DUAL)을 포함한 베이스라인 모델들을 일관되게 능가합니다.
전체 데이터 증류 성능 초과: 놀랍게도, IF-Beta는 훨씬 적은 데이터와 계산량으로도 전체 데이터셋에서 증류된 학생의 성능을 뛰어넘을 수 있게 합니다. 예를 들어, CIFAR-10에서 데이터의 70%만을 유지했을 때 95.69%의 정확도를 달성하여 전체 데이터 증류의 95.50%를 넘어섰습니다.
계산 효율성: 프루닝 과정의 비용을 고려하더라도, IF-Beta는 전체 데이터 증류와 동일한 성능에 도달하는 데 훨씬 적은 시간이 걸립니다 (예: CIFAR-10에서 0.806배의 비용).
일반화: 이 방법은 비-KD 표준 지도 학습 설정과 다양한 KD 목적 함수(Relational, FitNets, Attention Transfer)에도 일반화되며, 이는 데이터 부분 집합의 품질이 특정 증류 알고리즘과 독립적인 핵심 요소임을 보여줍니다.

의의 및 주장

본 논문은 훈련 데이터의 품질이 특정 증류 알고리즘의 세부 사항보다 종종 더 중요하다는 점을 주장합니다. 가장 정보가 많은 샘플을 효율적으로 식별하고 선택함으로써, IF-Beta는 데이터 큐레이션을 현대적 훈련 파이프라인을 개선하기 위한 원칙적이고 실질적인 레버로 위치시킵니다. 이 연구는 "다이어트 중인 증류(distilling on a diet)"—즉, 전체 데이터셋보다는 신중하게 프루닝된 부분 집합을 사용하는 것—가 더 나은 일반화와 더 낮은 계산 비용을 가져올 수 있음을 시사하며, 더 많은 데이터가 항상 더 나은 증류 결과를 보장한다는 가설에 도전합니다.