Neural Field Tokenizations with Hierarchy and Spatial Locality Priors

Resumo Técnico: Tokenizações de Campos Neurais com Hierarquia e Priores de Localidade Espacial

Declaração do Problema

Campos neurais parametrizam dados como funções contínuas de coordenadas para valores (ex: $f_\theta: [0, 1]^2 \to \mathbb{R}^3$), oferecendo um framework unificado para o aprendizado de representação entre modalidades. No entanto, abordagens agnósticas à modalidade existentes, como o Functa e seus descendentes, dependem fortemente de meta-aprendizado por amostra (ex: MAML ou auto-decodificação) para inferir representações latentes. Este paradigma apresenta três limitações críticas:

1. Gargalo de Escalabilidade: O meta-aprendizado requer o armazenamento de todo o grafo de computação para a otimização do loop interno para computar gradientes de segunda ordem. Isso limita severamente os tamanhos de lote (batch sizes) e a eficiência de memória, dificultando a escala para altas resoluções ou grandes conjuntos de dados.
2. Falta de Estrutura: Espaços latentes derivados de otimização carecem de priores estruturais explícitos, tornando o aprendizado ineficiente e as representações resultantes difíceis de usar para tarefas como classificação.
3. Compromisso (Trade-off): Existe uma tensão fundamental entre introduzir priores estruturais úteis (como localidade e hierarquia) e manter a agnosticidade à modalidade. A maioria dos métodos que introduz estrutura (ex: grades 2D, nuvens de pontos) sacrifica a capacidade de lidar com domínios de coordenadas arbitrários.

Metodologia: LH-NeF

Os autores propõem o LH-NeF (Locality-preserving Hierarchical Neural Fields), um framework que aprende representações tokenizadas de sinais contínuos de propósito geral sem comprometer a agnosticidade à modalidade. O método substitui a otimização por amostra por uma única passagem direta (forward pass) através de uma arquitetura codificador-decodificador aprendida.

1. O Tokenizador LH-NeF (Codificador)
O tokenizador baseia-se em Hierarchical Perceivers (HiP), mas os modifica para respeitar a localidade no espaço de coordenadas:

• Ordenação Preservadora de Localidade: Ao contrário do HiP padrão, que utiliza um escaneamento raster fixo (violando a localidade espacial no espaço de coordenadas), o LH-NeF reordena as observações de entrada $\{(x_i, v_i)\}$ usando uma permutação $\pi$ que preserva a localidade. Esta é derivada de curvas de preenchimento de espaço (ex: Morton/Z-order para domínios euclidianos, curvas de Hilbert S2 para dados esféricos). Isso garante que grupos contíguos na sequência de tokens correspondam a regiões espaciais compactas no domínio de coordenadas.
• Atenção Agrupada Hierárquica: A sequência reordenada é processada através de $L$ blocos de atenção agrupada. Os tokens são progressivamente fundidos de grupos de granularidade fina para grupos mais grosseiros.
• Metadados Espaciais: O output final é um conjunto de tokens agrupados $Y^{(L)}$ acompanhado de metadados de roteamento: o centroide $\mu_g$ (média de Fréchet das coordenadas no grupo) e a extensão $\lambda_g$ (raio espacial ou caixa delimitadora). Isso permite que a representação se adapte à densidade de amostragem e à geometria da entrada.

2. O Renderizador LH-NeF (Decodificador)
Dada a tokenização e uma coordenada de consulta $x$, o renderizador reconstrói o valor do campo $f_\theta(x)$ através de três etapas:

• Roteamento Suave por k-Grupos (k-Group Soft Routing): A consulta $x$ é roteada para os $k$ grupos mais próximos com base na distância entre $x$ e os centroides dos grupos $\mu_g$. Um kernel Gaussiano aprendível atribui pesos suaves $w_g(x)$ a esses grupos, garantindo continuidade e evitando mudanças abruptas entre grupos.
• Agregação por Atenção Cruzada: Para cada grupo roteado, a atenção cruzada é realizada entre o embedding da consulta e os tokens do grupo. Os outputs são agregados via uma soma ponderada.
• Modulação Condicionada à Geometria: O vetor de característica agregado é modulado usando um mecanismo do tipo FiLM condicionado à posição relativa da consulta $\tilde{x}$ dentro do grupo mais próximo (normalizada pela extensão do grupo). Isso fornece invariância à translação e escala, permitindo que a rede aprenda padrões locais que generalizam através de diferentes grupos.

Principais Contribuições

1. Identificação de um Compromisso: O artigo identifica a tensão entre escalabilidade, priores estruturais e agnosticidade à modalidade, argumentando que hierarquia e localidade podem ser introduzidas sem suposições específicas de modalidade.
2. Framework LH-NeF: Uma arquitetura inovadora que produz tokenizações de grupos espacialmente estruturadas. Ela condiciona um renderizador que decodifica campos em coordenadas arbitrárias via roteamento suave de grupos e atenção cruzada, eliminando a necessidade de otimização por amostra.
3. Eficiência e Desempenho: O método substitui o loop interno intensivo em memória do meta-aprendizado por uma única passagem direta, alcançando ganhos significativos em eficiência de treinamento enquanto iguala ou supera o desempenho do estado da arte.

Resultados Experimentais

Os autores avaliaram o LH-NeF em imagens (CIFAR-10, CelebA-HQ, ImageNet), formas 3D (ShapeNet) e campos climáticos (ERA5).

• Reconstrução: O LH-NeF atinge o desempenho de estado da arte entre os métodos agnósticos à modalidade. Por exemplo, no CelebA-HQ (64²), atinge 44.4 dB de PSNR com apenas 198K parâmetros, superando o Functa (38.1 dB, 2.6M params) e o ENF (42.2 dB, 522K params).
• Tarefas de Jusante (Downstream Tasks):
  • Geração: No CelebA-HQ, o LH-NeF atinge um FID de 9.7, superando métodos de campo neural generativos especializados e baselines agnósticos à modalidade.
  • Classificação: No ShapeNet16, um classificador leve sobre tokens congelados atinge 96.8% de acurácia, superando o ENF (96.6%) e o Functa (90.3%).
  • Previsão: No ERA5, o LH-NeF melhora em relação ao Functa na previsão de temperatura, embora fique atrás do ENF (que utiliza priores específicos de simetria).
• Eficiência: Ao remover o loop interno de meta-aprendizado, o LH-NeF usa ~42× menos memória e suporta 133× maiores tamanhos de lote (266 vs 2) comparado ao baseline agnóstico à modalidade mais forte (ENF) no mesmo hardware.

Significância e Alegações

O artigo afirma que o LH-NeF serve como um método fundamental escalável e agnóstico à modalidade para o aprendizado de representação de campos neurais. Ao alavancar hierarquia e localidade como vieses indutivos, o framework consegue unir com sucesso a eficiência da inferência por passagem direta com a flexibilidade dos campos neurais.

Os autores observam que, embora o método seja altamente eficaz para dados amostrados uniformemente (imagens, grades) e se adapte bem a amostragem irregular (nuvens de pontos), a preservação da localidade em dados altamente irregulares (ex: LiDAR) ocorre apenas em expectativa. Eles também sugerem que trabalhos futuros poderiam explorar o aprendizado de quais vieses generalizam para as tarefas em vez de projetar manualmente, e estender o mecanismo de condicionamento para explorar informações de todos os níveis do codificador, em vez de apenas o final.

O trabalho não pretende resolver todos os problemas de aprendizado de representação, mas demonstra que a tokenização estruturada pode substituir o ajuste de latente baseado em otimização sem sacrificar a generalidade necessária para aplicações multimodais.