Neural Field Tokenizations with Hierarchy and Spatial Locality Priors

기술 요약: 계층 구조 및 공간적 국소성 사전 확률을 갖춘 뉴럴 필드 토큰화 (Neural Field Tokenizations with Hierarchy and Spatial Locality Priors)

문제 정의

뉴럴 필드는 데이터를 좌표에서 값으로의 연속 함수(예: $f_\theta: [0, 1]^2 \to \mathbb{R}^3$)로 매개변수화하며, 이는 다양한 모달리티를 위한 통합된 표현 학습 프레임워크를 제공합니다. 그러나 Functa 및 그 후속 모델들과 같은 기존의 모달리티 불가지론적(modality-agnostic) 접근 방식은 잠재 표현을 추론하기 위해 샘플별 메타 학습(예: MAML 또는 오토 디코딩)에 크게 의존합니다. 이러한 패러다임은 세 가지 결정적인 한계를 가집니다:

1. 확장성 병목 현상: 메타 학습은 2차 미분(second-order gradients)을 계산하기 위해 내부 루프 최적화를 위한 전체 계산 그래프를 저장해야 합니다. 이는 배치 크기와 메모리 효율성을 심각하게 제한하여 고해상도 또는 대규모 데이터셋으로의 확장을 저해합니다.
2. 구조의 부재: 최적화로부터 유도된 잠재 공간은 명시적인 구조적 사전 확률(structural priors)이 부족하여 학습이 비효식적이며, 분류와 같은 다운스트림 태스크에 사용하기 어렵습니다.
3. 트레이드오프: 유용한 구조적 사전 확률(국소성 및 계층 구조 등)을 도입하는 것과 모달리티 불가지론을 유지하는 것 사이에는 근본적인 긴장이 존재합니다. 구조(예: 2D 그리드, 포인트 클라우드)를 도입하는 대부분의 방법은 임의의 좌표 도메인을 다룰 수 있는 능력을 희생합니다.

방법론: LH-NeF

저자들은 모달리티 불가지론을 해치지 않으면서 연속 신호의 범용적인 토큰화된 표현을 학습하는 LH-NeF (Locality-preserving Hierarchical Neural Fields) 프레임워크를 제안합니다. 이 방법은 샘플별 최적화를 학습된 인코더-디코더 아키텍처를 통한 단일 순전파(forward pass)로 대체합니다.

1. LH-NeF 토크나이저 (인코더)
토크나이저는 계층적 퍼셉터(Hierarchical Perceivers, HiP)를 기반으로 구축하되, 좌표 공간의 국소성을 존중하도록 수정되었습니다:

• 국소성 보존 순서(Locality-Preserving Ordering): 좌표 공간에서의 공간적 국소성을 위반하는 고정된 래스터 스캔(raster scan)을 사용하는 표준 HiP와 달리, LH-NeF는 국소성을 보존하는 순열 $\pi$를 사용하여 입력 관측치 $\{(x_i, v_i)\}$를 재정렬합니다. 이는 공간 채우기 곡선(예: 유클리드 도메인을 위한 Morton/Z-order, 구형 데이터를 위한 S2 Hilbert curve)으로부터 유도됩니다. 이를 통해 토큰 시퀀스의 연속적인 그룹이 좌표 도메인 내에서 조밀한 공간 영역에 대응하도록 보장합니다.
• 계층적 그룹 어텐션(Hierarchical Grouped Attention): 재정렬된 시퀀스는 $L$개의 그룹 어텐션 블록을 통해 처리됩니다. 토큰은 미세한 그룹에서 거친 그룹으로 점진적으로 병합됩니다.
• 공간 메타데이터: 최종 출력은 라우팅 메타데이터인 중심점 $\mu_g$ (그룹 내 좌표의 프레셰 평균)와 범위 $\lambda_g$ (그룹의 반경 또는 경계 상자)를 동반하는 그룹화된 토큰 집합 $Y^{(L)}$입니다. 이를 통해 표현이 입력의 샘플링 밀도와 기하학적 구조에 적응할 수 있게 합니다.

2. LH-NeF 렌더러 (디코더)
토큰화와 쿼리 좌표 $x$가 주어지면, 렌더러는 다음 세 단계를 통해 필드 값 $f_\theta(x)$를 재구성합니다:

• k-그룹 소프트 라우팅(k-Group Soft Routing): 쿼리 $x$는 $\mu_g$와 $x$ 사이의 거리를 기준으로 가장 가까운 $k$개의 그룹으로 라우팅됩니다. 학습 가능한 가우시안 커널이 이 그룹들에 소프트 가중치 $w_g(x)$를 할당하여 연속성을 보장하고 그룹 간의 급격한 전환을 방지합니다.
• 교차 어텐션 집계(Cross-Attention Aggregation): 라우팅된 각 그룹에 대해, 쿼리 임베딩과 그룹의 토큰 간에 교차 어텐션이 수행됩니다. 출력은 가중 합을 통해 집계됩니다.
• 기하학 조건부 변조(Geometry-Conditioned Modulation): 집계된 특징 벡터는 가장 가까운 그룹 내에서의 상대적 위치 $\tilde{x}$(그룹의 범위로 정규화됨)를 조건으로 하는 FiLM 스타일의 메커니즘을 통해 변조됩니다. 이는 평행 이동과 스케일링에 대한 불변성을 제공하여, 네트워크가 서로 다른 그룹 전반에 걸쳐 일반화될 수 있는 로컬 패턴을 학습할 수 있게 합니다.

주요 기여

1. 트레이드오프 식별: 본 논문은 뉴럴 필드 학습에서 확장성, 구조적 사전 확률, 그리고 모달리티 불가지론 사이의 긴장을 식별하며, 계층 구조와 국소성이 모달리티 특화된 가정 없이도 도입될 수 있음을 주장합니다.
2. LH-NeF 프레임워크: 공간적으로 구조화된 그룹 토큰화를 생성하는 새로운 아키텍처를 제시합니다. 이 프레처는 소프트 그룹 라우팅과 교차 어텐션을 통해 임의의 좌표에서 필드를 디코딩하는 렌더러를 조건화하며, 이를 통해 매 샘플마다 최적화를 수행할 필요를 제거합니다.
3. 효율성 및 성능: 메모리 집약적인 내부 루프를 단일 순전파로 대체함으로써, 최첨단 성능을 유지하거나 능가하면서도 훈련 효율성을 크게 향상시켰습니다.

실험 결과

저자들은 이미지(CIFAR-10, CelebA-HQ, ImageNet), 3D 형상(ShapeNet), 기후 필드(ERA5)에 대해 LH-NeF를 평가했습니다.

• 재구성(Reconstruction): LH-NeF는 모달리티 불가지론적 방법들 중 최첨단 성능을 달성했습니다. 예를 들어, CelebA-HQ (64²)에서 198K 개의 파라미터만으로 44.4 dB PSNR을 달성하여, Functa (38.1 dB, 2.6M params)와 ENF (42.2 dB, 522K params)를 능가했습니다.
• 다운스트림 태스크:
  • 생성(Generation): CelebA-HQ에서 LH-NeF는 FID 9.7을 달성하여, 특화된 생성 뉴럴 필드 방법 및 모달리티 불가지론적 베이스라인들을 앞질렀습니다.
  • 분류(Classification): ShapeNet16에서, 동결된 토큰을 이용한 경량 분류기는 96.8% 정확도에 도달하여 ENF (96.6%)와 Functa (90.3%)를 능가했습니다.
  • 예측(Forecasting): ERA5에서, LH-NeF는 온도 예측에서 Functa보다 개선된 성능을 보였으나, 대칭성 특화 사전 확률을 사용하는 ENF에는 뒤처졌습니다.
• 효율성: 메타 학습의 내부 루프를 제거함으로써, LH-NeF는 동일한 하드웨어에서 가장 강력한 모달리티 불가지론적 베이스라인(ENF) 대비 약 42배 적은 메모리를 사용하며, 133배 더 큰 배치 크기(266 vs 2)를 지원합니다.

의의 및 주장

본 논문은 LH-NeF가 뉴럴 필드 표현 학습을 위한 확장 가능하고 모달리티 불가지론적인 기초 방법론 역할을 한다고 주장합니다. 계층 구조와 국소성을 유도 편향(inductive bias)으로 활용함으로써, 이 프레임워크는 순전파 추론의 효율성과 뉴럴 필드의 유연성 사이의 간극을 성공적으로 메웁니다.

저자들은 이 방법이 균일하게 샘플링된 데이터(이미지, 그리드)에 매우 효과적이며 불규칙한 샘플링(포인트 클라우드)에도 잘 적응하지만, 매우 불규칙한 데이터(예: LiDAR)에서의 국소성 보존은 기대치(in expectation)로만 성립한다고 언급했습니다. 또한, 설계된 편향을 사용하는 대신 어떤 편향이 태스크 전반에 걸쳐 일반화되는지를 학습하는 방향이나, 마지막 레벨뿐만 아니라 모든 인코더 레벨의 정보를 활용하도록 조건화 메커니즘을 확장하는 것에 대한 향후 연구 가능성을 제시했습니다.

이 연구는 모든 표현 학습 문제를 해결한다고 주장하는 것이 아니라, 구조화된 토큰화가 다중 모달 애플리케이션에 필요한 일반성을 희생하지 않으면서도 최적화 기반의 잠재 피팅(latent fitting)을 대체할 수 있음을 입증합니다.