기술 요약: SAFE-KD

문제 정의

현대의 비전 백본(vision backbones)은 높은 정확도를 달 achieve하지만 상당한 추론 비용을 발생시켜, 엣지 디바이스나 지연 시간에 민선한 시스템에 배포하는 데 제약을 줍니다. 조기 종료 네트워크(Early-exit networks)는 중간 분류기(intermediate classifiers)를 부착하여 "쉬운" 입력에 대해 계산을 조기에 중단할 수 있게 함으로써 이 문제를 해결합니다. 그러나 현재의 배포 방식은 신뢰도(confidence)나 엔트로피(entropy) 임계값과 같은 휴리스틱한 중단 규칙에 의존하며, 이는 정식적인 보증(formal guarantees)이 결여되어 있습니다. 이러한 휴리스틱은 모델의 미교정(miscalibration)으로 인해 신뢰할 수 없으며 분포 변화(distribution shifts) 하에서 급격히 성능이 저하되므로, 정확도와 지연 시간 사이의 트레이드오프를 예측 가능하고 안전하게 수행하기 어렵습니다. 따라서 조기 종료의 리스크에 대해 유한 샘플 보증(finite-sample guarantees)을 제공하는 적응형 추론 정책이 절실히 필요합니다.

방법론

SAFE-KD(Risk-Controlled Early-Exit Distillation)는 보정된 종료(calibrated exits)와 통계적 보증을 가능하게 하기 위해 설계된 범용 멀티 엑싯(multi-exit) 래퍼입니다. 이 방법은 세 가지 핵심 구성 요소를 통합합니다:

1. 범용 멀티 엑싯 아키텍처 (Universal Multi-Exit Architecture)

SAFE-KD는 최소한의 아키텍처 가정을 바탕으로 현대적인 백본(CNN 및 Vision Transformer)에 경량화된 엑싯 헤드를 중간 깊이에 부착합니다:

• 배치(Placement): CNN(예: ResNet, ConvNeXt)의 경우, 선택된 잔차 블록 그룹(residual block groups) 뒤에 엑싯을 배치합니다. ViT(예: ViT-S, Swin-T)의 경우, 선택된 트랜스포머 블록 뒤에 엑싯을 부착하고 CLS 또는 평균 토큰 임베딩을 풀링합니다.
• 헤드 설계(Head Design): 작은 예측 헤드(Global Average Pooling 또는 Token Pooling 후 선택적 MLP/LayerNorm 및 선형 분류기)는 오버헤드를 최소화하면서도 잘 형성된 로짓(logits)을 생성하도록 보장합니다.
• 신호(Signal): 주요 신뢰도 신호는 최대 소프트맥스 확률(Maximum Softmax Probability, MSP)입니다.

2. 계층적 증류 (Hierarchical Distillation)

중간 엑싯의 취약성과 노이즈를 해결하기 위해, SAFE-KD는 지수 이동 평균(Exponential Moving Average, EMA) 교사(teacher)를 사용하는 계층적 증류 전략을 채택합니다:

• 분리된 지식 증류(Decoupled Knowledge Distillation, DKD): 모든 엑싯은 DKD를 사용하여 교사에 의해 감독됩니다. DKD는 타겟 클래스와 비타겟 클래스 성분을 증류 손실에서 분리하여 "다크 지식(dark knowledge)"을 더 잘 활용합니다.
• 심층-천층 일관성(Deep-to-Shallow Consistency): 일관성 손실은 초기 엑싯을 최종 엑싯과 정렬하여 네트워크 깊이에 따른 예측 안정성을 보장합니다.
• 목적 함수(Objective): 전체 손실은 표준 크로스 엔트로피, 교사에 대한 DKK 손실, 그리고 최종 엑싯에 대한 일관성 손실을 결합하며, 하이퍼파라미터 $\alpha$와 $\beta$에 의해 가중치가 부여됩니다.

3. 공형 리스크 제어를 통한 리스크 제어형 종료 (Risk-Controlled Exiting via Conformal Risk Control, CRC)

휴리스틱 임계값 대신, SAFE-KD는 홀드아웃 교정 데이터셋($D_{cal}$)을 사용하여 각 엑싯별로 보정된 중단 임계값을 계산하는 공형 리스크 제어(CRC)를 사용합니다:

• 선택적 리스크 정의: 조기 종료는 선택적 분류(selective classification)로 프레임화됩니다. 불확실성 점수는 $r_j(x) = 1 - \max_c p_j(c|x)$로 정의됩니다.
• 교정(Calibration): 사용자 지정 리스크 수준 $\delta$에 대해, CRC는 각 엑싯 $j$에 대한 임계값 $\hat{\tau}_j(\delta)$를 계산합니다. 이 임계값은 교환 가능성(exchangeability) 가정 하에, 선택적 오분류 리스크(엑싯 $j$에서 수용된 샘플들의 에러 확률)가 $\delta + O(1/|D_{cal}|)$ 내로 제한됨을 보장합니다.
• 추론 정책(Inference Policy): 테스트 시, 모델은 불확실성 점수 $r_j(x) \leq \hat{\tau}_j(\delta)$를 만족하는 가장 빠른 헤드 $j$에서 종료합니다. 만약 어떤 조기 엑싯도 조건을 만족하지 못하면, 모델은 최종 엑싯까지 진행합니다.

주요 기여

1. 범용 멀티 엑싯 래퍼: 백본의 특징을 보존하면서 최소한의 오버헤드로 CNN 및 ViT에 중간 엑싯을 추가하는 모듈형 설계.
2. 계층적 증류: 중간 예측의 정확도와 안정성을 향상시키기 위해 DKD와 심층-천층 일관성을 사용하는 학습 접근법.
3. 리스크 제어형 종료: 신뢰도 점수를 사용자 조절이 가능한 리스크 제한 임계값으로 변환하여, 취약한 휴리스틱을 통계적으로 유효한 테스트로 대체하는 교정 절차.
4. 실증적 검증: 다양한 데이터셋과 아키텍처에 걸쳐 평가되어, 개선된 정확도-연산 트레이드오프, 강력한 교정 능력, 그리고 손상(corruption)에 대한 강건성을 입증함.

실험 결과

저자들은 ResNet-50, MobileNetV3-S, EfficientNet-B0, ConvNeXt-T, ViT-S, Swin-T와 같은 백본과 CIFAR-10/100, STL-10, Oxford-IIIT Pets, Flowers102, FGVC Aircraft 데이터셋을 사용하여 SAFE-KD를 평가했습니다.

• 정확도 및 효율성: SAFE-KD는 표준 DKD의 최종 엑싯 정확도와 일치하거나 이를 능가하며, MultiExit 베이스라인들을 크게 앞섭니다. CIFAR-100에 대한 ViT-S의 경우, ERM 대비 +0.80%의 정확도 향상을 달성했습니다.
• 파레토 우위(Pareto Dominance): SAFE-KD는 고정된 신뢰도 및 엔트로피 기반 게이팅에 비해 엄격하게 우월한 파레토 프런티어를 달로합니다. 예를 들어, 타겟 리스크 $\delta=0.05$인 ResNet-50에 대해, SAFE-KD는 94.1%의 상대적 정확도를 유지하면서 계산 깊이를 41% 감소시켰습니다.
• 리스크 유효성: 실증적 선택적 리스크 곡선은 지속적으로 $y=x$ 대각선을 추적하거나 그보다 약간 아래에 머물러, 유한 샘플 유효성을 확인해 줍니다. 반면, 고정 신뢰도 0.95와 같은 휴리스틱 베이스라인은 종종 리스크 제약 조건을 위반하여, 관측된 리스크(예: Flowers102에서 12%)가 타겟을 훨씬 초과하는 결과를 초래했습니다.
• 강건성: CIFAR-10-C 손상 환경에서 SAFE-KD는 표준 모델보다 더 완만하게 성능이 저하되었습니다. 교정된 임계값은 손상된 샘플을 조기 엑싯에서 자연스럽게 거부하여, 모델이 더 깊고 강건한 레이어를 활용하도록 강제했습니다.
• 절제 연구(Ablation): DKD를 제거하면 상당한 정확도 저하(-2.1%)가 발생하고 보수적인 임계값(깊이 증가 0.85)을 강제하게 됩니다. 일관성을 제거하면 리스크 변동성이 증가합니다. 전체 방법론이 가장 낮은 깊이와 높은 정확도의 최적의 균형을 제공했습니다.

의의 및 주장

본 논문은 SAFE-KD가 적응형 추론을 휴리스틱 최적화 문제에서 신뢰할 수 있는 리스크 제어 시스템으로 변모시킨다고 주장합니다. 조기 종료의 에러율을 수학적으로 제한함으로써, 모델 출력과 시스템 레벨 요구사항 사이의 신뢰할 수 있는 인터페이스를 제공합니다. 이를 통해 시스템 설계자는 리스크 파라미터 $\delta$를 단순한 튜닝 가능한 하이퍼파라미터가 아닌 하드 제약 조건으로 취급할 수 있으며, 이는 예측 가능한 지연 시간과 에러 바운드가 필수적인 안전 필수(safety-critical) 배포 환경에서 조기 종료 네트워크를 실용적으로 만듭니다.

저자들은 교정 데이터셋에 대한 의존성(작은 데이터셋은 보수적인 임계값을 유발함)과 엄격한 분포 변화 시 깨질 수 있는 교환 가능성 가정에 대한 한계를 언급했습니다. 그러나 이러한 범위 내에서 SAFE-KD는 효율적이고 안전한 추론을 위한 원칙적인 프레임워크를 제공합니다.