QG-MIL: 의료 영상에서 도메인 비종속 Multiple Instance Learning을 위한 게이트드 Transformer 집계기

Attention 기반 MIL은 어텐션 가중치가 소수 인스턴스에 편중되어 예측이 과도하게 자신감 및 불안정해지는 문제가 존재한다. QG-MIL은 아키텍처 수준의 정규화와 게이팅을 통해 이러한 attention concentration을 구조적으로 완화하여 도메인 간(whole-slide pathology ↔ hematology) 일관된 성능 개선을 달성한다.

Multiple Instance Learning(MIL)은 각 샘플(슬라이드 또는 환자)이 다수의 인스턴스(패치 또는 세포)를 포함하고, 라벨은 bag 단위로만 주어지는 약지도 학습 설정이다. Attention 기반 MIL은 각 인스턴스에 가중치(어텐션)를 부여해 bag 표현을 가중합으로 얻고 분류를 수행한다. 이때 학습 과정에서 일부 인스턴스에 어텐션이 과도하게 집중(attention concentration)되면 다른 유의미한 인스턴스 정보가 무시되어 예측이 과도하게 자신감 있게 되고, 도메인 편차에 취약해진다. 기존 접근은 보조 손실이나 마스킹으로 이를 완화했지만 학습 파이프라인이 복잡해진다. QG-MIL은 어텐션 집중을 학습 규제 없이 아키텍처 설계로 억제한다. 구체적으로 입력을 RMSNorm으로 사전 정규화하고, 멀티헤드 어텐션에서 각 헤드의 Query와 Key를 별도로 정규화한 뒤 scaled dot-product를 계산한다. 그 결과(attention output)에는 추가 학습 가능한 게이트(W_g)를 적용해 각 헤드 또는 각 피처별로 sigmoid 게이트를 곱해 출력 값을 조절한다. 이 흐름은 수치적 안정성과 세밀한 출력 조절을 동시에 제공해 특정 인스턴스가 어텐션을 독점하는 현상을 완화한다. 아키텍처적 차이로 인해 얻어지는 변화는 두 측면에서 나타난다. 첫째, 학습 안정성: per-head Q/K 정규화와 RMSNorm pre-norm이 그래디언트 흐름을 개선해 학습 중 발산·불안정성을 줄인다. 둘째, 어텐션 분포: 게이팅이 Z(어텐션 출력)를 스케일링(헤드 단위 또는 요소 단위)하므로 소수 인스턴스의 극단적 점유를 억제해 attention mass가 더 고른 분포로 확산된다. 실제로 Breast 테스트셋에서 QG-MIL의 top-10% attention mass는 약 15%인 반면 ABMIL은 최대 54%에 달해 분포 차이가 수치로 확인된다.

전체 접근: QG-MIL은 인스턴스별 임베딩을 동일 차원 D로 투영한 뒤 L개의 gated transformer 블록을 연속으로 적용해 인스턴스 표현을 정제하고, 마지막에 gated attention pooling으로 bag 표현 z를 계산해 분류한다. 표준 구성은 D=512, dropout=0.25, 기본 깊이 L=2이며 전체 파라미터 수는 약 9M( Light variant는 2.4M)이다. 어텐션 및 정규화(수식 흐름): 인스턴스 특징 x_i가 주어지면 W_p를 통해 h_i = Dropout(RMSNorm(W_p x_i))를 얻는다 → multi-head로 Q,K,V를 선형투영하고 헤드별로 reshape한다 → 각 헤드에서 Query와 Key를 헤드 단위로 정규화(입력 값: Q,K) → 정규화된 Q와 K로 S = (Q^ K^{^T})/sqrt(d_h)를 계산한다(이 연산은 두 행렬의 내적을 스케일링해 유사도 점수 S를 산출한다) → softmax(S)를 적용해 각 헤드의 attention weight A를 얻는다 → A와 V를 곱해 Z(어텐션 출력)를 얻는다(결과의 의미: 각 토큰에 대한 헤드별 집계 표현). gating 메커니즘(수식 흐름): attention output Z_i,h,:가 입력으로 주어지면 헤드 단위 게이트는 z_{i,h} = W_{g,h} Z_{i,h,:}를 계산한다 → g_{i,h} = σ(z_{i,h})를 통해 [0,1] 스케일의 값 획득(결과의 의미: 각 헤드의 출력 세기 조절 인자) → O_{i,h,:} = g_{i,h} Z_{i,h,:}로 모듈레이션한다. 요소별(elementwise) 게이트 변형은 z_{i,h,:}=W_{g,h}Z_{i,h,:} 형태로 d_h 차원을 유지하고 g_{i,h,:}⊙Z_{i,h,:}로 각 피처별 스케일을 적용한다. 블록 구성 및 FFN: 각 블록은 pre-norm residual 레이아웃을 따르며 FFN은 SwiGLU 스타일이다. FFN 흐름: 입력 x → W_g x 와 W_u x를 계산 → φ(W_g x) ⊙ W_u x를 통해 활성화가 피처별 곱으로 적용된 출력을 생성 → W_d로 투영해 블록 출력을 얻는다(의미: 더 표현력이 큰 은닉 변환과 안정적 비선형성을 제공). 학습·평가 설정: 모든 실험은 환자 단위로 20% hold-out 테스트셋을 확보하고, 나머지 80%에서 5-fold cross-validation으로 모델 5개를 학습해 테스트 시 평균 확률 앙상블을 사용했다. 학습은 CrossEntropy, 최대 150 epochs, early stopping(검증손실 10 epoch 개선 없을 때), 모델 선택은 검증 손실 기준이며 임베딩 크기 D=512, dropout=0.25를 고정했다. 아키텍처 변형으로 Elementwise/Headwise gating, noGate, noQKnorm, LayerNorm 대체, Light/Deep 변형 등을 평가해 구성 요소별 기여를 분석했다.

메인 벤치마크 성능: QG-MIL 계열은 6개 병리·혈액학 벤치마크에서 평균 +6.1 mean macro F1 포인트 향상을 기록했다(논문 전체 평균값). 병리 진단에서는 Breast에서 QG-MIL Deep이 macro F1 98.97%를 달성했고, Lung에서는 peak macro F1 93.3%가 보고됐다. 전립선(prognostic)에서는 QG-MIL LayerNorm 변형이 평균 57.55%로 가장 높은 성능을 보였다. 혈액세포(hematology) 벤치마크에서도 AML subtyping 등에서 평균 5% 이상의 개선이 보고됐다. 어블레이션: 게이팅, per-head QK normalization, RMSNorm(pre-norm) 및 SwiGLU 각 요소가 데이터셋별로 다른 영향력을 보였다. 작은 코호트에서는 게이팅과 QK 정규화가 분산을 키워 성능 저하를 초래하는 경우가 있었고, 중대형 코호트에서는 전체 게이트 구조가 안정성과 성능을 향상시켰다. 또한 깊이(Deep variant)는 충분한 데이터가 있을 때 성능 향상에 기여했다. 어텐션 분포 및 해석성: 정량 지표로 QG-MIL은 attention distribution을 더 균일하게 유지했다. Breast 테스트에서 평균 top-10% attention mass는 QG-MIL 약 15% 대, ABMIL은 최대 54%로 보고됐다. Gini와 entropy 수치로도 ABMIL(Gini: 0.62±0.13, entropy: 0.90±0.05) 대비 QG-MIL(Gini: 0.16±0.07, entropy: 0.99±0.01)로 분포가 유의하게 고르다. 효율성 및 리소스: 논문에 따르면 realistic bag sizes 512 및 1024에 대해 QG-MIL의 연산량은 각각 약 7 GFLOPs 및 18 GFLOPs이며, 이는 RRT 및 TransMIL보다 낮은 부담을 보였다. 단, 게이팅 및 SwiGLU로 인해 학습 시 메모리와 시간이 증가한다고 보고됐다.

전체 아키텍처: 입력 인스턴스 x_i는 고정된 피처 추출기(예: UNI2-h, Prov-Gigapath, CONCHv1.5, DinoBloom 등)로부터 얻어지며 W_p로 투영되어 h_i = Dropout(RMSNorm(W_p x_i))을 생성한다. 이후 L개의 gated transformer 블록이 연속 적용되어 최종 인스턴스 표현을 얻고, gated attention pooling으로 bag 표현 z를 계산해 분류한다. 표준 설정은 D=512, L=2이며 모델 파라미터는 표준형 약 9M이다. 핵심 메커니즘의 수학적/알고리즘적 기반: 멀티헤드 어텐션에서 각 헤드의 Q,K를 헤드별로 정규화한 뒤 S = (Q^ K^{⊤})/√d_h를 계산하고 A = softmax(S)를 얻는다 → A와 V를 곱해 Z를 얻는다(의미: 헤드별 토큰 가중합). 게이팅은 Z에 대해 z=W_g Z → g=σ(z) → O=g⊙Z로 적용되며, elementwise variant는 g와 Z가 d_h 차원에서 곱해져 각 피처별 조절을 수행한다. 이 흐름은 확률적 또는 손실 기반 정규화 없이도 어텐션 분포를 완화한다. Prior work 대비 차별점: 기존 MIL 안정화는 보조손실(entropy regularization), 마스킹, self-supervised 사전학습, teacher-student distillation 등 학습 단계에서의 추가 규제를 사용했다. QG-MIL은 아키텍처 수준의 정규화와 게이팅으로 동일 목적을 달성해 파이프라인 복잡도를 늘리지 않고도 attention sink 문제를 해결한다. 구현 및 학습 세부사항: 모든 인코더는 frozen 상태로 사용했고, 학습은 CrossEntropy, 최대 150 epochs, early stopping(10 epochs patience), dropout=0.25로 수행했다. 검증 기반 모델 선택과 환자-단위 20% hold-out + 5-fold CV로 평가 안정성을 확보했고, 테스트 시 fold별 모델을 평균 확률로 앙상블했다. Ablation으로 LayerNorm 대체, noGate, noQKnorm, Elementwise 게이팅, Light(2.4M) / Deep(4-layer) 변형 등을 비교해 구성요소별 기여와 코호트 크기 의존성을 분석했다. 계산 비용: 논문은 bag size 512/1024에서 각각 약 7 GFLOPs와 18 GFLOPs가 소요된다고 보고했다. 학습 시 게이팅·SwiGLU로 인한 메모리/시간 증가는 존재하지만 추론 비용은 다른 비교 대상보다 합리적이라고 명시됐다.

논문에서 명시한 한계는 다음과 같다. (1) 작은 환자 코호트에서는 전체 게이팅 아키텍처가 분산을 키워 성능 저하를 유발해, 간단한 변형(noGate 등)이 더 효과적일 수 있다. (2) fine-grained gating과 SwiGLU 모듈은 학습 중 메모리와 시간 소모를 증가시킨다(추론 비용은 비교적 낮음). (3) 광범위한 하이퍼파라미터 최적화는 수행되지 않아 추가 튜닝으로 성능 향상이 가능할 여지가 있다.