MinD-Vis
Seeing Beyond the Brain: Conditional Diffusion Model with Sparse Masked Modeling for Vision Decoding CVPR2023 1National University of Singapore, Center for Sleep and Cognition, Centre for Translational Magnetic Resonance Research 2The Chinese University of
mind-vis.github.io
본 리뷰는 해석과 내용에 대한 분석을 함께 포함하고 있다.
CVPR2023에 선정된 논문으로, fMRI to Image 분야에서는 현재까지 state-of-the-art를 차지하고 있다. 이 논문을 참고하여 data reference를 EEG로 바꾼 DreamDiffusion 이라는 논문도 존재하고 있으니, 참고 해 보면 좋을 것 같다.
Abstract
시각 자극에 대한 뇌 신호를 해석하는 것은 인간의 시각계에 대한 이해를 깊게 하고 BCI(Brain Computer Interface)를 통한 인간과 컴퓨터 비전 사이 연결의 기반을 튼튼히 하는 것을 목표로 합니다. 하지만, 뇌파 속에 숨은 복잡한 표현들과 data annotation의 부족으로 올바른 의미를 포함하는(정답에 가까운) 높은 퀄리티의 이미지를 재구성 하는 것은 쉬운일이 아닙니다.
이 논문에서, 우리는 MinD-Vis: Sparse Masked Brain Modeling with Double-Conditioned Latent Diffusion Model for Human Vision Decoding(이중 조건 처리 된 Latent Diffusion Model을 활용한 희소 마스킹 뇌 모델링을 통한 인간 시각계 해석)을 제안합니다.
먼저, sparse하게(희소하게) 부호화한 일차 시각 피질(primary visual cortex)의 정보에 영감을 받아, 거대 잠재 공간(latent space)에 대한 마스킹 모델링을 활용하여 fMRI 데이터의 효과적인 self-supervised(자기 지도) 표현을 배울 수 있었습니다.(i. e. latent space를 predection 하는 self-supervised learning을 활용해서 fMRI data로부터 알맞은 feature과 reconstruction 능력을 갖추는 표현을 얻어내겠다. 이 논문에서는 MAE 활용함)
그리고 LDM(latent diffusion model)을 double-conditioning 방식을 사용하여 증강함으로써, 매우 적은 annotation pair를 가지고도 Mind-Vis는 뇌파를 통해 의미적으로도 일치하고 매우 그럴듯한 이미지를 재구성 할 수 있다는 것을 증명했습니다.
연구에서는 모델에 대한 질적, 양적 평가를 완료하였으며, semantic mapping과 generation quality의 측면에서 모두 state-of-the-art의 성능을 보이는 것을 확인 했습니다.
Introduction
“What you think is what you see”
인간의 지각과 선험적인 지식은 한 사람의 마음속에 깊게 얽혀있습니다. 세상에 대한 우리의 인지는 객관적 사물에 대한 자극이 아니라, 우리의 경험들과 세상을 지각하는 과정 뒤에 숨어있는 뇌의 복잡하고 정밀한 활동들에 의해 이루어집니다. 이러한 뇌의 활동을 이해하는 것과 부호화된 정보들을 해석하는 것은 인지 뇌과학 측면에서 가장 중요한 목표입니다. 이러한 끝을 알 수 없는 해석을 하기 위하여, 많은 연구에서 시각적 정보를 해독하는 것에 초점을 맞추고 있습니다.
비침습적이고 효과적인 방법으로 직접적인 뇌 활동을 측정하는 수단 중 하나인 fMRI는 image classes와 같은 시각적 정보를 해석하는 데에 주로 사용됩니다. 논문에서는 최근 딥러닝 모델들의 도움으로 fMRI를 활용하여 해당 신호에 대응되는 시각적인 자극을 그대로 복구 할 수 있다면 흥미로울 것이라고 가정 했습니다.
하지만 fMRI-image 쌍의 데이터셋 부족과 fMRI로부터 복잡한 신경계 활동을 해석하는 생물학적인 원리의 부족으로, 재구성된 이미지는 대체적으로 흐리고 불명확 하며, 원하는 의미 또한 잡아내지 못했습니다.
따라서 적은 데이터셋만으로도 시각적 자극과 뇌 활동을 연결할 수 있도록 fMRI에 대한 생물학적으로 명확하고 효과적인 표현을 아는 것이 매우 중요했습니다. 또한, 뇌신호에 있어서 개인차가 분명하게 존재하고, 이러한 차이는 같은 시각적인 자극에 있어서도 다른 뇌 활동 패턴을 보여줍니다.
이 논문에서는 위와 같은 상황에서 두가지의 방법을 사용하여 brain signal domain으로부터 image를 생성하는 과정을 제시한다.
1. Sparse-Coded Masked Brain Modeling (SC-MBM)
ViT, MAE와 동일하게 fMRI 신호를 1D sequence로 만들어서 Patch로 자른 후에, 일정 비율(이 논문에서는 75%) 만큼을 Masking 한 후 그 부분을 reconstruction 하도록 pre-training 과정을 구성한다. 뇌에 있는 시각피질(visual cortex)가 들어오는 자극에 대해서 모든 정보를 저장하는 것이 아니라, 가중합하여 정보를 복원하는 sparse coding을 사용하는 것에 영감을 받았다고 한다.(라고 하지만, 그냥 MAE에 brain signal을 적용하기 위해서 나름의 명분찾기(?)를 한 것 같다..)
2. Augmenting the latent diffusion model with double conditioning (DC-LDM)
Double Conditioning은 Latent Diffusion Model에 ViT Encoder가 압축한 Latent Vector를 Cross Attention하여 넣을 때, 시계열 데이터인 fMRI 데이터의 특성을 살리고자 Time Embedding을 Cross Attention과 같은 크기로 이어 붙여서 U-Net에 넣어준다.
Methodology
Stage A, SC-MBM
fMRI 데이터는 시각 피질(visual cortex)를 ROI로 설정 했을 때, 4500개 정도의 voxel이 나오게 되고, 그 voxel을 1차원 텐서로 쭉 펴주면 위와 같은 원본 데이터가 나오게 된다. 입력 형식이 MAE에서는 이미지의 패치 Sequence였다면, 여기서는 fMRI의 패치 Sequence라는 것 말고는 거의 차이가 없다.
몇가지 차이가 존재하는데, MAE에서는 패치 embedding dimenstion이 768이었지만, 여기에서는 1024로 증가 한 것, 그리고 1D Embedding을 진행 할 때, brain signal을 MAE의 입력에 맞는 형식으로 맞춰주는 코드가 추가 된 것 정도가 있다. Mask ratio가 75%일 때 가장 좋은 성능을 내는 것은, brain signal 또한 image와 같이 자연계의 신호이기 때문에, 높은 마스킹 비율이 필요했기 때문이라고 설명한다. 자세한 것은 MAE 논문을 참고 하길 바란다.
또한, 저자들은 embedding-to-patch-size ratio 즉, 임베딩 차원과 패치사이즈의 비율이 MAE에서는 16*16*3 : 768 = 1 : 1 이었지만, 논문에서는 1024로 증가시켰기 때문에, brain signal의 representation space가 증가했다고 주장한다.(정말 그런지는 잘 모르겠다)
결론적으로 Stage A가 가장 중요한 단계이지만, 사실상 BERT에서 Natural Language의 Masking 된 부분을 prediction 하는 방식으로 문장의 표현력을 모델이 학습했던 것 처럼, MAE의 입력 구조에 Brain signal을 끼워맞춰 brain signal에 대한 표현력을 습득한 사전학습 된 모델을 만든것이라고 할 수 있다.
Stage B, DC-LDM
DC-LDM은 간단하다. 그저 U-Net에 들어가는 CrossAttention 바로 다음 layer에 시간을 embedding해서 만든 layer를 concatination한다. 최종적인 Loss function은 다음과 같다.
Experiments
Dataset은 image-fMRI pair로 annotation이 필요가 없는 Stage A에서 unlabeled fMRI dataset인 Human Connectome Project의 dataset과 Generic Object Dataset(GOD)의 dataset을 사용했고, Stage B에서는 Brain, Object, Landscape Dataset(BOLD500)을 사용했다.
또한, SC-MBM 방식을 사용한 것이 사용하지 않았을 때보다 월등한 성능을 보여준다고 주장한다.
Conclusion
가장 최신의 그리고 최고의 기술들인 Latent Diffusion Model(Stable Diffusion)과 MAE를 활용해서 brain signal domain과 image data domain간의 cross generative model을 시도한 의미있는 논문이라고 생각한다. brain signal을 패치로 잘라서 LLM과 같이 pre-training한 후에, generative downstream task에 적용할 때는, supervised learning을 사용 하는 것이, 최근 딥러닝 논문들의 트렌드를 그대로 반영하려고 노력했다는 생각이 들었다. 하지만 아직 Dataset의 크기가 비교적으로 매우 작기 때문에, 과연 brain signal domain에서도 scalable한 model 학습이 가능할 것인지에 대한 생각이 든다.
Natural Language를 제외한 다양한 Time series와 같은 Sequential dataset들에 대한 scalable pre-training model에 대한 연구 중에 매우 흥미로운 분야를 다룬 논문이라고 생각한다.