[논문리뷰] Trellis: Structured 3D Latents for Scalable and Versatile 3D Generation

Paper: https://arxiv.org/abs/2412.01506 Project Page: https://trellis3d.github.io/ Github: https://github.com/microsoft/TRELLIS?tab=readme-ov-file

1. Introduction 3. Method 3.1. Structured Latent Representation 3.2. Structured Latents Encoding and Decoding 3.2.1. Visual feature aggregation 3.2.2. Sparse VAE for structured latents 3.2.3. Decoding into versatile formats 3.3. Structured Latents Generation 3.3.1. Rectified flow models 3.3.2. Sparse structure generation 3.3.3. Structured latents generation 3.4. 3D Editing with Structured Latents 3.4.1. Detail variation 3.4.2. Region-specific editing 4. Experiment Implementation Details 4.3. Ablation Study 4.3.1. Size of Structured Latents 4.3.2. Rectified Flow v.s. Diffusion 4.3.3. 모델 크기 (Model Size)5. Limitation

1. Introduction

3D를 위한 AIGC(AI Generated Content)는 많은 발전을 이루었으나, 아직 2D Image 만큼 실제 ready-to-use tool로 사용되지는 못하고 있다. 일반적으로 pixel grid로 표현되는 2D 이미지와는 다르게 3D의 경우는 mesh, point cloud, radiance fields, 3d gaussian 등으로 그 representation이 다양하기 때문이다. Mesh나 Implicit-fields를 representation으로 활용하는 3D 생성 연구의 경우 디테일한 표현에 약하고, 3DGS나 NeRF를 활용하는 연구들은 그럴듯한 geometry 추출에 어려움을 겪고 있다. 게다가, 각 representation의 고유한 구조적 특성 때문에 단일 아키텍쳐로 처리하기가 어렵다. 이는 2D 생성 모델들이 통합된 latent space에서 학습하며 표준화된 방식(Latent Diffusion Model, Flow)을 구축한 최근 사례와 대조적이다. 즉, 본 논문에서는 3D representation에 대한 근본적인 문제를 제기한다.

본 논문에서는 솔루션으로서 SLAT(Structured Latent Representation, 구조적 잠재 표현)을 제안하는데, SLAT은 희소 구조와 강력한 시각적 특징을 결합하고 다양한 3D 표현 형식(예: radiance fields, 3D gaussian, mesh)으로 디코딩할 수 있는 프레임워크를 제공한다. 먼저 Object 표면과 교차하는 활성 Voxel(3D 객체의 겉면에 해당하는 복셀)에 대한 local latent를 정의한다. 다음으로 3D 에셋을 dense하게 렌더링한 multi-view 이미지들에 대해 feature(디테일한 geometry 및 visual)를 추출하여(이때 DINO-v2가 활용됨) 정의한 local latent와 결합한다. 즉, SLAT은 활성 Voxel을 통해 필수적으로 고려되어야 할 표면 Voxel만을 1차 추리고, Multi-View 이미지에서 디테일한 구조적 특징과 비주얼 특징(텍스쳐, 색상 등)을 추출하여 활성 Voxel에 결합한 뒤, 이를 Local Latent로 만든 것이다. 최종적으로 디코더들이 SLAT을 다양한 고품질 3D 표현으로 매핑할 수 있게 된다.

또한, 본 논문에서는 SLAT(Structured Latent Representation)을 기반으로 텍스트 프롬프트나 이미지를 조건으로 활용하는 대규모 3D 생성 모델 TRELLIS를 제안한다. TRELLIS는 2단계 파이프라인으로 동작하며, 첫 번째 단계에서는 SLAT의 희소 구조(sparse structure)를 생성하고, 두 번째 단계에서는 비어 있지 않은 셀(non-empty cells)에 대한 잠재 벡터(latent vector)를 생성한다. Rectified flow transformer를 백본 모델로 사용하며, 약 50만 개의 신중히 수집된 3D 자산 데이터셋을 기반으로 최대 20억 개의 매개변수를 학습하였다.

3. Method

3.1. Structured Latent Representation

3D 에셋 에 대해, geometry와 appearance 정보를 unified structured latent 로 인코딩하는데, 이는 3D grid상의 local latent 집합으로 정의된다.

는 의 표면(surface)과 교차하는 3D grid 상 활성 복셀의 위치 index이고, 는 상응하는 복셀에 연결된 local latent를 뜻한다. N은 3D grid의 공간적 길이이며, L은 활성 복셀의 총 개수이다. 직관적으로, 활성 복셀 는 3D 에셋의 대략적인 구조를 나타내고, latent 는 외형 및 형태의 세부 사항을 나타낸다. 이러한 structured latent들은 3D 에셋 의 전체 표면을 포괄하며, 전체적인 형상과 세부 사항을 효과적으로 포착한다.

3D data의 sparsity(희소성)으로 인해, 활성 복셀의 개수는 3D grid의 총 사이즈보다 매우 작고(), 이는 상대적으로 높은 해상도로 생성될 수 있음을 의미한다. 기본적으로 N=64로 설정하며, 이 경우 L의 평균값은 20K이다.

3.2. Structured Latents Encoding and Decoding

SLAT을 위한 encoder 및 여러 타입(RF, 3DGS, Mesh)을 위한 decoder를 소개한다.

3.2.1. Visual feature aggregation

먼저 3D 에셋 를 복셀화된 feature 로 변환한다. 여기서 는 eq 1.에서 정의된 활성 복셀이고, 는 local region의 상세한 구조와 외형 정보를 담은 visual feature다.

각 활성 복셀의 를 유도하기 위해, 저자들은 의 dense한 multiview 이미지들에서 feature를 추출하여 합친다. Sphere에서 랜덤하게 샘플된 카메라 views에서 이미지들을 렌더링하고, DINOv2 encoder를 통해 feature map을 추출한다. 각 복셀은 multiview feature 맵의 상응하는 위치로 투영(project)되어 feature를 가져오며, 그 평균값을 로 사용한다(Fig.2의 좌측 상단). f는 structured latents 의 해상도()와 일치하도록 설정된다. 경험적으로, 이는 원본 3D 에셋을 높은 정확도로 복원하는데에 충분한데, DINOv2 특징의 강력한 표현 능력과 활성 복셀이 제공하는 대략적인 구조 덕분이다.

3.2.2. Sparse VAE for structured latents

복셀화된 feature 와 함께 3D 에셋 인코딩을 위한 transformer기반의 VAE 아키텍쳐를 소개한다.

구체적으로 인코더 는 먼저 를 structured latent 로 인코딩하고, 디코더 는 를 특정 3D representation으로 변환한다. 인코더와 디코더는 Reconstruction loss(디코딩된 3D 에셋과 원본 3D에셋에 대한)를 통해 학습되며, 이때 에 대해 정규 분포를 따르도록 KL-패널티를 적용한다.

인코더와 디코더는 동일한 Transformer 구조를 공유하며, Fig. 3a에 그 구조가 나타나 있다. sparse한 복셀을 처리하기 위해, 활성 복셀로부터 입력 특징을 serialize(직렬화)하고, 복셀 위치 기반 sinusoidal positional encodings(Transformer에 위치정보를 제공하기 위함, sin 함수 등 주기 함수를 활용한 encoding)을 추가하여 가변 길이 L의 토큰으로 만든다. 이는 이후 Transformer 블록에서 처리된다. Latents의 로컬 특성을 고려하기 위해, 3D 공간에서 Shifted Window Attention을 도입하여 로컬 정보 상호작용을 강화한다. 이는 전체 어텐션(full attention) 구현보다 효율성을 개선한다.

3.2.3. Decoding into versatile formats

SLAT(Structured Latents)은 다양한 3D Representation(3D Gaussians, Radiance Fields, Meshes)으로 디코딩을 지원하며, 각각의 디코더는 , , 으로 명명된다. 이러한 디코더는 출력 레이어를 제외하고 동일한 아키텍처를 공유하며, 각 표현 형식에 특화된 재구성 loss로 학습된다.

(a) 3D Gaussians

각 는 K개의 Gaussian으로 디코딩 된다. 위치 오프셋 , 색상 , 스케일 , 불투명도 , 회전 을 포함한다. Gaussian의 최종 위치 는 의 locality를 보존하기 위해 활성 복셀 근처로 제한된다 (). 재구성 loss는 렌더링된 가우시안과 GT 이미지간의 , , 이다.

(b) Radiance Fields

여기서 및 는 해상도의 local radiance volume을 나타내는 CP Decomposition이다. 재구성 Loss는 Gaussian과 유사하다.

(c) Meshes

는 Flexi-Cubes의 매개변수이며, 은 복셀의 8개 꼭짓점에 대한 signed distance 값이다. 디코더 뒤에는 2개의 컨볼루션 업샘플링 블록이 추가되어 최종 출력 해상도를 으로 증가시킨다. Mesh는 0-레벨 Iso Surface에서 추출되며, 렌더링된 depth (normal) 맵과 GT로부터 을 계산하여 재구성 loss로 활용한다. (실제로 코드 상에서는 Mesh Decoder 쪽에서 Texture도 처리하는 것으로 보이며 아래에 A.2.의 디테일을 추가해둔다.)

Mesh Decoder Detail

우리는 희소 구조의 공간 크기를 에서 으로 증가시키기 위해 Transformer 백본 뒤에 두 개의 희소 컨볼루션 업샘플러(sparse convolutional upsamplers)를 추가합니다.메쉬를 추출하기 위해, 우리는 주로 형상(geometry)에 초점을 맞추지만, 메쉬에 색상(color)과 법선 맵(normal maps)을 추가로 예측합니다. 따라서, 고해상도 활성 보셀의 최종 출력은 다음과 같습니다

메쉬 추출 과정

희소 구조를 메쉬 표면 추출이 가능한 밀집 그리드(dense grid)에 연결(attach)하여, FlexiCubes 기반으로 미분 가능한 표면을 추출합니다.

비활성 복셀에 대해서는 그리드 상에서 부호 거리 값을 1.0으로 설정하고, 다른 관련 속성은 0으로 설정합니다.

0-레벨 아이소표면(isosurface)을 기반으로 메쉬를 추출합니다.

추출된 메쉬의 각 꼭짓점은 그리드 꼭짓점의 속성(예: c,n)을 보간하여 계산됩니다.

우리는 Nvdiffrast를 사용하여 추출된 메쉬와 그 속성을 렌더링합니다. 이 과정에서 다음과 같은 출력물이 생성됩니다:

Foreground mask (M): 전경 마스크.

Depth map (D): 깊이 맵.

Normal map (N): 메쉬로부터 직접 파생된 법선 맵.

RGB image (C): 메쉬에서 추출된 색상 이미지.

Normal map (N): 예측된 법선 값으로부터 생성된 법선 맵.

학습 Objective는 아래와 같습니다.

실제로, 논문에서는 Gaussian을 사용하여 인코더와 디코더를 엔드투엔드로 학습한다. 다른 출력 형식에 대해서는 학습된 인코더를 고정하고 디코더만 새로 학습한다. Gaussian으로 학습된 SLAT은 다른 형식에서도 원본을 충실히 복원할 수 있어 확장성을 입증한다(표 1 참조). 더 자세한 구현 세부 사항은 Sec. A.2에 나와 있다.

3.3. Structured Latents Generation

SLAT을 생성하기 위해 두 단계로 이루어진 생성 파이프라인을 도입한다. 첫 번째 단계에서는 희소 구조(sparse structure)를 생성하고, 두 번째 단계에서는 해당 구조에 연결된 로컬 잠재 요소(local latents)를 생성한다. Latent를 모델링하기 위해 Rectified Flow 모델을 사용한다.

3.3.1. Rectified flow models

Rectified Flow 모델은 linear interpolation 기반의 forward process를 사용한다. , 여기서 는 time step, 는 데이터 샘플, 은 노이즈이다. 이 과정은 데이터 샘플과 노이즈 사이를 보간한다. backward process는 시간 종속적인 벡터 필드 로 표현되며, noisy한 샘플을 data distribution으로 옮긴다. 자세한 내용은 Flow 관련 논문을 참고하는 편이 좋을 것 같다.

3.3.2. Sparse structure generation

첫번째 단계에서는 sparse structure 를 생성한다. 이것을 neural net으로 가능하게 하기 위해, sparse한 활성 복셀을 dense한 바이너리 3D grid 으로 변환하는데, 활성이면 1 아니면 0으로 복셀 값을 세팅한다.

다이렉트로 dense grid 를 생성하는 것은 매우 비효율적이기 때문에, 3D 컨볼루션 블록이 포함된 단순한 VAE를 사용하여 를 저해상도 feature grid 로 압축한다. 는 주로 coarse한 형상을 나타내기 때문에, 이러한 압축은 손실이 거의 없으며 효율성을 크게 향상시킨다. 또한, 의 이산적인 값을 연속적인 특징으로 변환하여 Rectified Flow 학습에 적합하게 만든다.

다음으로 Transformer 백본 를 사용하여 를 생성한다.(Fig. 3b 참조). 입력으로 noisy grid를 직렬화하고, 위치 인코딩을 추가한 후 Transformer에 전달한다. Timestep 정보는 Adaptive Layer Normalization (adaLN)과 gating 메커니즘을 통해 통합된다. 텍스트 조건은 CLIP 모델을 사용하고, 이미지 조건은 DINOv2을 사용한다. 최종적으로 denoised된 feature grid 는 discrete grid 로 디코딩되고, 다시 활성 복셀 로 변환된다.

3.3.3. Structured latents generation

두 번째 단계에서는 Transformer 을 사용하여, 에 기반하여 latent 를 생성한다(Fig. 3c 참조). Sparse VAE 인코더(3.2.2)와 달리, noisy 잠재 요소를 직접 직렬화하지 않고, 더 짧은 시퀀스로 패킹하여 효율성을 높인다(예: DiT [67]와 유사한 방식). 이전 단계에서 생성한 sparse structure(3.3.2)를 활용하여, local 영역 내에서 latent를 패킹하고 직렬화 전에 다운샘플링 블록을 적용한다. 이후 Transformer 블록을 통해 시간 조정된(time-modulated) 처리를 수행한다. 마지막에 업샘플링 블록과 skip connections을 추가하여 공간 정보를 효과적으로 전달한다.

와 은 독립적으로 학습되며, 학습 후 SLAT 은 두 모델을 통해 순차적으로 생성된다. 또한 다양한 디코더(, , )를 통해 고품질 3D 에셋으로 변환된다.

3.4. 3D Editing with Structured Latents

유연한 3D 편집도 지원하며, 두 가지 간단한 튜닝 없는(tuning-free) 편집 전략을 제안한다.

3.4.1. Detail variation

structure와 latent가 분리되어 있기 때문에, 3D 에셋의 전체적인 대략적인 형상을 변경하지 않으면서 세부 사항을 조정할 수 있다. 이는 3D 에셋의 구조를 유지한 상태에서, 다른 텍스트 프롬프트를 사용하여 두 번째 생성 단계를 실행함으로써 쉽게 달성할 수 있다. 즉 Geometry가 고정된 상태에서의 텍스쳐 생성 용도로도 활용이 가능하다.

3.4.2. Region-specific editing

SLAT의 지역성(Locality) 특성 덕분에, 특정 영역의 복셀과 latent를 수정하고 다른 부분은 변경하지 않은 상태로 남겨둘 수 있다. 이를 위해, 우리는 Repaint [55]를 두 단계 생성 파이프라인에 맞게 조정했다. 편집할 복셀의 경계 상자(bounding box)가 주어지면, 해당 영역에서 새로운 콘텐츠를 생성하도록 Flow 모델의 샘플링 과정을 수정한다. 이 과정은 변경되지 않은 영역과 주어진 텍스트 또는 이미지 프롬프트를 조건으로 설정하여 이루어진다.

4. Experiment

Implementation Details

약 500K의 고품질 3D 에셋으로 구성된 데이터셋을 사용하여 Trellis 모델을 학습했다. 이 데이터셋은 Objaverse (XL) [16], ABO [13], 3D-FUTURE [20], HSSD [34]의 4개 공개 데이터셋에서 수집되었다. 각 에셋 대해 150개의 이미지를 렌더링하고 GPT-4o [1]를 사용하여 캡션을 생성했다. 텍스트와 이미지 프롬프트 모두에 대해 데이터 증강을 수행했으며, 텍스트는 다양한 길이로 요약되었고 이미지는 서로 다른 시야각(Field of View, FoV)으로 렌더링되었다.

우리는 AdamW [53] 옵티마이저를 사용했으며, 배치 크기는 256. 훈련은 총 400K Step 동안 진행되었으며, 64개의 A100 GPU(40GB 메모리)를 사용. 세 가지 모델 크기를 학습: Basic(342M 파라미터), Large(1.1B 파라미터), X-Large(2B 파라미터). 추론 과정에서는 샘플링 단계 수 50으로 설정하고 CFG 강도를 3으로 설정.

4.3. Ablation Study

4.3.1. Size of Structured Latents

구조화된 잠재 요소(SLAT)의 크기를 결정하기 위해 다양한 해상도와 채널 수를 사용하여 Sparse VAE를 학습했습니다. 표 3에 나타난 바와 같이, 해상도에서도 우수한 성능을 보여줬으나, 채널 수가 증가할수록 성능 향상은 미미했습니다. 그러나 해상도를 로 전환했을 때 성능이 크게 향상되었습니다. 우리는 효율성보다는 품질을 우선으로 고려하여 SLAT의 기본 설정으로 을 채택했습니다.

4.3.2. Rectified Flow v.s. Diffusion

Rectified Flow 모델과 널리 사용되는 Diffusion 모델을 비교했습니다(표 4 참조). 우리는 생성 파이프라인의 각 단계에서 Diffusion 모델을 Rectified Flow 모델로 대체하여 성능을 측정했습니다. 결과적으로, 생성 파이프라인의 어느 단계에서든 Diffusion 모델을 Rectified Flow 모델로 대체하면, 생성 품질과 입력 프롬프트와의 일치도가 모두 개선되었습니다.

4.3.3. 모델 크기 (Model Size)

우리는 모델 파라미터 수를 조정하여 성능을 평가했습니다(표 5 참조). 결과적으로, 모델 크기가 증가함에 따라 학습 데이터 분포와 Toys4k 데이터셋에서의 생성 성능이 지속적으로 향상되었습니다.

5. Limitation

모델이 3D 생성에서 강력한 성능을 보여주지만, 여전히 몇 가지 제한점이 존재함.

첫째, 2단계 생성 파이프라인의 비효율성. SLAT(structured latent representation)을 생성하기 위해 2단계 생성 파이프라인을 사용함. 이 파이프라인은 먼저 희소 구조(sparse structures)를 생성하고, 그 위에 로컬 잠재 요소(local latents)를 생성하는 방식임. 이러한 접근 방식은 단일 단계에서 완전한 3D 자산을 생성하는 end-to-end 방법에 비해 효율성이 낮을 수 있음.

둘째, Lighting 이슈. 이미지-기반 3D 생성 모델은 생성된 3D 자산에서 조명 효과(lighting effects)를 분리하지 않음. 그 결과, 참조 이미지의 그림자(shading)와 하이라이트(highlights)가 고정된 상태(baked-in)로 나타남. 이를 개선하기 위한 잠재적인 방법으로는, 훈련 중 이미지 프롬프트에 대해 더 강력한 조명 증강(lighting augmentation)을 적용하고, **Physically Based Rendering(PBR)**을 위한 재질(material) 예측을 모델에 추가하는 것이 있음. 이러한 부분은 향후 연구를 통해 탐구할 예정.