[논문리뷰] Gaussian Masked Autoencoders

arXiv 2025. [Paper] [Page]
Jathushan Rajasegaran, Xinlei Chen, Rulilong Li, Christoph Feichtenhofer, Jitendra Malik, Shiry Ginosar
Meta | UC Berkeley | Toyota Technological Institute at Chicago
6 Jan 2025

Introduction

정적 이미지에서도 레이어 표현을 통해 세계의 구조에 대해 더 많이 알 수 있다. 본 논문은 2.1D 레이어링을 통해 objectness, grouping, semantic 구조와 같은 높은 수준의 semantic 추상화를 self-supervised learning을 통해 공동 학습하는 것을 제안하였다. 아이디어는 개념적으로 간단하다. 픽셀 기반 self-supervised 표현 학습 방식인 Masked Autoencoder (MAE)를 기반으로, 학습된 latent 표현이 바람직한 중간 표현으로 이어질 수 있도록 한다. 핵심 통찰력은 3D Gaussian이 semantic 및 공간 이해로 이어질 수 있는 중간 이미지 표현에 적합한 후보라는 것이다.

3D Gaussian은 원래 최적화 기반 3D 재구성을 위해 제안되었다. 정사각형 픽셀과 같은 기하학적으로 균일한 표현과 달리, 이미지에 대한 크기, 위치 및 정보 분포는 동적으로 학습된다. 게다가 Gaussian 기반 표현은 픽셀 공간으로 다시 매핑하는 splatting 기반 이미지 렌더링 덕분에 end-to-end 학습에 적합하다. 따라서 MAE와 같은 self-supervised 프레임워크 내에서 이러한 중간 수준 표현을 공동으로 학습할 수 있다.

이 접근 방식을 Gaussian Masked Autoencoder (GMAE)라고 부르며, 3D 재구성 프레임워크가 아닌 시각적 표현 학습 프레임워크에서 Gaussian primitive를 탐색한 최초의 연구이다. GMAE는 표준 MAE 학습에 비해 무시할 수 있는 수준의 오버헤드만 추가하며, splatting은 계산 시간을 1.5%만 증가시킨다. 표현 학습 성능을 저하시키지 않으면서 GMAE는 zero-shot 능력에서 상당한 이점을 얻는다.

Gaussian 기반 이미지 표현은 여러 가지 내장된 장점이 있다. 표현의 불균일성은 이미지의 정보 밀도와 상관관계가 있는 표현 밀도의 공간적 분포로 이어진다. 3D Gaussian이 $z$축을 따라 이동하도록 허용함으로써, 모델은 한 이미지의 하나의 시점뿐만 아니라 수백만 개의 뷰를 관찰하여 세계의 기본 구조를 학습한다. 결과적으로, 튜닝 없이도 전경/배경 분할, 단순한 레이어링, edge detection를 깊이 불연속성으로 찾을 수 있다.

GMAE로 학습한 표현은 image classification과 object detection task에서 MAE와 유사한 성능을 보인다. 표현 품질은 사용된 Gaussian의 수에 따라 향상된다. 이러한 결과는 GMAE가 MAE를 증강하고 중간 수준 표현을 사용하여 이점을 얻을 수 있는 애플리케이션에서 더 나은 대안이 될 수 있음을 시사한다.

Method

본 논문은 self-supervised learning의 주류인 픽셀 기반 학습과 표현에 추가 속성을 부과할 수 있는 latent 기반 학습을 제안하였다. 핵심 통찰력은 end-to-end로 학습 가능한 3D Gaussian이 불균일한 속성으로 인해 중간 수준 이미지 표현에 적합한 후보라는 것이다.

많은 이미지 컬렉션이 주어지면 마스킹된 입력에서 전체 이미지를 재구성하기 위해 MAE를 학습시킨다. MAE 인코더는 마스킹된 이미지의 마스킹되지 않은 정사각형 패치를 학습된 임베딩으로 인코딩하는 방법을 학습하는 ViT이다. 그러나 MAE에서처럼 픽셀 패치를 직접 예측하는 대신, ViT 기반 디코더는 3D Gaussian의 색상, 3D 중심 위치, 크기, 방향을 예측한다. 그런 다음 이러한 Gaussian들을 splatting 기반의 미분 가능한 렌더러를 사용하여 이미지로 렌더링하고 픽셀 공간에서 MSE loss를 사용하여 전체 모델을 학습시킨다.

모델은 ViT 기반 인코더 모델, 가벼운 디코더 모델, 미분 가능한 렌더러를 가지고 있다. 주어진 이미지에 대해 먼저 $N$개의 패치로 patchify하고 마스킹 비율 $r$로 무작위로 마스킹하여 $n$개의 패치만 보이게 한다. ViT 인코더 모델은 마스킹되지 않은 패치만 보고 패치에서 latent 임베딩 $x_i \in \mathbb{R}^{d_\textrm{enc}}, i \in \{1, \ldots, n\}$으로 인코딩한다.

디코더에 $k$개의 학습 가능한 쿼리 토큰 $q_j \in \mathbb{R}^{d_\textrm{dec}}, j \in \{0, 1, \ldots, k\}$가 있다고 하자. $k$는 마스킹된 토큰의 수와 관계없이 어떤 값이든 될 수 있다. 인코더 latent를 $\hat{x}_i \in \mathbb{R}^{d_\textrm{dec}}$로 projection하고 쿼리 토큰과 concat한다.

\[\begin{equation} X_\textrm{dec} = \{\hat{x}_i\}_{i=1}^n \cup \{q_j\}_{j=1}^k \end{equation}\]

디코더는 $X_\textrm{dec}$ 토큰을 보고 각 쿼리 토큰에 대해 하나씩 총 $k$ 개의 Gaussian을 예측한다. 각 Gaussian은 14차원 벡터 $g_j = \{p, s, \phi, r, o\} \in \mathbb{R}^{14}$로 parameterize된다.

$k$개의 Gaussian을 얻은 후, 고정된 카메라 projection으로 평면에 splatting시킨 후 렌더링하여 이미지를 생성한다. Gaussian의 scale은 $c \cdot \textrm{sigmoid}(s)$로 제한하며, $c$는 Gaussian의 최대 크기이다. 렌더링 후, MSE loss를 사용하여 재구성된 이미지를 원래 입력 이미지와 비교한다.