Neural Field Tokenizations with Hierarchy and Spatial Locality Priors

Sintesi Tecnica: Tokenizzazioni di Campi Neurali con Gerarchia e Prior di Località Spaziale

Definizione del Problema

I campi neurali parametrizzano i dati come funzioni continue da coordinate a valori (ad esempio, $f_\theta: [0, 1]^2 \to \mathbb{R}^3$), offrendo un framework unificato per l'apprendimento di rappresentazioni attraverso diverse modalità. Tuttavia, gli approcci agnostici rispetto alla modalità esistenti, come Functa e i suoi discendenti, si affidano pesantemente al meta-learning per singolo campione (ad esempio, MAML o auto-decodifica) per inferire le rappresentazioni latenti. Questo paradigma presenta tre criticità fondamentali:

1. Collo di bottiglia della scalabilità: Il meta-learning richiede la memorizzazione dell'intero grafo computazionale per l'ottimizzazione dell'inner-loop al fine di calcolare i gradienti del secondo ordine. Ciò limita severamente le dimensioni dei batch e l'efficienza della memoria, ostacolando la scalabilità verso risoluzioni elevate o dataset estesi.
2. Mancanza di struttura: Gli spazi latenti derivati dall'ottimizzazione mancano spesso di espliciti prior strutturali, rendendo l'apprendimento inefficiente e le rappresentazioni risultanti difficili da utilizzare per task a valle come la classificazione.
3. Compromesso (Trade-off): Esiste una tensione fondamentale tra l'introduzione di utili prior strutturali (come località e gerarchia) e il mantenimento dell'agnosticismo rispetto alla modalità. La maggior parte dei metodi che introducono struttura (ad esempio, griglie 2D, nuvole di punti) sacrifica la capacità di gestire domini di coordinate arbitrari.

Metodologia: LH-NeF

Gli autori propongono LH-NeF (Locality-preserving Hierarchical Neural Fields), un framework che apprende rappresentazioni tokenizzate general-purpose di segnali continui senza compromettere l'agnosticismo rispetto alla modalità. Il metodo sostituisce l'ottimizzazione per singolo campione con un singolo passaggio in avanti (forward pass) attraverso un'architettura encoder-decoder appresa.

1. Il Tokenizer LH-NeF (Encoder)
Il tokenizer si basa sui Hierarchical Perceivers (HiP) ma li modifica per rispettare la località nello spazio delle coordinate:

• Ordinamento preservante la località: A differenza del normale HiP, che utilizza una raster scan fissa (violando la località spaziale nello spazio delle coordinate), LH-NeF riordina le osservazioni di input $\{(x_i, v_i)\}$ utilizzando una permutazione $\pi$ che preserva la località. Questa è derivata da curve space-filling (ad esempio, Morton/Z-order per domini euclidei, curve di Hilbert S2 per dati sferici). Ciò garantisce che gruppi contigui nella sequenza di token corrispondano a regioni spaziali compatte nel dominio delle coordinate.
• Attenzione raggruppata gerarchica: La sequenza riordinata viene elaborata attraverso $L$ blocchi di attenzione raggruppata. I token vengono progressivamente uniti da gruppi a grana fine a gruppi più grossolani.
• Metadati spaziali: L'output finale è un insieme di token raggruppati $Y^{(L)}$ accompagnato da metadati di routing: il centroide $\mu_g$ (media di Fréchet delle coordinate nel gruppo) e l'estensione $\lambda_g$ (raggio spaziale o bounding box). Ciò consente alla rappresentazione di adattarsi alla densità di campionamento e alla geometria dell'input.

2. Il Renderer LH-NeF (Decoder)
Data la tokenizzazione e una coordinata di query $x$, il renderer ricostruisce il valore del campo $f_\theta(x)$ attraverso tre passaggi:

• Soft Routing a k-Gruppi: La query $x$ viene instradata ai $k$ gruppi più vicini in base alla distanza tra $x$ e i centroidi dei gruppi $\mu_g$. Un kernel gaussiano apprendibile assegna pesi soft $w_g(x)$ a questi gruppi, garantendo la continuità ed evitando commutazioni brusche tra i gruppi.
• Aggregazione tramite Cross-Attention: Per ogni gruppo instradato, viene eseguita una cross-attention tra l'embedding della query e i token del gruppo. Gli output vengono aggregati tramite una somma pesata.
• Modulazione condizionata dalla geometria: Il vettore di feature aggregato viene modulato utilizzando un meccanismo di tipo FiLM condizionato sulla posizione relativa della query $\tilde{x}$ all'interno del gruppo più vicino (normalizzata rispetto all'estensione del gruppo). Ciò fornisce invarianza alla traslazione e alla scala, permettendo alla rete di apprendere pattern locali che generalizzano attraverso diversi gruppi.

Contributi Chiave

1. Identificazione di un compromesso: Il paper identifica la tensione tra scalabilità, prior strutturali e agnosticismo rispetto alla modalità, sostenendo che gerarchia e località possono essere introdotte senza assunzioni specifiche della modalità.
2. Framework LH-NeF: Un'architettura innovativa che produce tokenizzazioni di gruppo spazialmente strutturate. Essa condiziona un renderer che decodifica i campi ad arbitrarie coordinate tramite soft group routing e cross-attention, eliminando la necessità di ottimizzazione per singolo campione.
3. Efficienza e Performance: Il metodo sostituisce l'inner-loop ad alta intensità di memoria del meta-learning con un singolo passaggio in avanti, ottenendo guadagni significativi nell'efficienza dell'addestramento pur eguagliando o superando le prestazioni dello stato dell'arte.

Risultati Sperimentali

Gli autori hanno valutato LH-NeF su immagini (CIFAR-10, CelebA-HQ, ImageNet), forme 3D (ShapeNet) e campi climatici (ERA5).

• Ricostruzione: LH-NeF raggiunge le prestazioni state-of-the-art tra i metodi agnostici rispetto alla modalità. Ad esempio, su CelebA-HQ (64²), raggiunge 44.4 dB di PSNR con soli 198K parametri, superando Functa (38.1 dB, 2.6M params) ed ENF (42.2 dB, 522K params).
• Task a valle:
  • Generazione: Su CelebA-HQ, LH-NeF ottiene un FID di 9.7, superando i metodi neural field generativi specializzati e le baseline agnostiche rispetto alla modalità.
  • Classificazione: Su ShapeNet16, un classificatore leggero su token congelati raggiunge il 96.8% di accuratezza, superando ENF (96.6%) e Functa (90.3%).
  • Previsione (Forecasting): Su ERA5, LH-NeF migliora rispetto a Functa nella previsione della temperatura, sebbene resti indietro rispetto a ENF (che utilizza prior specifici per la simmetria).
• Efficienza: Rimuovendo l'inner-loop del meta-learning, LH-NeF utilizza circa 42 volte meno memoria e supporta 133 volte batch più grandi (266 vs 2) rispetto alla più forte baseline agnostica rispetto alla modalità (ENF) sullo stesso hardware.

Significato e Rivendicazioni

Il paper sostiene che LH-NeF possa servire come un metodo fondamentale, scalabile e agnostico rispetto alla modalità per l'apprendimento di rappresentazioni di campi neurali. Sfruttando la gerarchia e la località come bias induttivi, il framework riesce a colmare il divario tra l'efficienza dell'inferenza tramite forward-pass e la flessibilità dei campi neurali.

Gli autori osservano che, sebbene il metodo sia altamente efficace per dati campionati uniformemente (immagini, griglie) e si adatti bene a campionamenti irregolari (nuvole di punti), la preservazione della località su dati altamente irregolari (ad esempio, LiDAR) è valida solo in termini di aspettativa. Suggeriscono inoltre che il lavoro futuro potrebbe esplorare l'apprendimento di quali bias generalizzano meglio attraverso i task invece di progettare manualmente i bias, ed estendere il meccanismo di condizionamento per sfruttare le informazioni da tutti i livelli dell'encoder piuttosto che solo dall'ultimo.

Il lavoro non pretende di risolvere tutti i problemi di apprendimento delle rappresentazioni, ma dimostra che una tokenizzazione strutturata può sostituire l'adattamento del latente basato sull'ottimizzazione senza sacrificare la generalità richiesta per le applicazioni multi-modali.