Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration

기술 요약: 불확실성 보정을 위한 확산 모델에서 영감을 받은 트랜스포머 재구성

문제 제기
사전 훈련된 트랜스포머는 비전, 언어, 음성 도메인 전반에 걸친 대규모 기반 모델의 핵심 인프라가 되었습니다. 그러나 이러한 모델들은 일반적으로 특징 변환 스택을 통한 불확실성 전파를 위한 원칙적인 메커니즘이 결여되어 있습니다. 그 결과, 종종 잘못된 예측에 대해 높은 확신을 부여하여 고위험 응용 분야에서 안전상의 위험을 초래합니다. 불확실성 보정 (UC) 기법이 더 작은 네트워크에 존재하지만, 이러한 기법들은 종종 변분 추론이나 사후 확률 샘플링을 통해 모델 사후 확률을 근사하는 데 의존하며, 이는 대규모 모델의 경우 계산적으로 불가능하고 부정확해집니다.

어텐션 블록을 가우시안 프로세스 (GP) 로 재매개화하려는 기존 접근법들은 블록을 독립적으로 취급한다는 치명적인 한계에 직면해 있습니다. 이는 사전 훈련 중에 확립된 특징 변환 간의 상관관계를 고려하지 못합니다. 그 결과, 이러한 방법들은 종종 개선된 불확실성 보정이 크게 저하된 예측 성능을 대가로 치르는 트레이드오프를 강요합니다. 더 나아가, 확률적 어텐션 블록과 점 추정 구성 요소 (MLP 및 레이어 정규화와 같은) 의 교차 배치는 불확실성 전파의 흐름을 방해하여 잘 정의된 확률적 경로의 수립을 막습니다.

방법론
저자들은 원칙적인 불확실성 전파를 가능하게 하기 위해 사전 훈련된 트랜스포머를 통합된 확산 과정으로 재구성하는 DIRECTOR(Diffusion-Inspired REConfiguration of TransfORmers) 라는 프레임워크를 제안합니다.

1. 확률 경로로서의 구조적 재구성:
   이 방법은 트랜스포머 아키텍처를 재분할하여 각 변환 블록이 멀티헤드 셀프 어텐션 (MHSA) 모듈로 종료되도록 시작합니다. MHSA 와 가우시안 프로세스 간의 확립된 연결에 기반하여, 각 MHSA 블록의 출력은 중간 특징에 대한 가우시안 분포로 해석됩니다.
   결정론적 세그먼트 (MLP 및 레이어 정규화) 는 일반적으로 불확실성 흐름을 방해하지만, 이 방법에서는 이러한 세그먼트를 알 수 없는 사전 분포에서 관찰된 함수가 아니라, GP-재매개화 MHSA 에 의해 유도된 신경화된 가우시안 전이의 추가 매개변수화로서 재해석합니다. 이러한 재구성은 특징 변환의 전체 시퀀스가 가우시안 전이를 갖는 일관된 확률적 과정 $\{X_t\}_{t=0}^T$를 형성하도록 보장하여, 트랜스포머 아키텍처를 역시간 확산 과정과 효과적으로 정렬시킵니다.

2. 통합 확산 모델로의 증류:
   각 블록의 전이를 독립적으로 취급하는 대신, DIRECTOR 는 관찰된 특징 변환 시퀀스를 통합된 시공간 전이 커널로 증류합니다. 이는 역시간 과정을 학습하는 확산 기반 모델 (단일 블록 DiT 로 구현됨) 을 훈련하여 최종 특징 분포를 입력 데이터 분포로 매핑함으로써 달성됩니다.
   훈련 목적 함수는 두 가지 구성 요소를 결합합니다:

   • 불확실성 인식 손실 ($L_1$): 학습된 전이 커널을 고정된 사전 훈련된 백본에서 관찰된 전이와 정렬시키는 Kullback-Leibler (KL) 발산 항입니다. 이는 모델이 어텐션 블록 간의 내재적 상관관계를 포착하도록 보장합니다.
   • 성능 인식 손실 ($L_2$): 학습된 확산 과정에서 추출된 샘플에 적용되는 표준 예측 손실 (예: 교차 엔트로피) 입니다. 이는 재구성된 모델이 예측 정확도를 유지하거나 향상시키도록 보장합니다.

주요 기여

• 확률적 재구성: 이 논문은 사전 훈련된 트랜스포머를 신경화된 가우시안 전이의 시퀀스로 변환하는 새로운 구조적 재구성을 도입하여, 통합된 전이 커널로 매개변수화된 확산 과정으로 모델을 일반화합니다.
• 통합 불확실성 전파: 트랜스포머의 특징 경로를 단일 확산 모델로 증류함으로써, 이 방법은 사전 훈련 중에 확립된 블록 간 상관관계를 포착하여 전체 모델 사후 확률을 계산하는 비실현 가능성 없이 원칙적인 불확실성 전파를 가능하게 합니다.
• 효율적인 사후 보정: 이 접근법은 확률적 추론의 사후 임베딩으로 기능합니다. 통합 전이 모델은 원래 백본보다 훨씬 작습니다 (예: ViT 의 경우 6.24M 파라미터 대비 2.7M 파라미터), 이는 원래 모델의 예측 능력을 유지하면서 향상된 메모리 효율성을 제공합니다.

실험 결과
저자들은 비전 벤치마크 (CIFAR-10, CIFAR-100, CIFAR-10-C) 와 언어 벤치마크 (IMDB, CoLA, CoLA-OOD) 에서 DIRECTOR 를 평가했습니다.

• 보정과 정확도: DIRECTOR 는 기존 불확실성 인식 트랜스포머 (KEP-SVGP 변형을 포함) 와 표준 베이스라인 (Temperature Scaling, Monte Carlo Dropout, Deep Ensembles) 보다 일관되게 우수한 성과를 거두었습니다. CIFAR-10 에서 DIRECTOR 는 더 높은 정확도 (KEP-7/7 의 82.68% 대비 84.74%) 와 더 낮은 기대 보정 오차 (6.95% 대비 5.18%) 를 달성했습니다. 모든 데이터셋에서 유사한 개선이 관찰되어, 이 방법이 보정과 성능 간의 전형적인 트레이드오프를 해결함을 보여주었습니다.
• 강건성: 분포 변화 (CIFAR-10-C 및 CoLA-OOD) 하에서 DIRECTOR 는 사전 훈련된 모델에 비해 보정 지표를 크게 개선하면서 경쟁력 있는 예측 정확도를 유지했습니다.
• OOD 감지: 이 방법은 Out-of-Distribution 감지 작업에서 더 우수한 성능을 보여주었으며, 더 높은 AUROC 및 AUPR 점수를 달성했습니다.
• 확장성: 더 큰 모델 (ViT-B-16 및 Qwen-2.5) 에 대한 실험은 DIRECTOR 가 효과적으로 확장되어 원래 모델보다 적은 파라미터로 신뢰할 수 있는 보정을 제공함을 확인했습니다.

의의 및 주장
이 논문은 DIRECTOR 가 대규모 사전 훈련된 모델의 불확실성 보정을 위한 확장 가능하고 원칙적인 패러다임을 제공한다고 주장합니다. 트랜스포머의 특징 변환 스택을 확산 과정으로 재해석함으로써, 이 방법은 표현력 있는 불확실성 인식 구조를 모델의 내부 아키텍처에 직접 임베딩합니다. 저자들은 이 접근법이 신뢰성 있는 확신 점수가 경험적 정확도를 충실히 반영하도록 보장함으로써 위험 민감 응용 분야의 신뢰할 수 있는 배포를 가능하게 한다고 주장하며, 이는 예측 성능을 유지하거나 향상시키고 메모리 효율성을 개선하는 동시에 이루어집니다. 이 연구는 핵심 능력을 희생하지 않고 기반 모델의 신뢰성을 향상시키기 위한 새로운 방향을 열어줍니다.