Patch size and its effect on representations in MAEs

Motivation

• Known: MAE는 self-supervised representation learning에 널리 쓰이지만 전통적으로 고정된 patch size를 사용하며, FlexiViT 같은 연구는 supervised 세팅에서 kernel interpolation을 통해 다양한 patch size를 수용할 수 있음을 보였다.
• Gap: 의료 영상처럼 미세 텍스처부터 거친 의미적 랜드마크까지 여러 공간 스케일의 정보가 공존하는 도메인에서, 고정 patch size MAE는 patch보다 작은 물체를 인코딩하지 못하며, 스케일별로 별도 모델을 학습하는 것은 계산 비용이 과도하다는 문제가 있다.
• Why: 단일 사전학습 모델로 미세한 디테일과 거친 의미 구조를 동시에 추출할 수 있게 되면, 다양한 스케일의 병변·해부학적 구조를 다뤄야 하는 의료 영상 표현 학습의 효율성과 성능을 크게 개선할 수 있다.
• Approach: MNIST-on-ImageNet 벤치마크와 실제 lateral DXA bone scan에 대해, patch size를 무작위로 샘플링하며 학습하되 patch size를 MLP를 통해 LayerNorm의 scale/shift 파라미터에 매핑하여 encoder가 명시적으로 스케일을 인지하도록 하는 Mosaic MAE를 제안하고 FlexiMAE와 비교한다.

Achievement

Figure 2. Downstream linear probing performance on clinical

1. Patch size의 고역통과 필터 효과 실증: p<s일 때 linear probing accuracy가 p>s일 때보다 일관되게 높음을 보여 patch size가 인코딩 가능한 물체의 최소 스케일을 결정함을 확인했다.
2. Mosaic MAE 제안: FlexiViT의 동적 patch resizing(선형 투영 가중치 및 위치 임베딩의 동적 리사이즈)에 diffusion model/DiT 스타일의 adaptive LayerNorm 조건화를 결합해 self-supervised 생성 모델에 확장했다.
3. FlexiMAE 대비 우수한 표현력: 명시적 patch-size conditioning이 조건화되지 않은 variable-patch 베이스라인(FlexiMAE)보다 더 나은 downstream representation을 제공함을 clinical DXA 데이터에서 입증했다.

How

Figure 1. Left & Middle: MAE pre-training on the MNIST-on-ImageNet benchmark. Smaller patches (top) preserve fine-graine

• 4x4 corner 위치에 다양한 크기(6x6, 8x8, 14x14)의 MNIST 숫자를 Tiny ImageNet 배경 위에 overlay하여 multi-scale 벤치마크 구성
• patch size p ∈ {4,8,16}과 digit size s ∈ {8,12,14}의 모든 조합에 대해 9개의 개별 MAE를 학습하여 patch size의 필터링 효과 검증
• FlexiViT 방식으로 32x32 해상도에서 초기화된 linear projection 가중치를 학습 스텝마다 목표 patch size로 동적 리사이즈, 위치 임베딩은 bi-linear interpolation으로 시퀀스 길이에 맞게 리사이즈
• 연속적인 patch size p_i를 MLP를 통해 encoder/decoder 전 layer의 LayerNorm scale(β)·shift(µ) 파라미터로 매핑하는 adaptive LayerNorm conditioning 도입 (p(z|x,p) 명시적 모델링)
• decoder에 conditioning을 적용하는 것이 encoder의 표현 품질을 저하시킬 수 있다는 가설을 세우고 이를 실험으로 검증
• Base MAE(고정 patch), FlexiMAE(조건화 없는 variable patch), Mosaic MAE(조건화된 variable patch) 세 계열을 clinical lateral DXA(LVA) 데이터에서 linear probing으로 비교

Originality

• MAE의 patch size를 명시적으로 "spatial/high-pass filter"로 해석하고 이를 통제된 실험(MNIST-on-ImageNet)으로 정량적으로 입증한 개념적 기여
• FlexiViT의 supervised dynamic patch resizing 아이디어를 self-supervised 생성 모델(MAE)로 확장
• diffusion model(DiT)의 time-step conditioning 방식을 patch-size conditioning으로 변형하여 MAE의 LayerNorm에 적용한 adaptive LayerNorm 기반 scale conditioning은 기존 FlexiMAE 대비 새로운 조합

Limitation & Further Study

• workshop paper 특성상 본문이 발췌되어 실험 섹션(4.2 이후)과 정량적 결과 표가 부족하며, decoder conditioning의 영향에 대한 가설 검증 결과가 명확히 제시되지 않음
• 벤치마크가 MNIST-on-ImageNet이라는 인공적 데이터셋과 단일 의료 영상 모달리티(DXA)에 국한되어 일반화 가능성이 제한적
• patch size 샘플링 범위나 conditioning MLP 구조에 대한 ablation, 계산 비용(학습 시간·메모리) 비교가 부족해 보임
• 후속 연구로 더 다양한 의료 영상 모달리티(CT, MRI, 병리 슬라이드 등)와 다운스트림 태스크로의 확장, decoder conditioning의 체계적 분석이 필요

Evaluation

총평: Patch size가 MAE 표현의 스케일을 결정한다는 직관적이지만 중요한 관찰을 명확한 실험으로 뒷받침하고, 이를 해결하는 실용적인 architecture(Mosaic MAE)를 제안한 점에서 의료 영상 표현 학습에 유의미한 기여를 하는 workshop 논문이다. 다만 발췌된 본문만으로는 정량적 결과와 ablation의 완전성을 판단하기 어렵다.