Rethinking Patch Dependence for Masked Autoencoders (CrossMAE)

[Project Page]

[Github]

[arXiv](2024/01/25 version v1)

 

 

Abstract

Masked Autoencoder에서 디코더의 self-attention을 제거, 훨씬 빠르고 성능 좋음

 

 

 

CrossMAE

1. Preliminaries: Masked Autoencoders
2. Reconstruction with Cross-Attention
3. Partial Reconstruction
4. Inter-block Attention

 

Preliminaries: Masked Autoencoders

MAE

[논문공부] (자세한 리뷰) Masked Autoencoders are Scalable Vision Learners

 

 

Reconstruction with Cross-Attention

 

인코더의 최종 출력은 feature map의 가중합이며, 디코더의 causal self-attention을 생략한다.

 

Partial Reconstruction

Self-attention의 존재로 인해 디코딩 시 모든 masked token을 디코더 블록에 전달해야 했지만 cross-attention만 사용하면 각 토큰의 부분적인 재구성이 가능하다.

 

마스크 된 p개의 토큰 중 γ개의 토큰을 예측할 때, 모든 masked token을 예측하는 MAE보다 분산이 p/γ배 더 크다. 

따라서 학습률을 γ/p로 스케일링 해주면 모델 성능에 도움이 된다.

 

부분적인 재구성은 모든 토큰을 동시에 예측하는 MAE에 비해 계산 복잡성이 낮으며 표현의 품질도 저하되지 않았다.

 

Inter-block Attention

인코더의 최종 출력만 사용하는 것은 information bottleneck이다.

 

디코더 블록 k에서, 인코더 블록 {f_j} 간의 inter-block attention을 통해 얻은 가중합을 K, V로 사용한다.

즉, 디코더의 cross-attention 블록들은 각각 다른 encoder layer features의 가중합을 반영한다.

 

 

 

Experiments

 

Ablation