[논문 리뷰] Long-LRM (arXiv 2024) : 3DGS 학습 1.3초

Long-LRM: Long-sequence Large Reconstruction Model for Wide-coverage Gaussian Splats, Ziwen Chen, arXiv2024

 

V100에서 1.3초만에 3D Gaussian Splatting parameter를 추정하는 Generalizable 모델입니다. 입력 이미지 제약은 960x540 해상도 32개까지 사용한다는 점입니다. 기존 Generalizable 연구들은 입력 이미지 제약이 1~4개까지였던 것과 비교해보면 엄청난 개선입니다. 퀄리티도 최초 3D Gaussian Splatting모델과 비교했을 때 비슷하거나 더 좋습니다.

 

동작 원리를 살펴보겠습니다.

 

Overview

32개 이미지를 입력받아 Mamba2+Transformer 구성된 네트워크로 3D Gaussian Paramter를 출력하도록 설계됩니다.

DL3DV-10K 데이터셋에서 10,510개 scene으로 해당 네트워크를 학습합니다.

Rendering Loss를 기본으로 Depth, Opacity Regularization Loss를 적용합니다.

Token Merge와 Gaussian Pruning으로 사용 메모리를 줄였습니다.

 

 

기존 방법론 및 문제점

Generalizable 3D-GS 연구들은 여러Scene으로 학습해서 novel scene에 대해 빠르게 3D Reconstruction 하는 연구입니다. GS-LRM(Apr 2024), Instant3D(Nov 2023), MVSplat(Jul 2024) 연구가 있습니다. 장점은 빠르다는 점이지만 퀄리티가 좋지 않고 입력수가 제한되어 있습니다. 기존 연구 모두 backbone으로 Transformer가 사용되었는데, Transformer는 입력 token을 길게 가져 갈 수 없으므로, 저자는 이를 개선한다면 입력 이미지수를 늘일 수 있을 것이라고 생각했습니다.

 

Transformer의 Long-context reasoning을 다루는 연구 분야로 State-space models(SSM)가 있습니다. 이 중에 Mamba2 (2024) 라는 연구가 있으며, Transformer가 입력 길이에 대해 O(N^2)의 시간 복잡도를 가졌다면 Mamba2는 O(N)의 시간 복잡도를 가지며 성능은 상대적으로 좋지 않습니다. Mamba2는 one-direction으로 token을 scan하기 때문에, 이를 Vision 분야에 적용하기 위해 bi-direction으로 적용한 VMamba, LocalMamba 연구가 제안되었습니다. 하지만 느린 속도를 가지기 때문에 본 연구에서는 그대로 적용하지 않았다고 합니다. 

 

 

 

Network Design

Token화 : 960x540 해상도 이미지 32개를 8x8 patch로 나누어 260K token 으로 생성합니다.

Token Merge : 2x2 convolution으로 stride 2해서 token의 수를 1/4로 줄입니다.

Mamba2 Block : Mamba2를 one-direction이 아닌 bi-direction로 설계하여 7개의 block으로 구성합니다.

Transformer Block : Transformer를 1개 block으로 구성합니다.

분홍색 Box : 이를 3번 반복합니다. 총 24개 Block으로 구성됩니다. Token Merge는 9번째 Block앞에서만 수행됩니다.

Linear & Unpatchify : gaussian parameter를 출력합니다.

Gaussian Pruning : gaussian의 opacity가 0.001이하인 경우 제거합니다.

 

 

Training Detail

256x256 -> 512x512 -> 960x540 순으로 학습합니다. 차이를 보기 쉽게 서술식이 아닌 개조식으로 쓰겠습니다.

 

Stage1(256x256)

• 이미지 짧은쪽이 256되게 만들고 center cropping
• 연속한 64~128개를 랜덤으로 선택한 후, uniform하게 32개 이미지를 input으로 선택, 8개 이미지를 target으로 선택
• AdamW - Learning Rate 4E-4
• 256 batch size로 60K step 학습

Stage2(512x512)

• AdamW - Learning Rate 4E-5
• 64 batch size로 10K step 학습

Stage3(960x540)

• 원본이미지 사용
• 64~128개가 아닌 전체 이미지(200~300)에서 input/target 이미지를 선택
• 64batch, 10K step 학습
• Gaussian Pruning적용 : Gaussian의 opacity top 40%를 사용하고, 나머지 10%만을 랜덤하게 샘플링해서 사용 -> 메모리 절약
• FoV augmentation 적용: 이미지 0.77~1.0영역을 center crop
• (optional) 입력 이미지 16개로 finetuning

이미지 크기가 적을 때, corase한 특징을 보고, 이미지 크기가 커질수록 fine한 특징을 보게 하기 위해, learning rate를 줄이게 되고, step수와 batch size가 줄어 듭니다.

 

 

Loss

Rendering Loss

시각적인 pixel의 차이를 계산하는 Mean Squared Error(MSE) Loss부분과, LPIPS처럼 학습된 VGG에 이미지를 넣고 output vector의 차이를 계산하는 Perceptual Loss로 구성되어 있습니다. λ는 0.5입니다.

 

Depth Regularization for training stability

Large-scale로 학습을 하게 될 경우 불안정한 결과를 만들기 때문에, Gaussian Position을 optimization하는게 어려웠다고 합니다. rendering loss만으로 할 경우 3D 표면에 놓여있지 않은 floater들을 제거하기 위해, DepthAnything 모델로 획득한 disparity map(=1/depth)을 활용해서 regularization하는 방법을 사용했습니다. disparity의 median $t(d_i)$ 값을 기준으로  $\frac{1}{HW} \sum \left| d_i - t(d_i) \right|$ 식을 통해 normalization하여 depth를 smooth하게 만들어서 수렴이 잘되도록 설계했습니다.

 

Opacity Regularization for inference efficency

gaussian을 효율적으로 사용하기 위해서 추가한 term입니다. opacity는 0~1사이의 값으로, 1에 가까울수록 불투명해지고 0에 가까울수록 투명해집니다. 때문에 해당 수식을 통해 투명한 것들은 더욱 투명하게 만들게 됩니다.

 

Total Loss

전체 Loss입니다. 각각을 약 90%, 10%, 1% 의 weight를 갖게 설계되었습니다. 

 

 

Dataset

DL3DV (A Large-Scale Scene Dataset for Deep Learning-based 3D Vision, link)를 사용합니다. 10510개의 network training scene과 140개 benchmark 평가 scene으로 구성되어 있습니다. 1개 scene에서 200~300개 이미지 중 8번째마다 test데이터셋으로 사용하고, 나머지 이미지들에 대해, 카메라 중심을 K-means 하여 clustering 중심과 가까운 이미지를 사용하였습니다. inference결과에 대한 top 50% opacity에 대해 사용하고 나머지에 대해 제거하였습니다.

 

 

 

Experiments

feed-forwrad는 zero-shot training을 했는가에 대한 내용입니다. 기존 3DGS에 비해 학습속도를 600배 이상 단축시키고, 퀄리티도 유사한 수준임을 볼 수 있습니다. 

PSNR이 적은 모델에 대해서도 시각적으로 봤을 때, 높은 퀄리티를 보입니다.

 

 

Ablation Studies

하나의 표에 많은 실험을 넣어뒀습니다. 때문에 색으로 구분해서 박스를 그려봤습니다.

(빨간박스) Input views : 입력 이미지수가 많아지면 퀄리티가 향상됩니다.

(초록박스) Input Size : 입력 해상도가 높으면 퀄리티가 향상됩니다.

(파란박스) Transformer : 단순히 Transformer만으로 설계하면 메모리가 많이 사용됩니다. Mamba2만으로 설계하면 퀄리티가 좋지 않습니다. 논문에서 제안하는 7Mamba + 1 Transformer일 경우, 메모리가 적게 사용되면서 퀄리티가 높습니다. 

(보라박스) Token Merge : Token merge를 할 경우, 퀄리티는 감소하지만 메모리가 절약되어 속도가 빨라지고, 고해상도를 커버 할 수 있습니다. 

 

 

 

Regularization term에 대한 분석입니다.

regularization을 통해, floater들을 많이 제거하게 됩니다. 전체 Gaussian 중에 약30% Gaussian이 사라지게 되는 것을 볼 수 있는데, 이부분이 PSNR을 향상시키는 요소로 작용하는 것을 볼 수 있습니다.

sampling range에 따른 분석입니다. range를 어떻게 설정하느냐에 따라서 제거되는 Gaussian의 갯수가 차이가 발생한다는 것을 볼 수 있습니다.

 

 

Closing..

본 논문에 대해서 openreview 리뷰가 좋지 않네요. 전 개인적으로 입력 이미지수를 늘여서 일상 시나리오에서 일반 scene을 커버 할 수 있다는 점에서 굉장한 연구 결과라고 생각합니다. 이미지 각각에 대한 카메라 포즈를 추정하는 모델이 빠른 속도로 가능하다면, 서비스하기도 좋다고 생각됩니다. 공식 프로젝트 페이지(link)에 결과물을 보면 퀄리티가 아주 좋은 정도는 아니라서, 퀄리티가 아주 좋지 않아도 되는 서비스에 적합 할 것이라 생각됩니다.