계층적 트랜스포머를 이용한 무지도 3D 형태 추상화

초록

HiT(Hierarchical Transformer)는 3D 점군을 입력으로 받아, 압축 코드북과 교차‑어텐션을 활용해 다중 레벨의 파트 트리를 자동으로 생성한다. 각 레벨은 일정 수의 노드를 갖는 코드북으로 구성되며, 하위 레벨 파트는 상위 파트와의 어텐션 매트릭스를 통해 부드럽게 연결된다. 파트는 3D 볼록 프리미티브(컨벡스)로 매핑되고, 부모‑자식 관계는 점유 함수의 곱셈으로 공간적 포함을 강제한다. 재구성 손실과 정규화 항을 결합한 자기지도 학습으로 ShapeNet 전 55개 카테고리에서 의미 있는 계층적 세그멘테이션을 달성한다.

상세 분석

HiT는 기존의 무지도 파트 학습이 갖는 “고정된 트리 구조” 제약을 완전히 탈피한다는 점에서 혁신적이다. 논문은 먼저 점군을 ConvOccNet 기반 인코더로 변환해 3차원 격자 형태의 피처 매핑 Z⁽⁰⁾을 만든다. 이후 L개의 디코더 레이어가 순차적으로 쌓이며, 각 레이어 ℓ은 Nℓ개의 학습 가능한 코드북 C⁽ℓ⁾을 갖는다. 코드북은 기존 트랜스포머의 쿼리 역할을 수행해 이전 레이어의 피처 Z⁽ℓ⁻¹⁾와 어텐션을 계산한다. 어텐션 행렬 A⁽ℓ⁾는 소프트맥스 후 √D로 정규화되며, 이를 통해 각 하위 파트가 어느 상위 파트에 속할지 확률적으로 할당한다. 논문은 이 확률 할당을 Straight‑Through Estimator(STE)와 결합해 전방에서는 하드 원‑핫 선택을, 역전파에서는 부드러운 그래디언트를 유지하도록 설계했다.

기하학적 파트 표현은 CvxNet에서 영감을 받아 3D 볼록 프리미티브(다수의 반평면 교집합)로 매핑된다. 각 파트는 H개의 반평면 파라미터(법선, 오프셋, 블렌딩 가중치)와 rigid 변환(회전, 이동, 스케일)으로 정의되며, 점유 함수 ˜Oₛ(x)=σ(−Φₛ(x)) 형태의 시그모이드 SDF를 사용한다. 여기서 Φₛ는 로컬 좌표계에서의 로그합(exp) 형태로, 볼록성을 보장한다.

공간적 포함 제약은 부모 파트의 점유 함수 ˆOₚ와 자식 파트의 점유 함수 ˜Oₛ를 곱해 ˆOₛ(x)=ˆOₚ(x)·˜Oₛ(x) 로 구현한다. 이는 자식 파트가 부모 파트 내부에서만 활성화되도록 강제한다. 손실 함수는 네 부분으로 구성된다: (1) 레벨별 재구성 손실 L_recon, (2) 포함 손실 L_contain(부모‑자식 점유 차이 최소화), (3) CvxNet에서 차용한 볼록 프리미티브 정규화 L_cvxnet(과도한 중첩 방지), (4) 트리 균형을 위한 L_balance(각 레벨의 파트 수가 사전에 지정된 Nℓ에 가깝게 유지).

실험에서는 ShapeNet 전체 55개 카테고리에 대해 무지도 파트 세그멘테이션을 수행했으며, IoU 기준으로 기존 최첨단 방법(AE‑Net, DAE‑Net, RIMNet 등)을 모두 앞섰다. 특히 의자와 같은 카테고리에서 파트 수가 가변적인 경우, HiT는 동적으로 트리 구조를 조정해 파트 수가 달라지는 현상을 자연스럽게 포착한다. 시각적 결과는 3레벨 이상의 계층을 보여주며, 상위 레벨은 의미 있는 큰 부품(예: 의자 좌석, 다리), 하위 레벨은 세부 구조(예: 다리의 각 관절)로 분해된다.

핵심 기여는 (i) 코드북 기반의 정보 병목을 통해 파트 재사용성을 학습, (ii) 교차‑어텐션을 이용한 부드러운 부모‑자식 매핑, (iii) 볼록 프리미티브와 점유 곱셈을 통한 명시적 공간 포함, (iv) 고정된 트리 구조 없이 데이터에 의해 계층이 자동 생성된다는 점이다. 이러한 설계는 3D 형태 이해, 파트 기반 편집, 로봇 조작 계획 등 다양한 downstream 작업에 직접 활용 가능성을 열어준다.