Neural Field Tokenizations with Hierarchy and Spatial Locality Priors

Technische Samenvatting: Neural Field Tokenisaties met Hiërarchie en Ruimtelijke Localiteitsprioriteiten

Probleemstelling

Neural fields parametriseren data als continue functies van coördinaten naar waarden (bijv. $f_\theta: [0, 1]^2 \to \mathbb{R}^3$), wat een verenigd kader biedt voor representatie-leren over verschillende modaliteiten heen. Echter, bestaande modaliteits-agnostische benaderingen, zoals Functa en diens afstammelingen, vertrouwen zwaar op per-sample meta-learning (bijv. MAML of auto-decoding) om latente representaties af te leiden. Dit paradigma presenteert drie kritieke beperkingen:

1. Schaalbaarheid bottleneck: Meta-learning vereist het opslaan van de volledige computationele graaf voor de inner-loop optimalisatie om tweede-orde gradiënten te berekenen. Dit beperkt de batchgroottes en de geheugenefficiëntie ernstig, wat schalen naar hoge resoluties of grote datasets belemmert.
2. Gebrek aan Structuur: Uit optimalisatie afgeleide latente ruimtes missen vaak expliciete structurele priors, waardoor het leren inefficiënt is en de resulterende representaties moeilijk te gebruiken zijn voor downstream taken zoals classificatie.
3. Trade-off: Er is een fundamentele spanning tussen het introduceren van nuttige structurele priors (zoals localiteit en hiërarchie) en het behouden van modaliteits-agnosticisme. De meeste methoden die structuur introduceren (bijv. 2D-grids, puntenwolken) offeren het vermogen op om willekeurige coördinaat-domeinen te verwerken.

Methodologie: LH-NeF

De auteurs stellen LH-NeF (Locality-preserving Hierarchical Neural Fields) voor, een framework dat algemene, getokeniseerde representaties van continue signalen leert zonder het modaliteits-agnosticisme in gevaar te brengen. De methode vervangt per-sample optimalisatie door een enkele forward pass door een geleerd encoder-decoder architectuur.

1. De LH-NeF Tokenizer (Encoder)
De tokenizer bouwt voort op Hierarchical Perceivers (HiP) maar wijzigt deze om de localiteit in de coördinatieruimte te respecteren:

• Localiteit-behoudende Ordening: In tegen tegenstelling tot standaard HiP, die een vaste rasterscan gebruikt (wat de ruimtelijke localiteit in de coördinaatruimte schendt), herordent LH-NeF de input-observaties $\{(x_i, v_i)\}$ met behulp van een localiteit-behoudende permutatie $\pi$. Deze wordt afgeleid van space-filling curves (bijv. Morton/Z-order voor Euclidische domeinen, S2 Hilbert curves voor sferische data). Dit zorgt ervoor dat opeenvolgende groepen in de token-sequentie overeenkomen met compacte ruimtelijke regio's in het coördinaatdomein.
• Hiërarchische Gegroepeerde Attention: De geordende sequentie wordt verwerkt door $L$ blokken van gegroepeerde attention. Tokens worden progressief samengevoegd van fijnmazige groepen naar grovere groepen.
• Ruimtelijke Metadata: De uiteindelijke output is een set gegroepeerde tokens $Y^{(L)}$ vergezeld van routing metadata: het zwaartepunt $\mu_g$ (Fréchet-gemiddelde van coördinaten in de groep) en de omvang $\lambda_g$ (ruimtelijke radius of bounding box). Dit stelt de representatie in staat zich aan te passen aan de bemonsteringsdichtheid en geometrie van de input.

2. De LH-NeF Renderer (Decoder)
Gegeven de tokenisatie en een query-coördinaat $x$, reconstrueert de renderer de veldwaarde $f_\theta(x)$ via drie stappen:

• k-Groep Soft Routing: De query $x$ wordt gerouteerd naar de $k$ dichtstbijzijnde groepen op basis van de afstand tussen $x$ en de groepszwaartepunten $\mu_g$. Een leerbaar Gaussisch kernel wijst zachte gewichten $w_g(x)$ toe aan deze groepen, wat continuïteit waarborgt en abrupte overgangen tussen groepen voorkomt.
• Cross-Attention Aggregatie: Voor elke gerouteerde groep wordt cross-attention uitgevoerd tussen de query-embedding en de groeps-tokens. De outputs worden geaggregeerd via een gewogen som.
• Geometrie-geconditioneerde Modulatie: De geaggregeerde feature vector wordt gemoduleerd via een FiLM-achtige mechanisme, geconditioneerd op de relatieve positie $\tilde{x}$ van de query binnen de dichtstbijzijnde groep (genormaliseerd door de omvang van de groep). Dit biedt invariantie voor translatie en schaling, waardoor het netwerk lokale patronen kan leren die generaliseren over verschillende groepen.

Belangrijkste Bijdragen

1. Identificatie van een Trade-off: Het artikel identificeert de spanning tussen schaalbaarheid, structurele priors en modaliteits-agnosticisme in neural field learning, en betoogt dat hiërarchie en localiteit kunnen worden geïntroduceerd zonder modaliteit-specifieke aannames.
2. LH-NeF Framework: Een nieuwe architectuur die ruimtelijk gestructureerde gegroepeerde tokenisaties produceert. Het conditioneert een renderer die velden op willekeurige coördinaten decodeert via soft group routing en cross-attention, waardoor de noodzaak voor per-sample optimalisatie wordt geëlimineerd.
3. Efficiëntie en Prestaties: De methode vervangt de geheugenintensieve inner-loop van meta-learning door een enkele forward pass, wat aanzienlijke winst oplevert in trainingsefficiëntie terwijl de state-of-the-art prestaties worden geëvenaard of overtroffen.

Experimentele Resultaten

De auteurs evalueerden LH-NeF op afbeeldingen (CIFAR-10, CelebA-HQ, ImageNet), 3D-vormen (ShapeNet) en klimaatvelden (ERA5).

• Reconstructie: LH-NeF bereikt state-of-the-art prestaties onder de modaliteits-agnostische methoden. Bijvoorbeeld, op CelebA-HQ (64²) bereikt het 44.4 dB PSNR met slechts 198K parameters, waarmee het Functa (38.1 dB, 2.6M params) en ENF (42.2 dB, 522K params) overtreft.
• Downstream Taken:
  • Generatie: Op CelebA-HQ bereikt LH-NeF een FID van 9.7, waarmee het gespecialiseerde generatieve neural field methoden en modaliteits-agnostische baselines overtreft.
  • Classificatie: Op ShapeNet16 bereikt een lichtgewicht classifier op bevroren tokens een nauwkeurigheid van 96.8%, waarmee het ENF (96.6%) en Functa (90.3%) overtreft.
  • Voorspelling: Op ERA5 verbetert LH-NeF de temperatuurvoorspelling ten opzichte van Functa, hoewel het achterblijft bij ENF (dat symmetrie-specifieke priors gebruikt).
• Efficiëntie: Door de meta-learning inner-loop te verwijderen, gebruikt LH-NeF ~42× minder geheugen en ondersteunt het 133× grotere batches (266 vs. 2) vergeleken met de sterkste modaliteits-agnostische baseline (ENF) op dezelfde hardware.

Betekenis en Claims

Het artikel claimt dat LH-NeF dient als een schaalbare, modaliteits-agnostische fundamentele methode voor neural field representatie-leren. Door hiërarchie en localiteit te gebruiken als inductieve biases, slaagt het framework erin de kloof te overbruggen tussen de efficiëntie van forward-pass inferentie en de flexibiliteit van neural fields.

De auteurs merken op dat hoewel de methode zeer effectief is voor uniform gesampleerde data (afbeeldingen, grids) en goed aanpasbaar is aan onregelmatige sampling (puntenwolken), de localiteit-behoud op zeer onregelmatige data (bijv. LiDAR) slechts in verwachting geldt. Ze suggereren ook dat toekomstig werk zou kunnen verkennen welke biases generaliseren over taken in plaats van het handmatig ontwerpen ervan, en het uitbreiden van het conditioneringsmechanisme om informatie van alle encoder-niveaus te exploiteren in plaats van alleen het laatste niveau.

Het werk claimt niet het probleem van alle representatie-leren op te lossen, maar demonstreert dat gestructureerde tokenisatie de optimalisatie-gebaseerde latente fitting kan vervangen zonder het algemene karakter op te offeren dat vereist is voor multi-modale toepassingen.