[논문리뷰] PointDiT: Pixel-Space Diffusion for Monocular Geometry Estimation

ICML 2026. [Paper] [Page]
Haofei Xu, Rundi Wu, Philipp Henzler, Nikolai Kalischek, Michael Oechsle, Fabian Manhardt, Marc Pollefeys, Andreas Geiger, Federico Tombari, Michael Niemeyer
Google | ETH Zurich | University of Tubingen | Microsoft | KE:SAI | Technical University of Munich
2 Jul 2026

Introduction

Monocular geometry estimation에 대한 기존 접근 방식은 크게 두 가지 카테고리로 나뉜다. 첫 번째는 deterministic한 regression 모델이다. 이러한 방법은 ViT와 convolution을 결합한 복잡한 하이브리드 아키텍처에 의존하며, 학습을 정규화하기 위해 복잡한 loss를 필요로 한다. 또한, task의 본질적인 모호성 때문에 모델은 출력 분포의 평균을 예측하는 경향이 있어, 특히 복잡한 장면 영역에서 고주파 디테일이 부족한 과도하게 부드러운 geometry를 생성하는 경우가 많다.

두 번째 카테고리는 latent diffusion model (LDM)을 사용하여 이러한 모호성을 해결하고자 한다. 이러한 방법들은 VAE를 통해 point map을 latent space로 압축해야 한다. Geometry 데이터의 무한한 범위와 상대적인 데이터 희소성으로 인해 VAE의 재구성 품질과 out-of-distribution 일반화 능력이 제한될 수 있다. 또한, 이러한 latent space를 구축하려면 정교한 tokenizer 설계가 필요한 경우가 많다. 결과적으로 인코딩 및 디코딩 과정에서 정보 손실이 빈번하게 발생하여 미세한 geometry를 정확하게 재구성하는 데 어려움을 겪는다. 또한, VAE 재구성 충실도와 diffusion 생성 능력 사이에는 최적화 목표가 상충되기 때문에 근본적인 trade-off가 존재한다.

본 논문에서는 그러한 아키텍처적 오버헤드와 복잡한 loss가 불필요함을 보여준다. 저자들은 JiT에서 영감을 받아, point map space에서 직접 학습하는 pixel-space diffusion 프레임워크를 제안하였다. 이를 통해 VAE 없이 diffusion의 확률적 특성을 활용하여 모호한 영역을 모델링할 수 있다. 본 아키텍처는 point map 패치에서 작동하는 단순한 ViT로 설계되었다. JiT의 핵심 요소는 flow matching에서 일반적으로 사용되는 $v$-prediction 대신, 깨끗한 이미지을 직접 예측하는 $x$-prediction이다. 본 논문은 JiT의 연구 결과를 이미지 생성 외의 다른 분야로 확장하여, 이 목표가 geometry 데이터에서도 매우 효과적이며 point map 추정 성능을 크게 향상시킨다는 것을 보여준다.

Geometry 예측을 가이드하기 위해, diffusion model은 입력 RGB 이미지에 기반하여 컨디셔닝된다. 모델은 단순한 linear patchification에서도 이미 우수한 성능을 보이지만, 강력한 prior를 도입하면 성능이 크게 향상된다. 저자들은 강력한 feature 추출 모델인 DINOv3를 채택했다. 일반 ViT backbone을 DINOv3 토큰으로 컨디셔닝함으로써 표현 학습에서 얻은 강력한 prior와 diffusion 학습을 결합할 수 있다.

PointDiT는 투명한 object와 같이 모호한 상황에서도 선명한 경계를 생성하고 깊이를 정확하게 구분하는 데 탁월한 성능을 발휘한다. 또한 1-step diffusion만으로도 매우 경쟁력 있는 결과를 달성하며, 샘플링 step을 추가할수록 구조적 디테일이 더욱 향상된다.

Method

본 논문에서는 단일 RGB 이미지 $\textbf{c} \in \mathbb{R}^{H \times W \times 3}$로부터 dense point map $\textbf{x} \in \mathbb{R}^{H \times W \times 3}$을 예측하는 문제를 다룬다. 이러한 단일 이미지 설정의 내재적인 모호성을 모델링하기 위해, ViT를 사용한 flow matching 프레임워크를 제안하였다. 입력 이미지를 조건으로 간단한 Gaussian noise 분포를 point map 분포로 전송하는 방법을 학습시킨다.

1. Point Map Generation with Flow Matching

Flow Matching

저자들은 단일 이미지로부터 point map 생성을 모델링하기 위해 flow matching 공식을 채택하였다. Flow matching은 noise 분포 $p_0$를 데이터 분포 $p_1$으로 지속적으로 변환하는 상미분 방정식(ODE)을 학습시킨다.

시간 $t \in [0, 1]$에서의 state를 \(\textbf{z}_t\)라고 하자. \(\textbf{z}_t\)는 noise 샘플 \(\boldsymbol{\epsilon} \sim p_0 = \mathcal{N}(\textbf{0}, \textbf{I})\)와 GT 데이터 샘플 $\textbf{x} \sim p_1$ 사이의 linear interpolation으로 정의된다.

\[\begin{equation} \textbf{z}_t = t \cdot \textbf{x} + (1 - t) \boldsymbol{\epsilon} \end{equation}\]

$t = 0$은 순수한 noise에 해당하고 $t = 1$은 깨끗한 데이터에 해당한다. 이 확률 경로를 생성하는 velocity field \(\textbf{v}_t (\textbf{z}_t)\)는 state의 시간 미분으로 주어진다.

\[\begin{equation} \textbf{v}_t = \frac{d \textbf{z}_t}{dt} = \textbf{x} - \boldsymbol{\epsilon} \end{equation}\]

이 선형 경로는 각 샘플 쌍 $(\textbf{x}, \boldsymbol{\epsilon})$에 대해 일정한 velocity를 유도하여 noise 분포와 데이터 분포 사이의 직선 이동을 보장한다.

Image Conditioned Flow Matching

본 논문에서는 조건부 분포 $p(\textbf{x} \vert \textbf{c})$를 모델링하기 위해 이 프레임워크를 확장하였다. 구체적으로, 조건부 velocity field \(\textbf{v}_\theta (\textbf{z}_t, t \vert \textbf{c})\)를 학습시킨다. $\textbf{c}$에 대한 컨디셔닝을 통해 모델은 이미지의 공간적 컨텍스트를 활용하여 geometry 모호성을 해결하고 flow를 타겟 point map으로 유도한다.

Point map 정규화

일반적인 RGB 이미지와 달리, dense point map은 장면 도메인(ex. 실내 vs. 실외)에 따라 좌표 범위가 달라진다. 데이터 $\textbf{x}$의 스케일이 noise의 스케일을 훨씬 초과하는 경우, 데이터 신호가 거의 0에 가까운 timestep에서도 interpolation 경로를 지배하게 된다. 이는 noise가 데이터 구조를 파괴하는 것을 막아 diffusion 학습을 불안정하게 만든다.

이러한 문제를 완화하기 위해, 학습 전에 point map을 정규화한다. 각 point map에 대해 중심점 $\boldsymbol{\mu}$와 중심점에서 각 점까지의 평균 유클리드 거리로 정의되는 스칼라 scale factor $s$를 계산한다. 정규화된 데이터 $\tilde{\textbf{x}}$는 다음과 같다.

\[\begin{equation} \tilde{\textbf{x}} = \frac{\textbf{x} - \boldsymbol{\mu}}{s} \end{equation}\]

이 정규화 과정을 통해 데이터는 noise와 유사한 스케일로 조정되어 안정적인 flow matching이 가능해진다. 모델은 이 정규화된 공간에서 학습되었으므로, 모델의 point map 예측 결과는 affine-invariant하며, 즉 모르는 scale 및 shift 값까지도 복원 가능하다.

하늘 처리

실외 장면에서 하늘의 깊이가 사실상 무한대에 가깝다는 점을 고려하여, $\boldsymbol{\mu}$와 $s$를 계산할 때 하늘 영역을 제외하고, 정규화된 프레임에서 하늘 영역을 고정된 반지름 3의 가상 구에 projection한다. 이 반지름은 실제 깊이를 나타내는 가상적인 값이므로, 하늘 픽셀을 완전히 마스킹하는 대신 학습 loss에서 가중치를 낮춘다.

작지만 0이 아닌 가중치를 유지함으로써 모델은 하늘을 멀리 떨어진 배경으로 간주하며, supervision이 없는 상황에서도 하늘의 geometry가 임의적으로 변하지 않도록 한다. 동시에 하늘의 가상 깊이 값이 최적화 과정에서 지배적인 영향을 미치지 않도록 한다. 결과적으로, 다양한 실내 및 실외 데이터셋에서 안정적인 공동 학습이 가능하다.

Inference 시에는 예측된 3D point 중 norm이 2.9를 초과하는 point는 제거한다.

2. Architecture

본 논문에서는 \(F_\theta\)를 pixel-space DiT 역할을 하는 ViT로 구현한다. 네트워크는 noise가 포함된 point map \(\textbf{z}_t\), 현재 timestep $t$, 조건 이미지 $\textbf{c}$를 입력으로 받는다. 기존의 flow matching 모델들이 일반적으로 velocity를 예측하는 것과는 달리, JiT에서 영감을 받아 noise가 제거된 point map을 예측하도록 학습되었다.

Point map patchification

\(\textbf{z}_t \in \mathbb{R}^{H \times W \times 3}\)은 입력 이미지와 동일한 공간 해상도를 가지므로 픽셀 수준에서 직접 ViT를 적용하는 것은 비용이 너무 많이 든다. 따라서 ViT 시퀀스 길이를 픽셀에서 패치로 줄이기 위해 point map을 겹치지 않는 $p \times p$ 패치의 그리드로 분할하여 patchify한다. 이렇게 하면 $N = (H/p) \times (W/p)$개의 패치가 생성되고, 각 패치는 $3p^2$ 크기의 벡터로 flatten된다. 그런 다음 이러한 벡터는 학습 가능한 linear projection $\phi$를 통해 임베딩 차원 $D$로 매핑되어 point map 토큰 \(\textbf{T}_z\)를 얻는다.

\[\begin{equation} \textbf{T}_z = \phi (\textrm{Patchify}(\textbf{z}_t)) \in \mathbb{R}^{N \times D} \end{equation}\]

이미지 컨디셔닝

조건 이미지 $\textbf{c}$는 생성에 대한 구조적 guidance를 제공한다. \(\textbf{z}_t\)와 달리 $\textbf{c}$는 깨끗하기 때문에 강력한 사전 학습된 표현을 활용하여 인코딩할 수 있다. 구체적으로, 고정된 DINOv3 인코더를 사용하여 패치 토큰을 추출하며, 여러 layer의 DINOv3 feature를 결합하는 것이 효과적이다. 특히, DPT head의 layer 선택 방식을 따라 균일하게 배치된 4개의 중간 layer를 사용한다.

이러한 feature를 융합하기 위해 정교한 convolution에 의존하는 DPT와 달리, 단순히 채널 차원을 따라 토큰을 concat하여 저수준 디테일에서 고수준 추상화에 이르는 풍부한 feature 계층 구조를 포착한다. 이를 통해 이미지 표현 \(\textbf{T}_c \in \mathbb{R}^{N \times 4D}\)가 생성된다. 공간적 정렬을 보장하기 위해 point map과 DINOv3 branch 모두에 대해 동일한 패치 크기 $p = 16$과 임베딩 차원 $D$를 사용한다.

이미지와 point map 융합

\(\textbf{T}_c\)와 \(\textbf{T}_z\) 사이의 공간적 정렬을 고려하여, channel-wise concat을 통해 두 모달리티를 융합하여 DiT에 대한 입력 \(\textbf{T}_\textrm{in} \in \mathbb{R}^{N \times 5D}\)를 생성한다. \(\textbf{T}_\textrm{in}\)을 $5D$ 차원에서 임베딩 차원 $D$로 projection한다. 그런 다음, 시퀀스는 multi-head self-attention과 MLP로 구성된 Transformer block들을 통해 처리된다.

깨끗한 point map 예측

DiT는 정제된 토큰 시퀀스 \(\textbf{T}_\textrm{out} \in \mathbb{R}^{N \times D}\)를 출력한다. Dense point map을 복원하기 위해, 각 토큰에 linear projection head를 적용하여 차원을 $D$에서 $3p^2$로 다시 projection한다. 그런 다음 unpatchification 연산을 통해 원래의 공간 그리드 $(H/p) \times (W/p) \times p \times p \times 3$으로 재배열하고 최종적으로 원본 해상도의 point map $\hat{\textbf{x}} \in \mathbb{R}^{H \times W \times 3}$을 생성한다.

3. Training

Noise schedule

학습 중 timestep $t \in [0, 1]$를 샘플링하기 위해 JiT를 따라 logit-normal distribution을 사용한다. 구체적으로, $\mu = −0.8$, $\sigma = 0.8$인 $z \sim \mathcal{N}(\mu, \sigma^2)$를 추출하고 sigmoid function을 $[0, 1]$로 매핑한다.

Sigmoid function은 경계에 점근적으로만 접근하기 때문에 모델은 정확한 순수 noise $t=0$에서 학습되지 않는다. inference는 항상 $t = 0$에서 시작되므로 학습-테스트 불일치가 발생하며, 모델은 flow 궤적을 시작하는 데 어려움을 겪을 수 있다. 이 문제를 해결하기 위해 rectified sampling 전략을 채택했다. 0.1의 확률로 logit-normal 샘플링을 무시하고 $t = 0$을 명시적으로 설정한다. 이를 통해 모델은 inference 시작 시 접하는 순수 noise 분포에 맞춰 보정된다.

Velocity loss

네트워크 \(F_\theta\)는 깨끗한 point map \(\hat{\textbf{x}}\)를 예측하지만, JiT를 따라 velocity space에서 최적화한다. 학습 과정에서 \(\textbf{z}_t\)를 구성하고 네트워크 예측 \(\hat{\textbf{x}} = F_\theta (\textbf{z}_t, t, \textbf{c})\)를 얻는다. 그런 다음 interpolation 경로를 재배열하여 \(\hat{\textbf{x}}\)를 velocity \(\hat{v}_t\)로 변환한다.

\[\begin{equation} \hat{\textbf{v}}_t (\textbf{z}_t, t) = \frac{\hat{\textbf{x}} - \textbf{z}_t}{1 - t} \end{equation}\]

수치적 안정성을 위해 분모의 $(1-t)$는 0.05로 clipping된다.

Flow matching loss는 추정된 velocity \(\hat{\textbf{v}}_t\)와 GT velocity \(\textbf{x} - \boldsymbol{\epsilon}\) 사이의 MSE를 최소화한다.

\[\begin{equation} \mathcal{L}_\textrm{fm} = \mathbb{E}_{\textbf{x}, t, \boldsymbol{\epsilon}} \left[ \frac{1}{M} \sum_{i=1}^M w_i \| \hat{\textbf{v}}_{t,i} - (\textbf{x}_i - \boldsymbol{\epsilon}_i) \|_2^2 \right] \end{equation}\]

($w_i$는 픽셀별 하늘 가중치, 하늘 픽셀의 경우 $w_i = 0.01$, 그렇지 않은 경우 $wi = 1$).

Relative Point Loss

Flow matching loss만으로도 이미 효과적인 모델을 만들 수 있다. 하지만 모델은 VAE 없이 원본 데이터 공간에서 직접 깨끗한 출력을 예측하기 때문에 필요에 따라 이 출력에 보조 loss를 적용하는 것이 간단하다. 저자들은 이러한 유연성을 보여주기 위해 relative point loss를 추가했다. Point map은 높은 동적 범위를 가지며, 멀리 있는 점은 norm이 커서 표준 error metric에 큰 영향을 미치고, 가까운 디테일은 상대적으로 덜 중요하게 여겨진다. 따라서 각 픽셀별 오차를 타겟 점의 크기로 정규화하여 로컬 디테일의 재구성을 강조한다.

\[\begin{equation} \mathcal{L}_\textrm{rel} = \mathbb{E}_{\textbf{x}, t, \boldsymbol{\epsilon}} \left[ \frac{1}{M} \sum_{i=1}^M w_i \frac{\| \hat{\textbf{x}}_i - \textbf{x}_i \|_1}{\| \textbf{x}_i \|_2 + \xi} \right] \end{equation}\]

($w_i$는 픽셀별 하늘 가중치, $\xi$는 작은 상수)

Total loss

최종 loss는 두 loss의 가중 합이다.

\[\begin{equation} \mathcal{L} = \mathcal{L}_\textrm{fm} + \lambda \mathcal{L}_\textrm{rel} \end{equation}\]

($\lambda = 0.1$)

전체 모델은 end-to-end로 학습된다. MoGe와 같이 다양한 정규화 loss들에 의존하는 기존 방법들과 달리, PointDiT는 flow matching loss와 간단한 보조 loss 1개만 사용한다.

4. Inference

Inference 시에는 입력 이미지 $\textbf{c}$를 조건으로 순수 noise \(\textbf{z}_0 \sim \mathcal{N}(\textbf{0}, \textbf{I})\)로부터 $\textbf{x}$를 복원하기 위해, $t=0$부터 $t=1$까지 ODE \(d \textbf{z}_t = \textbf{v}_\theta (\textbf{z}_t, t \mid \textbf{c}) dt\)를 푼다. 표준 Euler solver를 사용하며, step 크기는 $\Delta t$이다. 각 step $t$에서, 깨끗한 데이터 \(\hat{\textbf{x}}\)를 예측하고, velocity \(\hat{\textbf{v}}_t\)를 도출하고, state를 업데이트한다.

\[\begin{equation} \textbf{z}_{t + \Delta t} \leftarrow \textbf{z}_t + \Delta t \cdot \hat{\textbf{v}}_t \end{equation}\]

이 반복적인 과정을 통해 샘플은 학습된 선형 궤적을 따라 이동하여 최종 point map을 재구성한다.

놀랍게도, 모델은 1-step inference만으로도 경쟁력 있는 결과를 얻을 수 있으며, 추가 step을 거치면 동일한 모델을 사용하여 품질이 더욱 향상된다. 이는 예측된 point map과 조건 이미지 간의 픽셀 단위 정렬 덕분이다. 각 출력 위치는 해당 이미지 feature에 의해 크게 결정되므로, noise에서 목표 geometry로의 변환이 거의 직접적인 매핑이 되어 1-step에서도 정확한 결과를 얻을 수 있으며, 추가 step은 주로 디테일을 개선하는 데 사용된다. 또한, 모델을 랜덤 noise 대신 모든 값을 0으로 초기화함으로써 inference 시점에 deterministic하게 예측할 수도 있다.

Experiments

• 데이터셋
  • 256$\times$256 사전 학습: SceneNet-RGBD
  • 512$\times$512 fine-tuning: Hypersim, VKITTI2, UrbanSyn, Synscapes, TartanAir, OmniWorldGame, EDEN, IRS, Dynamic Replica, MVSSynth, TartanGround

1. Evaluation Results

다음은 다른 모델들과 point map, depth map, 경계 품질(BF1)을 비교한 결과이다.

다음은 다른 모델들과 point map을 비교한 예시들이다.

다음은 1-step inference 결과이다.

다음은 샘플링 step에 따른 depth map을 비교한 결과이다.

2. Ablation and Analysis

다음은 deterministic regression과의 비교 결과이다.

다음은 ablation 결과이다.

다음은 패치 크기에 따른 비교 결과이다.