기술 요약: 제한된 데이터로 효율적인 토큰화를 위한 스펙트럼 비전 트랜스포머

문제 정의

현재 의료 영상 분야에서 비전 트랜스포머 (ViT) 의 적용은 막대한 데이터 요구량으로 인해 방해받고 있습니다. 기존의 공간적 ViT 는 국소 패치 토큰화와 자기-주의 (self-attention) 메커니즘에 의존하는데, 이는 과적합을 방지하고 견고한 수렴을 보장하기 위해 일반적으로 수만 장의 훈련 이미지를 필요로 합니다. 임상 환경에서는 데이터셋이 종종 작고, 선별적이며, 분포 변화에 노출되어 있어 표준 ViT 는 일반화에 어려움을 겪습니다. 전이 학습이나 토큰 기반 증류와 같은 기존 완화 전략은 종종 방대한 외부 데이터셋이나 사전 훈련된 동반 모델에 의존하는데, 이는 분포 변화에 대한 취약성과 규제적 도전을 초래합니다. 결과적으로, 침습적 시술, 공격적인 질환, 희귀 질환의 예후를 예측하기 위한 ViT 의 임상적 전환은 제한된 상태입니다.

방법론

저자들은 표준 공간적 패치 토큰화를 스펙트럼 분해 접근법으로 대체하는 새로운 아키텍처인**스펙트럼 비전 트랜스포머 (Spectral ViT)**를 제안합니다. 원시 이미지 픽셀에서 잠재 토큰으로의 선형 프로젝션을 학습하는 대신, Spectral ViT 는 이미지 토큰화를 위해 고정된 기저 (주요 실험에서는 특히 주성분 분석, PCA) 를 활용합니다.

핵심 아키텍처

1. 스펙트럼 토큰화: 이미지들은 훈련 데이터에서 유도된 고정된 주성분 고유벡터 집합 ($w_i$) 으로 투영됩니다. 결과적으로 생성된 스칼라 계수 ($s_i = v \cdot w_i$) 가 시퀀스 토큰으로 사용됩니다.
2. 계층적 임베딩: 그리드 좌표로 토큰이 순서화되는 공간적 ViT 와 달리, 스펙트럼 토큰은 그 순위 (설명된 분산) 에 따라 순서화됩니다. 분산 계층 구조 가중치 ($\omega_i = 1/i$) 와 스펙트럼 위치 임베딩 ($e_{pos}$) 이 임베딩 공간으로 투영되기 전에 계수에 적용됩니다.
3. 자기-주의: 스펙트럼 임베딩의 시퀀스는 표준 자기-주의를 거칩니다. 이를 통해 모델은 선형 분해 방법만으로는 불가능했던 스펙트럼 성분 간의 비선형적이고 작업 관련 상호작용을 학습할 수 있습니다.
4. 기저 일반화: PCA 가 주요 초점이지만, 이 프레임워크는 공간 불변성을 위한 푸리에 기저와 비유클리드 다양체 위 위상 보존을 위한 라플라시안 기저를 포함한 다른 기저들을 수용하도록 설계되었습니다.

복잡도 분석

Spectral ViT 는 이미지 해상도 ($m$) 를 트랜스포머 시퀀스 길이 ($n$) 에서 분리합니다.

• 공간적 ViT: 시퀀스 길이가 부피에 비례하여 선형적으로 증가합니다 ($N_P \propto V^3$), 이는 $O(V^6)$의 2 차적 주의 복잡도로 이어집니다.
• 스펙트럼 ViT: 시퀀스 길이는 패치 수보다 수직으로 작은 스펙트럼 성분의 수 ($n$) 로 고정됩니다. 이는 트랜스포머 백본 복잡도를 $O(n^2)$로 줄여 3 차원 의료 영상에 대한 계산 비용을 크게 낮춥니다.

이론적 특성

• 최적 저랭크 근사: PCA 기저는 주어진 랭크에 대해 최소 제곱 재구성 오차를 제공합니다.
• 에너지 압축: PCA 는 고정된 시퀀스 길이에 대해 신호 분산을 최대화하여, 중복 정보를 제거함으로써 공간적 패치에 비해 더 높은 신호 대 잡음비 (SNR) 를 제공합니다.
• 공간 불변성: 푸리에 기저를 사용할 때, 모델은 스펙트럼 계수가 공간적 위치가 아닌 크기 모드를 나타내므로 엄격한 공간 (이동) 불변성을 달성합니다.

주요 기여

1. 새로운 아키텍처: 제한된 데이터 환경에 특히 적합하도록 설계된, 계승된 계층적 임베딩을 가진 스펙트럼 분해를 통해 이미지를 토큰화하는 ViT 프레임워크의 도입.
2. 이론적 정당성: 스펙트럼 기저의 선택에서 비롯된 계산 복잡도 감소와 최적의 신호 대 잡음비 증명.
3. 실증적 검증: 시뮬레이션, 공개, 임상 데이터셋에 대한 종합적 평가로, 표준 ViT, 주의가 있는 합성곱 신경망 (CNN), Swin ViT, 그리고 로지스틱 회귀에 비해 훨씬 적은 매개변수로 동등하거나 우수한 성능을 보임.
4. 해석 가능성: 모델의 중요도 맵 (saliency maps) 은 표준 모델들이 놓칠 수 있는 생물학적으로 관련 있는 구조 (예: 심부 뇌 자극 예측에서의 상소뇌다리) 를 강조합니다.

실험 결과

저자들은 Spectral ViT 를 세 가지 다른 작업에 대해 평가했습니다:

1. 패턴 분류 (시뮬레이션 데이터):

   • 다양한 신호 대 잡음비 (SNR) 를 가진 이진 분류 작업에서, Spectral ViT 는 저데이터 환경에서 공간적 ViT 보다 우수한 성능을 보였습니다.
   • 고 SNR 에서는 32 개의 샘플만으로 동등한 성능에 도달했습니다. 낮은 SNR 에서는 공간적 ViT 가 Spectral ViT 의 성능과 일치하기 위해 1,000 에서 100,000 개 이상의 샘플이 필요했습니다.

2. 객체 분류 (분포 변화):

   • 훈련 세트에서 객체 거리가 위치와 우연히 상관관계를 가진 작업에서, 공간적 ViT 는 공간적 위치에 대한 의존성으로 인해 실패했습니다.
   • Spectral ViT 는 푸리에 기저를 사용하여 공간 불변성을 활용함으로써 우연한 위치 상관관계를 무시하고 완벽한 분리를 달성했습니다.

3. 성별 분류 (IXI 데이터셋):

   • T1 가중치 IXI 데이터셋 (566 명의 피험자) 을 사용하여 Spectral ViT 는 약 13,000 개의 매개변수로 AUC 0.842를 달성했습니다.
   • 이는 표준 공간적 ViT (AUC 0.691) 와 컴팩트 공간적 ViT (AUC 0.616) 보다 우수한 성능이었습니다.
   • 데이터가 등록되지 않았음에도 불구하고, 주의가 있는 U-net 과 Swin ViT 에 비해 매개변수 백만 개당 AUC 측면에서 우수한 효율성을 보여주었습니다.

4. 신경자극 후보 분류 (임상 데이터):

   • 103 명의 환자로부터의 정량적 감수성 지도 (QSM) 를 사용하여 심부 뇌 자극 (DBS) 결과 예측에 적용되었습니다.
   • Spectral ViT 는 로지스틱 회귀와 임상 트랜스포머를 포함한 테스트된 모델들 중에서 가장 높은 균형 정확도 (0.783) 와 특이도 (0.977) 를 달성했습니다.
   • 중요도 맵은 DBS 에 의해 조절되는 알려진 신경 운동 회로인 상소뇌다리 (SCP) 백질 다발을 올바르게 식별한 반면, 다른 모델들은 흑질의 철에 집중했습니다.

중요성과 주장

이 논문은 Spectral ViT 가 데이터가 부족한 의료 영상 분야에서 트랜스포머 기반 아키텍처를 배포할 수 있는 실현 가능한 경로를 제공한다고 주장합니다. 스펙트럼 분해 (특히 PCA) 의 전역적 귀납적 편향을 활용함으로써, 이 모델은 막대한 데이터셋이나 복잡한 사전 훈련 없이도 특징 추출을 효율적으로 매개변수화합니다.

저자들은 그들의 접근 방식이 다음과 같다고 강조합니다:

• 데이터 요구 감소: 수만 장이 아닌 수백 장의 이미지로 데이터셋에서 효과적인 훈련을 가능하게 합니다.
• 효율성 향상: 성능을 유지하거나 향상시키면서 계산 복잡도와 매개변수 수를 극적으로 줄입니다.
• 강건성 강화: 공간적으로 불변한 토큰화 (푸리에 기저 사용 시) 를 통해 분포 변화 및 우연한 상관관계와 관련된 문제를 완화합니다.
• 해석 가능성 제공: 토큰의 스펙트럼 특성은 임상 지식과 일치하는 생물학적으로 관련 있는 특징을 식별할 수 있게 합니다.

저자들은 이 방법이 고정된 선형 프로젝션을 통해 편향을 도입하지만, 의료 AI 의 데이터 부족이라는 병목 현상을 효과적으로 해결하며, 현재 ViT 의 임상적 광범위한 채택을 방해하는 계산 및 데이터 장벽을 완화할 수 있다고 결론지었습니다. 그들은 향후 연구가 임상 구현 전에 최적의 기저 구성과 전이 학습 기술을 통한 검증에 초점을 맞춰야 한다고 지적합니다.