Hi-DREAM: Brain-Inspired Hierarchical Diffusion for fMRI-to-Image Reconstruction via ROI Encoder and VisuAl Mapping

기술 요약: Hi-DREAM

문제 정의

기능적 자기공명영상(fMRI) 데이터로부터 자연 이미지를 재구성하는 것은 신경 활동과 현대적 생성 모델이 사용하는 구조적 및 의미적 표현 사이의 간극을 메우는 것을 요구한다. 최근의 확산 기반 디코더(diffusion-based decoders)들은 시각적 품질을 개선해 왔으나, 결정적인 한계가 여전히 존재한다: 기존의 대부분의 방법들은 fMRI 신호를 하나의 전역적 임베딩(global embedding)으로 압축하여 생성 모델에 조건화(conditioning)한다. 이러한 접근 방식은 인간 시각 피질의 계층적 조직을 무시하며, 이로 인해 서로 다른 피질 영역이 특정 측면(예: 국소적 구조 vs. 고차원 의미론)에 어떻게 기여하는지 조사하는 것을 어렵게 만든다. 또한, 시각 피질을 단일한 실체로 취급하는 것은 초기(망막 위치 정보), 중간(윤곽/부분), 후기(의미론) 시각 영역이 이미지 생성 과정의 서로 다른 단계에서 수행하는 고유한 역할을 활용하지 못한다.

방법론

저자들은 시각적 계층 구조에 따라 fMRI 조건화를 구조화하는 프레임워크인 Hi-DREAM(Brain-Inspired Hierarchical Diffusion)을 제안한다. 이 방법은 잠재 확산 백본(latent diffusion backbone)에 통합된 두 가지 주요 구성 요소로 이루어져 있다.

1. 계층적 ROI 어댑터 (Hierarchical ROI Adapter)

Hi-DREAM은 단일 임베딩 대신, 시각적 관심 영역(ROI)을 기준으로 fMRI 반응을 세 가지 별도의 스트림으로 그룹화한다:

• 초기 스트림(Early Stream): 국소적이고 망막 위치 정보적인 구조와 관련된 V1/V2를 포함한다.
• 중간 스트림(Middle Stream): 윤곽 및 객체의 부분과 관련된 V3/hV4를 포함한다.
• 후기 스트림(Late Stream): 범주 수준의 의미론과 관련된 LOC/FFA를 포함한다.

ROI 어댑터는 이러한 시행별(trial-wise) 반응과 피험자별 마스크를 다중 스케일 피질 피라미드로 변환한다. 이는 다음을 채택한다:

• 스케일 공간 구축(Scale-space construction): 소프트 ROI 마스크를 U-Net 해상도(64, 32, 16)에 맞게 다운샘플링한다.
• 학습 가능한 게이팅(Learnable gating): 깊이 의존적 가중치($w_{r,s}$)와 혼합 가중치($\alpha_{g,s}$)를 통해 각 U-Net 스케일이 초기, 중간, 후기 증거의 특정 혼합물을 받을 수 있도록 한다.
• 멀티 헤드 잠재 어텐션(Multi-Head Latent Attention, MHLA): 이 모듈은 세 스트림 간의 협력적 상호작용을 모델링하여, 한 ROI 그룹의 정보가 특정 디노이징 깊이에서 다른 그룹이 활용되는 방식에 영향을 미칠 수 있도록 한다.

2. ROI 조건부 컨트롤넷 (ROI-Conditioned ControlNet)

경량화된 컨트롤넷 브랜치는 어댑터에서 생성된 해부학적 인지 프라이어(anatomy-aware priors)를 확산 백본에 주입한다.

• 깊이 매칭 주입(Depth-Matched Injection): 컨트롤넷은 피질 피라미드를 힌트($\xi_s$)로 변환하여, 잔차(residual) 및 FiLM(Feature-wise Linear Modulation) 퓨전을 통해 해당 U-Net 깊이에 주입한다.
• 선택적 가이드(Selective Guidance): 초기 스트림은 얕은 층을 가이드하여(공간적 구조 보존), 후기 스트림은 깊은 층을 가이드하여(의미론 정교화) 이미지를 정교화한다.
• 시간적 스케줄링(Temporal Scheduling): 가이드 강도는 확산 시간 동안 스케줄링되어, 구조적 힌트가 초기 디노이징 단계에서 지배적이고 의미적 힌트가 후기 단계에서 이미지를 정교화하도록 한다.

훈련 목적 함수는 표준 확산 손실과 함께 픽셀 일관성 및 엣지 보존을 보장하기 위한 구조적 정규화 항($L_1$, 그래디언트 보존, SSIM)을 결합한다.

주요 기여

1. Hi-DREAM 프레임워크: fMRI 조건화를 인간의 시각 피질의 계층적 조직을 모방하여 초기, 중간, 후기 시각 ROI 스트림으로 구조화하는 새로운 확산 프레임워크를 제시한다.
2. 아키텍처 설계: 피질 반응을 압축된 다중 스케일 프라이어로 변환하는 ROI 어댑터와, 이 프라이어를 특정 U-Net 깊이에 주입하여 전체 3D 볼륨 조건화의 계산 오버헤드 없이 해부학적 인지 가이드를 가능하게 하는 깊이 매칭 ROI-ControlNet을 설계하였다.
3. 해석 가능성 및 속성 규명(Interpretability and Attribution): 이 프레임워크는 피질 스트림과 생성 깊이 사이의 명확한 매핑을 제공한다. 절제 연구(Ablation studies)는 계층 인식 조건화가 효과적일 뿐만 아니라 해석 가능하다는 것을 보여주며, 서로 다른 ROI 스트림이 상호 보완적인 기여를 한다는 것을 입증한다.

실험 결과

이 방법은 피험자 1, 2, 5, 7을 대상으로 **Natural Scenes Dataset (NSD)**에서 평가되었다.

• 정량적 성능: Hi-DREAM은 Inception 및 CLIP 점수를 포함한 고차원 의미론적 지표에서 최첨단(SoTA) 성능을 달이트며, MindEye2, Psychometry, MindBridge와 같은 이전 방법들을 능가했다. 또한 강력한 저차원 구조적 충실도를 유지하여 최고의 AlexNet(2) 점수와 두 번째로 높은 PixCorr 및 Alex넥(5) 점수를 달성했다.
• 정성적 성능: 재구성된 이미지는 기존 확산 디코더들보다 객체의 정체성, 장면 레이아웃, 그리고 두드러진 시각적 속성(예: 색상 충실도 및 형태 기하학)을 더 일관되게 보존했다.
• 절제 연구(Ablation Studies):
  • 스트림 구성: 세 스트림(초기 + 중간 + 후기)을 모두 사용했을 때 구조적 및 의미론적 지표의 최적의 균형을 보여, 이들의 상호 보완적 성격을 확인했다.
  • 계층 검증: ROI를 스트림에 무작위로 할당하거나 계층을 역전시켜(초기 ROI를 깊은 층에 주입) 입력했을 때 성능이 크게 저하되었으며, 이는 피질 계층을 U-Net 깊이와 일치시키는 것의 중요성을 입증한다.
  • 모듈 기여도: ROI 어댑터나 MHLA 모듈을 제거하면 성능이 대폭 하락하여, 공간적 프라이어와 스트림 간의 교차 조정이 모두 필수적임을 증명했다.
  • 깊이 정렬: 주입된 에너지 분석을 통해, 모델이 얕은 층에서는 초기 스트림에, 중간 층에서는 중간 스트림에, 깊은 층에서는 후기 스트림에 의존하도록 학습되었음을 확인했다.

의의 및 주장

본 논문은 Hi-DREAM이 fMRI 신호를 차별화되지 않은 전역적 임베딩으로 취급하는 것에서 벗어나, 구조화된 뇌 모사 접근 방식을 활용하는 패러다임의 전환을 나타낸다고 주장한다. 조건화 신호를 시각 피질의 자연스러운 계층 구조 및 확산 모델의 깊이와 정렬함으로써, 이 프레임워크는 구조적 무결성을 유지하면서도 우수한 의미론적 재구성을 달성한다.

저자들은 이 프레임워크가 해석 가능성을 제공한다는 점을 강조한다: 각 조건화 스트림은 정의된 피질 ROI 그룹에 대응하므로, 연구자들이 특정 뇌 영역이 이미지 생성의 서로 다른 단계에 어떻게 영향을 미치는지 조사할 수 있다. 현재 구현은 피험자 특정적이며 압축된 시각 ROI 세트에 집중되어 있지만, 이 연구는 더 신경학적으로 정렬되고 효율적이며 해석 가능한 뇌-이미지 디코딩을 위한 토대를 마련한다. 저자들이 식별한 향후 방향에는 미세한 시각적 세부 사항 개선, 배측(dorsal) 및 운동 민감 영역을 포함한 해부학적 범위 확장, 그리고 피험자 간 일반화 전략 개발이 포함된다.