자기회귀 MRI 재구성을 위한 다음 가속 스케일 예측

왜 중요한가

MRI 재구성은 데이터가 불충분할 때 ill-posed 문제다. 고가속도 조건에서 픽셀 단위 예측은 고주파 해상 정보를 흐리게 만드는 경향이 있다. 본 연구는 reconstruction을 discrete multi-scale latent space로 옮기고, next-acceleration-scale 예측으로 해결 공간을 제약하여 extremely sparse measurements에서도 선명한 해상 유지가 가능하다고 보인다. Discrete latent formulation은 LLM의 토큰 예측과도 자연스러운 연결고리를 제공한다.

핵심 기여

Discrete autoregressive MRI reconstruction

다중 가속도 스케일의 토큰 hierarchies에서 next-acceleration-scale 토큰을 예측하는 Discrete autoregressive MRI 재구성 프레임워크를 제안한다. 이는 극단적 undersampling하에서도 해부학적으로 의미 있는 고주파 세부를 보존하는 inductive bias를 제공한다.

AQ-VAE with shared codebook across scales

AQ-VAE는 다중 가속도 레벨 간 공유 코드북을 갖고, 입력 레벨별로 토큰 그리드를 조정하여 가속이 큰 입력은 토큰 수가 적고, 해상도가 낮은 입력은 더 풍부한 세부를 제공하도록 구성한다.

Cross-attentive transformer for next-scale prediction

AQ-VAE 인코더에서 추출한 다중 해상도 피쳐를 교차 주의로 주입하는 Transformer를 설계하여, 첫 번째 스케일에서부터 다음 스케일의 토큰을 강하게 가이드한다.

On-policy privileged information distillation

훈련 중 Privileged Context로 Fully-sampled MR 이미지를 활용하는 교사-학생 구조를 도입하고, 역 KL 손실로 학생의 롤아웃 경로에 대한 샘플링 분포를 정렬한다. inference 시 입력은 바뀌지 않는다.

Empirical validation on fastMRI under extreme undersampling

R = 32에서 빠른 MRI(brain) 데이터셋 fastMRI를 대상으로 Cartesian-X/Y, Radial, Gaussian VD 패턴에서 PSNR/SSIM/LPIPS 등의 벤치마크에서 개선을 보인다.

핵심 아이디어 이해하기

기존 VAR은 단일 해상도에서 재구성 토큰을 순차적으로 생성하는 방식이다. 본 연구는 다중 입력 해상도에서 additive latent pyramid를 구성하고, 높은 가속도에서 얻은 정보가 낮은 가속도에서 더 풍부한 정보를 보완하도록 구성한다. AQ-VAE는 가속도 레벨 간 정보를 공유하는 코드북으로 토큰화를 수행하고, Cross-Attentive Transformer가 상위 레벨의 맥락 정보를 이용해 하위 레벨의 다음 토큰을 예측한다. 최종적으로 On-Policy Privileged Information Distillation은 학습 시 교사가 priviledged context를 활용해 롤아웃 경로의 질을 개선하고, inference 시에 발생하는 배포 불일치를 줄인다.

방법론

AQ-VAE 구성: Z_k를 각 가속도 k에서 얻고, 코드북 벡터로 Q_k를 양자화한다. 질의 조건으로 acceleration rate와 sampling pattern를 인코더에 주입하고, 이들 latent를 fusion한 후 디코더로 재구성한다. 토큰 격자는 32×32에서 11×11에서 16×16까지 레벨별로 차등한다. 학습 손실은 reconstruction(SSIM + perceptual LPIPS), L_com(코드북 커밋먼트), L_adv, L_perceptual, L_ssim 등으로 구성한다. 학습은 100에폭, EMA 기반 코드북 업데이트, rotation trick를 사용한다. Distillation은 L_theta = E_q∼p_theta[ 1/|K| sum_k KL(p_theta(·|Q≤k) || p_phi(·|Q≤k, x)) ] 형태의 역 KL을 사용한다.

주요 결과

실험은 fastMRI 멀티코일 뇌 데이터셋에서 R=32로 평가한다. Cartesian-X/Y에서 PSNR, SSIM, LPIPS 측정치가 기존 대비 개선된다. 예를 들어 ES Cartesian-X에서 Proposed의 PSNR은 T1=24.61, T2=19.16, FLAIR=21.29이며, SSIM은 T1=0.76, T2=0.58, FLAIR=0.60, LPIPS는 T1=0.16, T2=0.24, FLAIR=0.22이다. Radial에서 PSNR은 26.00, 20.51, 22.68로 나타났으며 Gaussian-VD에서 Distilled 모델은 Gaussian-VD에서 PSNR 27.86, 22.29, 24.08로 보고된다. 전반적으로 Cartesian 설정에서 PSNR/LPIPS/SSIM의 향상을 보이고, LPIPS 기반의 perceptual 지표에서도 우수한 성능을 보인다.

기술 상세

아키텍처: AQ-VAE(코드북 4096, latent dimension 32, 채널 폭 160, 6-level hierarchy), Token grids: 11×11→16×16, Encoder/Decoder: ResNet 계열 기반에 conditioning FiLM, 다중 스케일 융합. Cross-Attentive Transformer: depth 16, embed_dim 1024, 16 heads, 64×64/32×32/16×16 해상도 피처를 교차-주의로 주입. 학습: base 모델은 100에폭, distillation은 100에폭; AdamW; EMA 코드북 업데이트; rotation trick; BiomedCLIP 기반 디스크리미네이터. 손실: LSSIM, Ladv, Lperc, Lcom; Lcom은 0.25, β=0.25, Z_k와 코드북의 MSE.

실무 활용

극단적 undersampling에서도 고주파 구조를 보존하는 MRI 재구성 파이프라인으로 적용 가능하다.

고속 MRI 프로토콜에서 고해상도 재구성 필요 시
다양한 샘플링 패턴에 대한 견고한 재구성 모델이 필요할 때
의료 영상 합성이나 데이터 보강을 위한 페더레이티드 학습 전처리로 활용

코드 공개 여부: 공개

코드 저장소 보기

키워드

Visual Autoregressive ModelingDiscrete Latent Spacecodebook tokensprivileged information distillationon-policyfastMRIextreme undersampling

용어 해설

AQ-VAE: — AQ-VAE는 입력이 서로 다른 해상도(가속도 수준)에서 얻어진 다중 입력 Latent를 공유 코드북으로 양자화하고, 각 스케일의 정보를 합성해 재구성하는 다중 입력 벡터 양자화 변분 오토인코더이다. 32차원 잠재 표현과 4096 코드북 크기를 가지며, 가속도 수준 및 샘플링 패턴에 따라 인코더의 특징을 조정하는 조건부 인코더를 사용한다.
Cross-Attentive VAR Transformer: — 다중 해상도 특성(64×64, 32×32, 16×16)을 교차 주의(attention)로 주입하는 Transformer로, 각 스케일에서 다음 스케일 토큰을 예측하도록 설계되었다. 각 레이어는 엔코더의 다중 해상도 피쳐를 교차-주의로 받아들여 고해상도 디테일과 전역 구조의 일치를 동시에 달성한다.
Privileged Information Distillation(Privileged Information Distillation): — 훈련 시점에만 이용 가능한 특권 정보를 활용해 학생 모델의 롤아웃에 대한 타깃 분포를 개선하는 온-정책(distilled) 정보 증류 기법이다. 교사는 완전 샘플링된 MR 이미지와 학생의 현재 롤아웃을 함께 보고 분포를 정제하며, 역 KL 손실을 통해 학생이 규정된 경로를 따르도록 유도한다.
비주얼 autoregressive 모델링(Visual Autoregressive Modeling): — 이미지를 Discrete Latent Token의 시퀀스로 표현하고, 각 토큰을 순차적으로 예측하는 방식으로 고해상도 이미지를 생성하는 접근이다. VAR의 확장을 MRI 재구성에 적용하여, 다중 가속도 스케일 간Token 계층에서 Next-Scale 토큰을 예측한다.