SegviGen: 3D 부품 분할을 위한 3D 생성 모델의 재구성

기존 3D 부품 분할은 대규모 정답 데이터가 필요하거나 2D 투영 시 경계가 뭉개지는 한계가 있었다. SegviGen은 사전 학습된 3D 생성 모델의 지식을 활용해 아주 적은 데이터로도 정교한 분할을 가능하게 한다. 이는 고비용의 3D 데이터 구축 부담을 줄이면서도 산업용 모델링의 정밀도를 높이는 중요한 기술적 진보이다.

기존 3D 부품 분할은 점(Point)이나 면(Face)에 라벨을 붙이는 분류 문제로 접근했다. 하지만 3D 데이터는 주석을 달기가 매우 어렵고 비용이 많이 들어, 모델이 물체의 복잡한 구조를 충분히 학습하기 어려웠다. 특히 2D 이미지를 3D로 투영하는 방식은 시점 간 불일치로 인해 부품 간 경계가 뭉개지는 고질적인 문제가 있었다. SegviGen은 이 문제를 해결하기 위해 '이미 물체를 잘 만드는 법을 아는' 3D 생성 모델(Trellis.2)을 가져온다. 생성 모델은 수많은 3D 데이터를 보며 물체의 구조와 질감을 이미 이해하고 있다. 연구진은 이 모델에게 '물체를 새로 만들지 말고, 기존 물체의 부품별로 서로 다른 색을 칠해봐'라고 시키는 Colorization 방식을 도입했다. 이 방식은 생성 모델이 가진 날카로운 경계 표현 능력을 분할 작업에 그대로 이식한다. 결과적으로 아주 적은 양의 정답 데이터만으로도 모델은 물체의 논리적 구조를 파악해 정밀하게 부품을 나눌 수 있게 된다. 이는 마치 숙련된 조각가에게 색칠만 가르쳐서 부품을 구분하게 하는 것과 같아, 처음부터 해부학을 가르치는 것보다 훨씬 효율적이다.

SegviGen은 사전 학습된 3D VAE와 Flow Matching 기반의 Transformer 구조를 기반으로 한다. 입력된 3D 자산을 잠재 공간(Latent Space)으로 인코딩한 뒤, 이를 조건(Condition)으로 삼아 부품을 나타내는 색상을 생성하도록 설계했다. 이 과정에서 물체의 기하학적 구조를 유지하면서도 부품 간의 논리적 구분을 색상으로 표현한다. 분할 작업은 세 가지 모드로 구성된다. 사용자 클릭 좌표를 입력받아 특정 부품을 추출하는 대화형 분할, 랜덤 색상 팔레트를 이용해 전체 부품을 구분하는 전체 분할, 그리고 2D 세그멘테이션 맵을 조건으로 주어 3D 결과물을 제어하는 가이드 분할이다. 모든 모드는 동일한 아키텍처 내에서 Task ID Embedding을 통해 전환된다. 학습 시에는 Flow Matching 목적 함수를 사용한다. 정답 색상 데이터 y와 노이즈 epsilon 사이의 선형 보간 yt를 입력으로 하여, 모델 fθ가 노이즈 잔차를 예측하도록 한다. [yt, 기하학적 특징 z, 태스크 조건 C를 입력으로] → [Transformer 연산을 통해 노이즈를 예측하고] → [실제 노이즈와의 차이를 최소화하여] → [모델의 가중치를 업데이트]하는 구조이다.

대화형 부품 분할 실험에서 SegviGen은 PartObjaverse-Tiny와 PartNeXT 데이터셋 모두에서 압도적인 성능을 보였다. 특히 단 한 번의 클릭(IoU@1)만으로도 PartNeXT에서 54.86%의 IoU를 기록하며, 기존 모델인 Point-SAM(23.90%) 대비 약 31%p의 성능 향상을 달성했다. 전체 부품 분할(Full Segmentation)에서도 우수성이 입증됐다. 외부 2D 가이드 없이도 PartNeXT에서 55.40%의 IoU를 기록해 기존 최고 성능 모델들을 15% 이상 앞질렀다. 2D 가이드를 추가할 경우 성능은 71.53%까지 상승하여 정밀한 제어가 가능함을 보여주었다. 데이터 효율성 측면에서 SegviGen은 단 0.32%의 라벨링된 데이터만으로 학습되었음에도 불구하고, 대규모 데이터를 사용한 기존 모델들보다 더 선명한 경계와 정확한 부품 인식을 보여주었다. 이는 3D 생성 모델의 사전 학습된 지식이 분할 작업에 매우 효과적으로 전이됨을 의미한다.

아키텍처는 Trellis.2의 구조를 계승하며, Sparse Convolution VAE(SC-VAE)를 통해 3D 자산을 콤팩트한 잠재 표현으로 압축한다. 이 잠재 공간 위에서 Flow Matching 알고리즘을 적용해 기하학적 구조와 부품 색상을 동시에 최적화한다. 특히 옴니 복셀(Omni-voxel) 표현을 사용하여 기하학과 외형을 통합적으로 모델링한다. 핵심 메커니즘은 Part-wise Colorization이다. 분할 레이블을 직접 예측하는 대신, RGB 색상 공간에서 부품별 고유 색상을 생성하도록 유도한다. 이는 생성 모델의 가중치를 최대한 재사용하면서도 복잡한 부품 경계를 학습하는 데 유리하다. 또한 RoPE(Rotary Positional Embedding)를 활용해 3D 좌표 정보를 효과적으로 주입한다. Task ID Embedding을 도입하여 단일 모델이 세 가지 서로 다른 분할 시나리오를 처리할 수 있도록 했다. 각 태스크는 고유한 임베딩 벡터로 구분되며, 이는 DiT(Diffusion Transformer) 백본의 적응형 레이어(Adaptive Layers)에 주입되어 모델의 행동을 제어한다. 학습 시에는 여러 태스크의 샘플을 인터리빙(Interleaving)하여 공유 백본이 태스크별 특성을 동시에 학습하도록 유도한다. 구현 측면에서는 8개의 NVIDIA A800 GPU에서 단 8시간 만에 학습을 완료했다. 이는 기존의 무거운 3D 학습 방식 대비 매우 효율적이며, 12단계의 추론(Inference)만으로도 고품질의 분할 결과를 얻을 수 있도록 최적화되었다. 최적화 도구로는 AdamW를 사용했으며 학습률은 1e-4로 설정되었다.