Bases of Steerable Kernels for Equivariant CNNs: From 2D Rotations to the Lorentz Group

Probleemstelling

In Convolutionele Neurale Netwerken (CNNs) is het integreren van symmetrie-priors (zoals rotatie- of Lorentz-invariantie) cruciaal voor de prestaties, vooral bij taken zoals het voorspellen van atomaire krachten of het analyseren van relativistische systemen. Een "stuurbaar" (steerable) equivariant CNN vereist dat de convolutie-kern $K$ voldoet aan een specifieke stuurbaarheidsbeperking (steerability constraint).

Voor een groep $G$ die inwerkt op invoer- en uitvoer-kenmerkkaarten via representaties $\rho_{in}$ en $\rho_{out}$, moet de kern $K$ voldoen aan:
$$K(g \cdot x) = \frac{1}{|\det g|} \rho_{out}(g) K(x) \rho_{in}(g)^{-1}$$

Het traditionele probleem is het vinden van een volledige basis van oplossingen voor deze vergelijking voor verschillende groepen en tensor-types. Bestaande methoden (zoals die in [2] en [9]) maken vaak gebruik van:

1. Clebsch-Gordan-coëfficiënten: Deze zijn computatierijk en moeilijk te berekenen voor bepaalde groepen (zoals de Lorentz-groep).
2. Harmonische basisfuncties: Deze vereisen vaak een complexe basiswissel naar een "gekoppeld" (coupled) basisstelsel en terug.

Dit leidt tot complexe implementaties en hoge rekentijd, en maakt het moeilijk om deze methoden toe te passen op niet-compacte groepen zoals de Lorentz-groep $SO^+(1,3)$.

Methodologie

De auteur presenteert een alternatieve, directere strategie om de ruimte van oplossingen voor de stuurbaarheidsbeperking te vinden. De kern van de methode is het vermijden van Clebsch-Gordan-coëfficiënten en het werken direct met de representaties van de invoer en uitvoer.

De aanpak bestaat uit drie stappen:

1. Vaststellen van een referentiepunt: Kies een punt $x_0$ op een baan (orbit) van de groep $G$. Definieer de stabilisatorgroep (isotropegroep) $H = \{g \in G \mid g \cdot x_0 = x_0\}$.
2. Oplossen van de invariantie bij $x_0$: De beperking voor de kern op het punt $x_0$ vereist dat $K(x_0)$ een homomorfisme is tussen de restricties van de representaties $\rho_{in}$ en $\rho_{out}$ tot de subgroep $H$.
   • Omdat $H$ vaak een kleinere, eenvoudiger groep is (bijv. $SO(2)$ binnen $SO(3)$), kunnen deze homomorfismen eenvoudig worden gevonden met behulp van Schur's Lemma.
   • Dit resulteert in een basis van blokmatrices die de irreducibele componenten van $H$ koppelen.
3. Sturen (Steering): Zodra de algemene vorm van $K(x_0)$ bekend is, wordt de volledige kern $K(x)$ voor willekeurige $x$ in de baan verkregen door de stuurvergelijking toe te passen:
   $$K(g \cdot x_0) = \rho_{out}(g) K(x_0) \rho_{in}(g)^{-1}$$
   Hierbij wordt $x$ geschreven als $g \cdot x_0$.

Deze methode omzeilt de noodzaak om naar een gekoppelde basis te wisselen en Clebsch-Gordan-coëfficiënten te berekenen. De oplossingen worden direct uitgedrukt in termen van de matrixelementen van de oorspronkelijke representaties (zoals Wigner D-matrices).

Belangrijkste Bijdragen

• Algemene Strategie: Een uniforme methode om bases van stuurbare kernen te construeren voor elke Lie-groep (compact en niet-compact) zonder Clebsch-Gordan-coëfficiënten.
• Expliciete Bases: Het paper levert expliciete, "klaar-voor-gebruik" formules voor zowel reële als complexe representaties voor:
  • SO(2) en O(2): 2D rotaties en spiegelingen.
  • SO(3) en O(3): 3D rotaties en spiegelingen (essentieel voor moleculaire dynamica).
  • SO(1, 3) (Lorentz-groep): De eerste systematische afleiding van stuurbare kernen voor de Lorentz-groep, relevant voor relativistische fysica.
• Behandeling van Niet-Compacte Groepen: De methode wordt succesvol toegepast op de Lorentz-groep, waarbij onderscheid wordt gemaakt tussen:
  • Massieve deeltjes: Gebaseerd op hyperboloïden met stabilisator $SO(3)$.
  • Massaloze deeltjes: Gebaseerd op lichtkegels met stabilisator $ISO(2)$ (waarbij de translatie-deel wordt genegeerd voor fysisch relevante representaties).
• Projector-Interpretatie: De auteur toont aan dat de gevonden intertwiners (kernen) vaak kunnen worden geïnterpreteerd als projectoren op invariant deelruimtes binnen ruimtetijd-tensoren (bijv. projecteren op de ruimtelijke componenten loodrecht op een tijdsvector).

Resultaten

Het paper levert concrete formules voor de basis-elementen van de kernen:

• Voor SO(3) worden de kernen uitgedrukt in termen van Wigner D-matrices en reële sferische harmonischen. De dimensie van de leerbare parameter-ruimte wordt nauwkeurig bepaald voor reële versus complexe representaties.
• Voor de Lorentz-groep worden kernen afgeleid voor verschillende spin-waarden:
  • Spin 0: Projectie op de tijdscomponent ($u_\mu u_\nu$).
  • Spin 1: Projectie op de ruimtelijke componenten loodrecht op de snelheidsvector ($\Delta_{\mu\nu} = \eta_{\mu\nu} - u_\mu u_\nu$).
  • Spin 1/2 en 3/2: Gebruikmakend van Dirac-matrices, ladingsconjugatie ($C$) en projectoren op de fysieke deelruimtes (zoals Rarita-Schwinger velden).
  • Massaloze velden: Projectoren gebaseerd op een lichtachtige vector $n$ en een hulp-vector $\bar{n}$, die gauge-invariantie respecteren.

De afgeleide bases zijn volledig en vormen een vectorruimte die equivalent is aan die verkregen via traditionele methoden, maar dan in een meer direct bruikbare vorm.

Significantie

• Efficiëntie: Door Clebsch-Gordan-coëfficiënten te omzeilen, wordt de berekeningslast voor het ontwerpen van equivariante netwerken aanzienlijk verlaagd, vooral voor complexe groepen.
• Toegankelijkheid: De methode maakt gebruik van minimale technische hulpmiddelen (stabilisatorgroepen en Schur's Lemma), waardoor het toegankelijker is voor een breder publiek dan de zware representatietheorie die vaak nodig is.
• Fysica-toepassingen: De uitbreiding naar de Lorentz-groep opent de deur voor het ontwerpen van CNNs die direct toepasbaar zijn op relativistische systemen (zoals deeltjesfysica en cosmologie), waar eerdere methoden beperkt waren of te complex waren.
• Implementatie: De methode suggereert een flexibeler manier om afsnijdingen (cut-offs) toe te passen op de feature maps (onafhankelijk voor input en output) in plaats van op de gekoppelde basis, wat de expressiviteit van het netwerk kan verhogen.

Kortom, dit paper biedt een krachtig, universeel raamwerk voor het bouwen van symmetrische neurale netwerken die verder gaan dan de gebruikelijke rotaties in 2D/3D, en bereidt de weg voor toepassing in de relativistische fysica.