Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers
이 논문은 Deep Ensembles 의 높은 계산 비용 문제를 해결하면서도 불확실성 정량화 성능을 유지하거나 향상시키기 위해, 어텐션 헤드를 가지치기하고 새로운 멀티헤드 어텐션 구조로 병합하는 'Hydra Ensembles'라는 효율적인 트랜스포머 앙상블 방법을 제안합니다.
원저자:Firas Gabetni, Giuseppe Curci, Andrea Pilzer, Subhankar Roy, Elisa Ricci, Gianni Franchi
1. 문제: "모든 걸 다 아는 천재"는 너무 비싸고, "혼자 일하는 사람"은 실수를 잘 합니다.
AI 모델 (딥러닝) 은 의료나 자율주행 같은 중요한 분야에서 실수하면 큰일 납니다. 예를 들어, AI 가 암 진단을 내릴 때 "99% 확신"이라고 말했지만 실제로는 틀렸다면 치명적이죠.
기존의 해결책 (Deep Ensembles): 가장 확실한 방법은 동일한 업무를 하는 전문가 3~5 명을 고용해서 각자 독립적으로 진단을 내리게 한 뒤, 그 결과를 합치는 것입니다. 이렇게 하면 실수를 줄이고 "내가 얼마나 확신하는지"도 정확히 알 수 있습니다.
단점: 전문가 5 명을 고용하려면 인건비 (계산 비용) 가 5 배나 들고, 사무실 공간 (메모리) 도 5 배가 필요합니다. AI 모델이 거대해질수록 이 비용은 감당하기 어려워집니다.
기존의 다른 방법 (가지치기/Pruning): 비용을 아끼려고 전문가 5 명을 고용하는 대신, 한 명의 전문가에게서 불필요한 지식 (머리) 을 잘라내서 5 개의 다른 전문가를 만들어내는 시도가 있었습니다.
문제: 하지만 단순히 지식을 잘라내면, 그 전문가들이 실제 상황 (예상치 못한 환자) 에 대해 "내가 모른다"라고 제대로 말하지 못하게 됩니다. 즉, 확신은 있는데 틀리는 '과신' 상태가 되어버립니다.
2. 해결책: "히드라"처럼 머리를 잘라내고 다시 합치다
이 논문은 **히드라 (Hydra)**라는 신화 속 괴물을 비유로 사용합니다. 히드라는 머리를 하나 자르면 그 자리에서 두 개의 머리가 다시 자라나는 괴물입니다.
연구자들은 이 아이디어를 AI 에 적용했습니다:
**하나의 거대한 AI(Transformer)**를 준비합니다.
이 AI 는 여러 개의 "주의 집중 헤드 (Attention Heads)"라는 작은 뇌 영역으로 구성되어 있습니다.
이 헤드들을 전략적으로 잘라내어 (Pruning) 서로 다른 3 개의 "작은 AI"를 만듭니다.
중요한 건, 단순히 잘라내는 게 아니라 어떤 헤드를 잘라낼지 신중하게 선택한다는 점입니다. (예: "이 헤드는 이상한 그림을 감지하는 데 특화되어 있으니 남겨두자" 등)
이렇게 만들어진 3 개의 작은 AI 를 하나의 거대한 AI 안에 다시 합칩니다.
3. 핵심 기술: "공유된 사무실"에서 일하는 3 명의 전문가
여기서 가장 놀라운 부분은 합치는 방식입니다.
기존 방식: 전문가 3 명이 각자 따로 사무실을 쓰고, 따로 일하고, 따로 결과를 내면 비용이 3 배입니다.
히드라 방식: 3 명의 전문가가 **하나의 큰 사무실 (단일 모델)**에 모여서 일합니다.
MLP(신경망의 기본 처리부): 3 명이 같은 책상과 도구를 공유합니다. (비용 절감)
Attention Heads(주의 집중부): 각자가 가진 고유한 "시각 (Head)"만 유지합니다. 한 사람은 "색깔"에 집중하고, 다른 사람은 "형태"에 집중하는 식으로 서로 다른 관점을 유지합니다.
이렇게 하면 3 명의 전문가가 합쳐진 지능을 가지면서도, 1 명의 전문가가 일하는 속도로 결과를 낼 수 있습니다.
4. 왜 이 방법이 좋은가요? (실생활 예시)
속도: 기존에 3 대의 컴퓨터로 3 번 계산해야 했던 일을, 1 대의 컴퓨터로 한 번에 처리합니다. 속도가 거의 3 배 빨라집니다.
정확도: 단순히 지식을 잘라낸 것보다, **특화된 지식 (회로)**을 잘라내어 합쳤기 때문에, 예상치 못한 상황 (예: 훈련 데이터에 없던 새로운 질병) 에 대해 "이건 잘 모르겠어요"라고 정확하게 경고할 수 있습니다.
재교육 불필요: 처음부터 3 대의 컴퓨터를 새로 조립할 필요 없이, 이미 훈련된 거대한 AI 하나만 가지고도 이 기술을 적용할 수 있습니다.
5. 결론: "가볍지만 똑똑한" AI 의 탄생
이 논문은 "불확실성을 측정하는 AI 를 만들려면 무조건 무겁고 비싸야 한다"는 고정관념을 깨뜨렸습니다.
비유하자면:
예전에는 "3 명의 명수를 고용해서 3 번 진단받아야 안전하다"고 생각했습니다.
이제는 **"한 명의 명수에게서 3 가지의 다른 관점 (시각) 을 추출해서, 한 번에 3 번의 진단을 동시에 내리게 하는 기술"**을 개발했습니다.
결과는 3 명을 고용한 것과 비슷하거나 더 좋으면서, 비용은 1 명을 고용한 것과 같습니다.
이 기술은 자율주행차가 길을 잃었을 때, 혹은 AI 가 의료 진단을 내릴 때 "내가 이 답을 확신할 수 없다"라고 정직하게 말하게 만들어, 더 안전하고 신뢰할 수 있는 AI 시대를 여는 중요한 발걸음이 될 것입니다.
1. 문제 정의 (Problem)
딥러닝 모델, 특히 트랜스포머 기반의 대규모 모델 (Foundation Models) 은 안전이 중요한 분야 (의료, 자율주행 등) 에 배포될 때 **불확실성 정량화 (Uncertainty Quantification, UQ)**가 필수적입니다.
현재의 한계: 가장 신뢰할 수 있는 UQ 방법인 Deep Ensembles는 여러 개의 독립적으로 학습된 모델을 집계하여 높은 정확도와 보정 (Calibration) 성능을 제공하지만, 학습 비용, 메모리 사용량, 추론 시간이 모델 수에 비례하여 급증하여 확장성이 매우 낮습니다.
기존 대안의 문제: 기존 효율적인 앙상블 방법 (MC Dropout, MIMO, BatchEnsemble 등) 은 재학습 비용은 줄이지만 여전히 전체 모델을 다시 학습해야 하거나, 가중치 공유로 인해 앙상블 멤버 간의 다양성이 부족합니다.
가지치기 (Pruning) 의 함정: 단순히 불필요한 가중치를 제거하는 기존 가지치기 방법은 정확도는 유지할 수 있으나, 노이즈가 있는 데이터나 분포 외 (OOD) 데이터에서 모델의 보정 능력을 해치고 불확실성 추정을 왜곡시킬 수 있음이 본 논문에서 이론적 및 실험적으로 증명되었습니다.
2. 제안 방법: Hydra Ensembles (Methodology)
저자들은 Hydra Ensembles라는 새로운 프레임워크를 제안합니다. 이는 단일 사전 학습된 트랜스포머 모델에서 어텐션 헤드 (Attention Heads) 를 구조적으로 가지치기하여 다양한 하위 네트워크를 생성하고, 이를 하나의 효율적인 모델로 병합하는 방식입니다.
핵심 구성 요소 및 과정:
구조적 가지치기 (Structured Head Pruning):
단일 트랜스포머 백본에서 여러 개의 사본을 생성합니다.
각 사본마다 어텐션 헤드를 제거하여 서로 다른 하위 네트워크 (서브네트워크) 를 만듭니다.
가지치기 전략:
Taylor Pruning: 손실 함수의 1 차 테일러 급수를 기반으로 중요도가 낮은 헤드를 제거 (데이터셋이 없는 경우).
Circuit Extraction (Headmap): 불확실성 (OOD 감지) 에 중요한 회로 (Circuit) 를 식별하여 헤드를 선택적으로 제거 (불확실성 검증 데이터가 있는 경우).
병합 (Merging) 메커니즘:
Fused Multi-Head Attention (MHA): 가지치기된 각 모델의 남은 어텐션 헤드를 Grouped Fully-Connected (GFC) 레이어를 사용하여 하나의 MHA 블록으로 통합합니다. 이는 여러 모델의 어텐션 연산을 병렬로 처리하여 단일 모델과 유사한 추론 속도를 보장합니다.
Merged MLP: 각 모델의 MLP 가중치를 평균화하여 단일 MLP 로 병합합니다.
효율성:
이 방식은 여러 모델을 순차적으로 실행하는 기존 앙상블과 달리, **단일 순전파 (Single Forward Pass)**로 앙상블 예측을 수행합니다.
추가적인 학습 (Fine-tuning) 이 없거나 최소한으로만 수행 가능하며, 사전 학습된 모델을 그대로 활용할 수 있습니다.
3. 주요 기여 (Key Contributions)
가지치기와 불확실성의 관계 규명: 이론적 증명 (Proposition 1) 과 실험을 통해, 단순한 가지치기는 노이즈가 있는 데이터에서 성능 저하를 유발하고 보정을 해칠 수 있음을 보였습니다. 반면, 구조화된 어텐션 헤드 가지치기는 모델의 내부 표현을 유지하면서 다양성을 확보할 수 있음을 입증했습니다.
Hydra Ensembles 프레임워크 개발: 대규모 트랜스포머 모델 (ViT, BERT, CLIP) 에 특화된 최초의 가지치기 기반 UQ 프레임워크를 제안했습니다. 이는 재학습 없이도 Deep Ensembles 수준의 다양성을 확보합니다.
성능과 효율성의 동시 달성:
Deep Ensembles 와 유사하거나 더 나은 UQ 성능 (정확도, 보정, OOD 감지) 을 달성합니다.
추론 비용은 단일 모델과 거의 동일하게 유지합니다 (Deep Ensembles 대비 약 3 배 빠른 추론).
4. 실험 결과 (Results)
논문은 이미지 분류, 텍스트 분류, 제로샷 (Zero-shot) 분류 등 다양한 태스크에서 Hydra Ensembles 를 평가했습니다.
지도 학습 이미지 분류 (ImageNet-1K, CIFAR-100):
ViT-B/16 기반 실험에서 Hydra Ensembles 는 Deep Ensembles 와 유사한 정확도를 유지하면서, **OOD 감지 성능 (AUROC, FPR95)**에서 기존 효율적 앙상블 방법들보다 우월한 성능을 보였습니다.
추론 비용: BF16 정밀도에서 단일 모델 대비 1.07 배의 추론 시간만 소요되었으며, Deep Ensembles 는 약 3 배의 시간이 걸렸습니다.
지도 학습 텍스트 분류 (SST-2, BERT):
BERT 기반 실험에서 Hydra Ensembles 는 Deep Ensembles 와 유사한 정확도와 보정 성능을 보였으며, OOD 감지 성능은 Deep Ensembles 보다 AUROC +2.8%, FPR95 -7.6% 개선되었습니다.
제로샷 이미지 분류 (OpenCLIP-ViT):
추가 학습 없이 (Zero-shot) 수행된 실험에서 Hydra Ensembles 는 ViLU (현재 SOTA 방법) 를 능가하는 성능을 기록했습니다.
ImageNet-1K OOD 벤치마크에서 AUROC +1.3%, FPR95 -3.5%, AUPR +4.0% 개선 효과를 보였습니다.
5. 의의 및 결론 (Significance)
확장 가능한 UQ: 대규모 기초 모델 (Foundation Models) 에 적용하기 어려운 Deep Ensembles 의 계산 비용을 획기적으로 줄이면서도, 높은 수준의 불확실성 추정을 가능하게 합니다.
구조적 다양성의 중요성: 단순한 가중치 제거가 아닌, 어텐션 헤드의 구조적 다양성을 유지하는 것이 UQ 성능에 결정적임을 보여주었습니다. 특히 'Circuit' 기반 가지치기는 OOD 감지에 매우 효과적입니다.
실용성: 재학습 비용 없이 기존 모델을 효율적으로 변환하여 안전이 중요한 애플리케이션에 배포할 수 있는 실용적인 솔루션을 제공합니다.
이 논문은 **효율성 (Efficiency)**과 **신뢰성 (Reliability/UQ)**이라는 상충되는 목표를 동시에 달성하기 위한 새로운 패러다임을 제시하며, 현대 딥러닝 시스템의 안전성 강화에 중요한 기여를 하고 있습니다.