ImmuVis: Hyperconvolutional Foundation Model for Imaging Mass Cytometry

기술 요약: ImmuVis

문제 제기
이미징 질량 세포계수법 (IMC) 은 금속 표지 항체를 측정함으로써 조직의 고처리량, 아세포 수준 공간 프로파일링을 가능하게 합니다. 그러나 IMC 데이터를 모델링하는 데 있어 근본적인 과제는 고정된 채널 공간의 부재입니다. 정적인 RGB 채널 세트를 갖는 자연 이미지와 달리, IMC 패널은 물리적 채널 간섭과 코호트별 실험 설계로 인해 연구 간에 크게 변합니다. 이 이질성은 표준 비전 백본의 핵심 가정인 고정된 입력 차원을 요구한다는 점에 위배됩니다. 기존 접근법인 U-Net 은 고정된 마커 패널로 제한되며, 최근의 트랜스포머 기반 모델 (예: VirTues, EVA) 은 마커별 공간 토큰과 식별 인코딩을 연결하여 가변 입력을 처리하려 시도합니다. 그러나 이러한 토큰 기반 방법들은 종종 중요한 국소 마커 의존성을 무시하며, 마커 수와 공간 해상도가 증가함에 따라 높은 계산 비용을 초래합니다. 또한, 증거가 약할 때 가상 염색 (측정되지 않은 마커 예측) 은 잘 정의되지 않은 문제가 됩니다. 현재 모델들은 종종 보정된 불확실성 추정을 결여하고 있어, 생물학적으로 타당한 재구성과 오류 있는 예측을 구별하기 어렵게 만듭니다.

방법론
저자들은 가변 입력 학습, 계산 효율성, 불확실성 정량을 통합하도록 설계된 ImmuVis라는 기초 모델 계열을 소개합니다. 핵심 혁신은 마커 적응형 하이퍼컨볼루션 아키텍처입니다.

1. 하이퍼컨볼루션 아키텍처:

   • 마커 무관 인코딩: 입력 마커 채널은 먼저 공유 컨볼루션 스템 ($Enc_{ma}$) 에 의해 독립적으로 처리되어 국소 공간 패턴을 포착합니다.
   • 조건부 하이퍼컨볼루션: 고정된 가중치 대신, 모델은 학습된 마커 임베딩에 조건부로 컨볼루션 커널을 생성하는 하이퍼네트워크를 사용합니다. 마커 집합 $I_e$에 대해, 각 마커 $i \in I_e$에 대한 학습된 커널 생성기 $\phi_e(i)$를 연결하여 하이퍼커널 $H^e_{I_e}$를 구성합니다. 이를 통해 모델은 아키텍처 변경 없이 관찰된 마커의 특정 하위 집합에 맞게 특징 추출 연산자를 동적으로 적응시킬 수 있습니다.
   • 전체 마커 백본: 하이퍼컨볼루션의 출력은 표준 컴퓨터 비전 백본 (예: ConvNeXt v2, Swin Transformer, ViT) 이 처리할 수 있는 통합된 고정 크기 임베딩입니다.
   • 대칭 디코더: 디코더는 인코더를 반영하여, 잠재 표현을 특정 목표 마커로 매핑하기 위해 마커 조건부 하이퍼컨볼루션을 사용하며, 평균 발현과 로그 분산 (불확실성) 을 모두 예측합니다.

2. 학습 전략:

   • 데이터셋: 모델들은 현재까지 가장 큰 IMC 데이터셋인 IMC17M에서 사전 학습되었습니다. 이 데이터셋은 28 개 코호트에 걸쳐 24,405 개의 이미지, 265 개의 마커, 1700 만 개 이상의 패치로 구성됩니다.
   • 목적: 모델들은 가aussian 이분산 회귀 손실을 사용한 마스킹 채널 재구성 작업으로 학습됩니다. 이 목적 함수는 목표 픽셀을 독립적인 가우시안 분포로 모델링하여 네트워크가 평균 ($\mu$) 과 분산 ($\sigma^2$) 을 모두 예측할 수 있게 합니다. 분산 예측은 재구성 오차와 상관관계가 있는 보정된 불확실성 추정을 제공합니다.
   • 전처리: 저자들은 이온 수 역학을 처리하기 위한 arcsinh 변환, IMC 특유의 노이즈를 줄이기 위한 버터워스 저역 통과 필터링, 그리고 최소 -99 백분위수 스케일링을 포함한 특정 전처리 파이프라인을 사용합니다.

주요 기여

• 하이퍼컨볼루션 설계: 학습된 마커 임베딩에서 컨볼루션 커널을 생성하는 새로운 아키텍처로, 재학습이나 아키텍처 수정 없이 단일 모델이 임의의 마커 하위 집합을 처리할 수 있게 합니다.
• 효율적인 기초 모델 계열: 확립된 백본을 활용하는 세 가지 변형 (ImmuVisConv, ImmuVisSwin, ImmuVisViT) 의 구현입니다. 이러한 모델들은 토큰 연결 기반의 트랜스포머 대안들에 비해 현저히 낮은 계산 비용으로 최첨단 성능을 달성합니다.
• 불확실성 인식 학습: 희귀 표현형이나 비정형 패턴에 내재된 가상 염색의 신뢰성 문제를 해결하는 이분산 가능도 목적 함수를 통합하여 픽셀 단위 불확실성 추정을 생성합니다.
• 대규모 사전 학습: 다양한 조직학과 마커 패널에 걸쳐 강건하고 전이 가능한 표현을 학습할 수 있도록 IMC17M 의 큐레이션과 활용을 촉진합니다.

결과

• 가상 염색: ImmuVis 변형들은 가상 염색 작업에서 최첨단 베이스라인 (VirTues, EVA) 보다 우수한 성능을 발휘합니다. 구체적으로, ImmuVisSwin과 ImmuVisConv는 가장 낮은 평균 제곱 오차 (MSE) 와 가장 높은 피어슨 상관관계를 달성하여 토큰 기반 방법들에 비해 공간 구조 보존이 우수함을 입증했습니다.
• 불확실성 보정: 이분산 헤드는 재구성 오차와 강력하게 상관관계가 있는 (피어슨 상관 계수 > 0.92) 불확실성 추정을 생성합니다. 이러한 모델들은 의미 있는 픽셀 단위 신뢰도 신호를 제공하며, ImmuVisSwin 은 가장 낮은 희석 오차 면적 (AUSE) 을 달성했습니다.
• 표현 학습: 하류 작업에서 ImmuVis 임베딩은 세포 유형 분류 (ImmuVisConv 는 Macro-F1 0.709 달성) 및 임상 예측 (ImmuVisSwin 은 Macro-F1 0.641 달성) 에서 최첨단 성능을 보여주어 VirTues 및 EVA 변형들을 능가했습니다.
• 계산 효율성: ImmuVis 모델들은 현저히 낮은 계산 비용으로 작동합니다. ImmuVisConv 는 순방향 통과당 약 89.9 GFLOPs 를 요구하며, 이는 VirTues (1049.3 GFLOPs) 보다 한 자릿수 낮고 EVA (3590 GFLOPs) 보다 거의 40 배 낮습니다. 추론 시간도 이에 비례하여 더 빠릅니다 (예: VirTues 의 28.5 ms 대비 13.4 ms).

의의
본 논문은 패널 이질성과 계산 효율성의 중요한 제약을 해결하는 실제 IMC 모델링을 위한 실용적인 프레임워크로 ImmuVis 를 위치시킵니다. 재학습 없이 임의의 마커 하위 집합에서 작동할 수 있는 단일 기초 모델을 가능하게 함으로써, ImmuVis 는 대규모 공간 조직 프로파일링과 레거시 코호트의 사후 업그레이드를 촉진합니다. 보정된 불확실성의 통합은 임상 응용에서 위양성이나 낮은 신뢰도 예측이 표시되어야 하는 가상 염색에 대한 신뢰성 인식 접근법을 제공하며, 이는 매우 중요합니다. 저자들은 이러한 연산자 수준의 채널 적응성이 멀티플렉스 이미징 및 그 이상에서 기초 모델을 구축하기 위한 일반적인 레시피를 제공하며, 단일 아키텍처 내에서 유연성, 효율성, 신뢰성을 결합한다고 제안합니다.