자연 이미지 오토인코더가 장기 역학 모델링을 위해 fMRI 볼륨을 콤팩트하게 토큰화할 수 있는가?

고차원 fMRI 데이터는 막대한 메모리 요구량 때문에 긴 시간의 뇌 활동을 한꺼번에 분석하기 어려웠다. 이 논문은 일반 자연 이미지로 학습된 2D 오토인코더가 뇌 영상의 핵심 정보를 매우 효율적으로 압축할 수 있음을 증명하여, 기존보다 훨씬 긴 시간 범위의 뇌 역학을 Transformer로 모델링할 수 있는 길을 열었다.

기존의 fMRI 분석은 뇌 영상을 작은 격자(Voxel) 단위로 직접 처리했다. 하지만 3D 영상이 시간에 따라 쌓이는 4차원 데이터 특성상, Transformer의 Self-Attention 연산에 필요한 메모리가 시퀀스 길이의 제곱에 비례하여 급증하는 문제가 있었다. 이로 인해 기존 모델들은 뇌 활동의 아주 짧은 순간(약 20초 내외)만 관찰할 수 있다는 한계가 있었다.

TABLeT은 이 문제를 '극단적인 압축'으로 해결한다. 핵심 아이디어는 수백만 개의 Voxel로 이루어진 3D 뇌 영상을 단 27개의 벡터 토큰으로 줄이는 것이다. 이를 위해 연구진은 수천만 장의 일반 사진으로 학습된 2D 이미지 압축 모델(DCAE)을 가져와 뇌 영상을 가로, 세로, 높이 세 방향으로 잘라 압축한 뒤 결합했다. 딥러닝의 Embedding 개념을 활용해 복잡한 공간 정보를 고차원 벡터 공간의 점으로 치환한 셈이다.

결과적으로 모델이 처리해야 할 데이터의 양이 획기적으로 줄어들면서, 동일한 하드웨어에서 기존보다 10배 이상 긴 시간의 뇌 활동을 한꺼번에 분석할 수 있게 되었다. 이는 뇌의 장기적인 변화나 복잡한 인지 과정을 파악하는 데 결정적인 역할을 한다.

TABLeT의 핵심은 2D DCAE(Deep Compression Autoencoder)를 활용한 슬라이스 기반 토큰화이다. 3D fMRI 볼륨 X(1xDxHxW)를 RGB 구조로 복제한 뒤, 깊이(D), 높(H), 너비(W)의 세 축을 기준으로 각각 2D 슬라이스 뭉치로 취급하여 인코더에 통과시킨다. 각 슬라이스는 32배의 공간 압축률을 가진 인코더를 거쳐 잠재 표현 Z로 변환된다.

추출된 세 방향의 잠재 표현들은 32x32 크기의 패치 단위로 그룹화되어 연결(Concatenate)된다. [32C' x D/32 x H/32 x W/32] 형태의 텐서로 재구성되며, 최종적으로 각 프레임당 27개의 토큰이 생성된다. 각 토큰은 3072차원의 임베딩 벡터를 가지며, 이는 뇌의 특정 3D 그리드 위치에 대응하는 정보를 담고 있다.

Transformer 인코더는 이 토큰 시퀀스를 입력받아 시공간적 관계를 학습한다. 효율적인 연산을 위해 Grouped Query Attention(GQA)과 Rotary Positional Encoding(RoPE)을 적용했다. 학습 시에는 Masked Token Modeling(MTM)을 사용하는데, 입력 토큰의 50%를 [MASK] 토큰으로 교체한 뒤 원래의 토큰 값을 예측하도록 L1 Loss를 계산하여 최적화한다.

UK-Biobank(UKB), Human Connectome Project(HCP), ADHD-200 등 대규모 데이터셋에서 실험을 수행했다. 성별 분류 태스크에서 TABLeT은 97.7%(UKB), 93.8%(HCP)의 정확도를 기록하며 기존 SOTA 모델인 SwiFT와 대등하거나 우수한 성능을 보였다. 특히 지능(Intelligence) 예측과 ADHD 진단처럼 장기적인 맥락이 중요한 작업에서 입력 프레임 수(T)를 늘릴수록 성능이 뚜렷하게 향상되는 경향을 확인했다.

효율성 측면에서 TABLeT은 압도적인 수치를 기록했다. NVIDIA RTX A6000 GPU 1장 기준, SwiFT가 메모리 한계로 50프레임까지만 처리 가능했던 반면, TABLeT은 동일 메모리 예산에서 약 384프레임까지 처리가 가능했다. 50프레임 기준 SwiFT 대비 메모리 사용량은 7.33배 적었으며, 에포크당 학습 시간은 3.8배 단축되었다.

Ablation Study를 통해 2D 자연 이미지로 학습된 인코더가 뇌 영상 전용으로 학습된 3D 인코더보다 오히려 더 나은 복원 품질(PSNR, SSIM)과 기능적 연결성 보존 능력을 보임을 입증했다. 이는 대규모 자연 이미지 데이터셋에서 학습된 저수준 특징 추출기가 의료 도메인에서도 강력한 범용성을 가짐을 시사한다.

TABLeT 아키텍처는 12개의 Transformer 레이어, 14개의 Attention 헤드, 2개의 KV 헤드로 구성된다. 입력 시퀀스 길이는 기본적으로 256프레임(총 6,912개 토큰)을 처리하도록 설계되었다. 2D DCAE는 수정되지 않은 체크포인트(dc-ae-f32c32-in-1.0)를 사용하여 고정(Frozen) 상태로 유지하며, 이는 별도의 뇌 영상 미세조정 없이도 충분한 표현력을 제공한다.

토큰화 과정에서 세 축의 정보를 통합하는 'Aggregation of 3 Axes' 기법은 특정 방향의 슬라이싱에서 발생할 수 있는 정보 손실을 방지한다. 실험 결과, 단일 축 토큰화보다 세 축을 모두 사용했을 때 성능 편차가 줄어들고 안정적인 결과를 얻었다. 또한, Integrated Gradients(IG)를 이용한 해석 결과, 모델이 성별 분류 시 내측 전전두엽(mPFC)과 시상(Thalamus) 등 실제 뇌 과학적으로 유의미한 영역에 집중하고 있음이 확인되었다.

각 fMRI 프레임을 독립적으로 토큰화하기 때문에 미세한 시간적 역학 정보가 일부 손실될 수 있다. 또한 모든 토큰을 평면적으로(Jointly) 처리하므로 토큰 간의 명시적인 3차원 공간 구조나 시간적 순서에 대한 기하학적 제약이 부족하다는 점이 한계로 지적된다.