기술 요약: LARGO (결손 모달리티 처리를 위한 저랭크 하이퍼네트워크)

문제 정의

다중 모달리티 이미지 분석, 특히 의료 영상 (예: 뇌종양 분할, 뇌졸중 병변 검출) 에서 장비 제약, 스캔 시간 제한, 또는 환자 상태 등으로 인해 사용 가능한 모든 모달리티를 획득하는 것은 종종 불가능합니다. 이로 인해 추론 시 불완전한 모달리티 세트가 발생합니다. 명시적인 아키텍처 설계 없이는 표준 딥러닝 모델이 이러한 결손 입력을 처리하는 데 어려움을 겪습니다.

기존 접근 방식은 전문성과 확장성 사이의 근본적인 트레이드오프에 직면해 있습니다:

• 단일 모델 방법 (예: 제로 임피테이션, 공유 인코더) 은 하나의 네트워크가 모든 $2^N - 1$개의 모달리티 조합을 처리하도록 강제하여, 종종 특정 구성에서의 성능을 희생합니다.
• 전용 모델 (각 조합마다 별도의 모델 학습) 은 전문성을 극대화하지만 저장 공간과 계산 비용이 기하급수적으로 증가합니다.
• 생성 및 증류 방법은 임상적으로 위험한 특징을 환각 (hallucinate) 할 위험이 있거나 복잡한 다단계 학습을 요구합니다.

대부분의 현재 솔루션은 결손 입력을 처리하기 위해 견고한 표현을 설계하는 **특징 공간 (feature space)**에서 작동합니다. 본 논문은 서로 다른 모달리티 부분집합을 위한 별도 모델을 학습하는 데 내재된 중복성을 해결하기 위해 **가중치 공간 (weight space)**에서 작동하는 것을 제안합니다.

방법론: LARGO

저자들은 $2^N - 1$개의 결손 모달리티 전용 모델 군을 단일 네트워크로 압축하는 하이퍼네트워크 프레임워크인 LARGO(Low-rAnk hypeRnetwork for missinG mOdalities) 를 제안합니다.

핵심 아키텍처

1. 하이퍼네트워크 공식화: $M = 2^N - 1$개의 별도 가중치 세트를 학습하는 대신, LARGO 는 모든 전용 모델 $\{F_{\theta(m)}\}_{m=1}^M$을 암묵적으로 포함하는 단일 하이퍼네트워크 $H_\phi$를 구성합니다.
2. 텐서 분해 (CPD): 이 방법은 동일한 모달리티에 대한 접근 권한을 공유하는 모델들이 유사한 패턴을 학습하여 매개변수 중복성이 있음을 관찰합니다. LARGO 는 모든 $M$개의 전용 모델의 합성곱 가중치를 새로운 "모델 차원"을 따라 적층하고 정준 다항식 (Canonical Polyadic, CP) 텐서 분해를 적용합니다.
   • 가중치 텐서 $W \in \mathbb{R}^{M \times C_{in} \times C_{out} \times K}$는 인자 행렬 $A, B, C, D$로 분해됩니다.
   • 인자 행렬 $A$는 모달리티 조합을 인덱싱하고, $B, C, D$는 각각 입력 채널, 출력 채널, 커널 차원을 나타냅니다.
3. 저랭크 재매개변수화:
   • 내부 레이어: 합성곱 레이어는 CP 분해를 사용하여 재매개변수화됩니다. 랭크 $R$은 단일 전용 모델과 동등한 매개변수 예산을 유지하도록 선택되어 ($1/M$ 압축 비율 달성).
   • 스템 (Stem) 및 헤드 (Head) 레이어: 이러한 레이어는 입력/출력 유연성을 유지하기 위해 압축되지 않지만, 내부 레이어에 비해 매개변수 기여도는 미미합니다.
4. 추론: 추론 중 특정 모달리티 조합 $m$에 대해, 하이퍼네트워크는 인자 행렬을 수축 (특히 인자 행렬 $A$의 $m$번째 행 선택) 하여 전용 합성곱 가중치 $W^{(m)}$을 재구성합니다. 이는 다른 모달리티 부분집합에 사용되지 않는 가중치를 계산하는 것을 방지합니다.

학습 절차

• 손실 함수: 모델은 표준 분할 손실 ($L_{Dice} + L_{CE}$) 을 사용하여 학습됩니다.
• 전략: 학습은 랜덤 모달리티 드롭아웃과 전체 모달리티 가이드를 결합합니다. 각 반복에서 특정 결손 모달리티 시나리오에 대한 손실을 최소화하기 위해 부분집합 $m$이 샘플링되는 반면, 전체 모달리티 입력을 사용하는 항은 네트워크가 학습 내내 완전한 데이터에 의해 안내되도록 보장합니다.
• 초기화: 인자 행렬 $B, C, D$는 Kaiming 정규 분포로 초기화되고, 행렬 $A$는 1 로 초기화됩니다. 가중치는 사전 학습된 가중치에서 파생되는 것이 아니라 처음부터 학습됩니다.

주요 기여

1. 새로운 하이퍼네트워크 공식화: LARGO 는 CP 분해를 사용하여 $2^N - 1$개의 결손 모달리티 전용 모델 군 전체의 커널 가중치를 단일의 매개변수 효율적인 표현으로 공동 재매개변수화하는 최초의 방법입니다.
2. 가중치 공간 작동: 특징 공간 방법과 달리 LARGO 는 직접 가중치 공간에서 작동하여, 보조 손실, 교사 네트워크, 또는 다단계 학습 없이 기하급수적인 모델 군을 단일 압축 가중치 매니폴드로 압축합니다.
3. 최첨단 성능: 이 방법은 단일 전용 모델과 비교 가능한 매개변수 수를 유지하면서 의료 영상 벤치마크에서 우수한 결과를 달성합니다.

실험 결과

데이터셋 및 시나리오

• BraTS 2018: 4 개 모달리티 (FLAIR, T1, T1c, T2), 15 개 결손 모달리티 시나리오. 작업: 뇌종양 분할 (WT, TC, ET).
• ISLES 2022: 3 개 모달리티 (DWI, ADC, FLAIR), 7 개 결손 모달리티 시나리오. 작업: 뇌졸중 병변 분할.
• avMNIST: 2 개 모달리티 (이미지, 오디오), 3 개 시나리오. 작업: 분류 (비의료 이질적 모달리티에 대한 개념 증명).

정량적 성능

• BraTS 2018: LARGO 는 전체 종양 (WT) 에서 15 개 시나리오 중 15 개, 종양 코어 (TC) 에서 11 개, 강화 종양 (ET) 에서 14 개의 시나리오에서 1 위를 기록했습니다.
  • 베이스라인 (mmFormer, M³AE, ShaSpec, SimMLM) 대비 평균 Dice 개선: WT +0.68%, TC +2.53%, ET 에서 유의미한 향상.
  • LARGO 는 평균 Dice 로 WT 87.05%, TC 79.35%, ET **62.83%**를 달성했습니다.
• ISLES 2022: LARGO 는 전체 7 개 시나리오에서 1 위를 기록했습니다.
  • 평균 Dice: 62.67% (다음 최상위인 M³AE 의 60.14% 보다 우수).
  • 어려운 단일 모달리티 사례에서 주목할 만한 향상 (예: ADC 만: 42.49% 대 40.03%).
• avMNIST: LARGO 는 가장 높은 평균 정확도 (97.00%) 를 달성하여 가장 가까운 경쟁자보다 +0.44% 향상되어 이질적인 비의료 모달리티에 대한 적용 가능성을 입증했습니다.

효율성 및 복잡성

• 매개변수 효율성: LARGO 는 모든 15 개의 전용 모델을 암묵적으로 표현하면서도 단일 전용 모델과 유사한 매개변수 수 (예: BraTS 2018 에서 22.6M) 를 유지합니다. 이는 ShaSpec(187.7M 매개변수) 보다 훨씬 효율적이며 다른 베이스라인과 비교 가능합니다.
• 계산 오버헤드: 동적 가중치 재구성은 단일 nnU-Net 백본에 비해 GFLOPs 를 24.6% 증가시키고 최대 VRAM 을 25.1% 증가시킵니다. 저자들은 결손 모달리티 모델이 일반적으로 추가 연산을 도입하므로 이는 수용 가능하다고 지적합니다.
• 확장성: 이 방법은 $N \lesssim 10$개의 모달리티에 대해 실용적이며, 기하급수적으로 증가하는 압축되지 않은 스템/헤드 레이어가 전체 매개변수의 미미한 부분에 불과하기 때문입니다.

중요성 및 주장

본 논문은 LARGO 가 전용 모델의 전문성과 단일 모델 접근법의 확장성 사이의 간극을 해소한다고 주장합니다. 텐서 분해를 통해 모달리티별 모델 간의 합성곱 커널 구조적 유사성을 활용함으로써, LARGO 는 기하급수적인 저장 및 계산 확장 필요성을 제거합니다.

저자들은 LARGO 를 다음과 같은 원칙적인 하이퍼네트워크 프레임워크로 위치시킵니다:

1. 생성적 임피테이션의 "환각" 위험을 회피합니다.
2. 공유 인코더 아키텍처의 성능 타협을 피합니다.
3. 복잡한 학습 파이프라인 (예: 지식 증류) 없이 최첨단 결과를 달성합니다.

avMNIST 에 대한 개념 증명은 이 방법의 잠재력이 의료 영상을 넘어 결손 데이터가 있는 이질적 모달리티가 포함된 모든 작업으로 확장될 수 있음을 시사합니다. 저자들은 향후 작업이 더 표현력 있는 비선형 하이퍼네트워크 공식화 및 트랜스포머 기반 아키텍처로의 확장을 탐구할 것이라고 결론지었습니다.