Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

Technische Samenvatting: Versnelling van Natuurlijke Gradiëntafstijging voor PINNs met Geautomatiseerde Numerieke Lineaire Algebra

Probleemstelling
Physics-Informed Neural Networks (PINNs) en aanverwante op neurale netwerken gebaseerde oplossers voor Partiële Differentiaalvergelijkingen (PDE's) hebben veelbelovend bewezen meshvrije benaderingen te kunnen bieden, met name in hoge dimensies. Hun praktische toepassing wordt echter belemmerd door optimalisatieproblemen. Standaard eerste-orde optimalisatoren zoals Adam vertonen vaak stagnatie bij oplossingen met lage nauwkeurigheid als gevolg van de slecht geconditioneerde verlieslandschap. Hoewel Natuurlijke Gradiëntafstijging (NGD) is opgekomen als een robuust alternatief van de tweede orde dat rekening houdt met de geometrie van de functieruimte, is de toepassing ervan op grote neurale netwerken computatief onhaalbaar.

De primaire bottleneck in NGD is de noodzaak om lineaire systemen op te lossen die de Gramiaan-matrix (of Fisher-informatiematrix), $G(\theta)$, bevatten. Directe inversie of factorisatie van $G(\theta)$ brengt een kubieke kost $O(p^3)$ met zich mee met betrekking tot het aantal parameters $p$, waardoor het onhaalbaar wordt voor grote netwerken. Matrixvrije benaderingen die de Conjugate Gradient (CG)-methode gebruiken, vermijden expliciete matrixassemblage, maar lijden onder trage convergentie wanneer $G(\theta)$ slecht geconditioneerd is, een veelvoorkomend verschijnsel bij PDE-discretisaties.

Methodologie
De auteurs stellen een modulair kader voor dat matrixvrije NGD versnelt door technieken uit Geautomatiseerde Numerieke Lineaire Algebra (RandNLA) te integreren om effectieve preconditioners te construeren voor de interne CG-oplosser.

1. Generaliseerde Matrixvrije NGD: Het werk breidt matrixvrije NGD uit tot een bredere klasse van metrieken die relevant zijn voor PDE's. Door gebruik te maken van automatische differentiatie en een specifieke aanname over de structuur van de metriek (op integralen gebaseerde vormen), tonen de auteurs aan dat matrix-vectorproducten (matvecs) met de Gramiaan efficiënt kunnen worden berekend zonder de volledige matrix te assembleren. Dit geldt zowel voor sterke formuleringen (PINNs) als voor zwakke formuleringen (Finite Element Interpolated Neural Networks, of FEINNs).

2. RandNLA-Preconditionering: Om de slechte conditionering van $G(\theta)$ aan te pakken, introduceert het artikel twee nieuwe preconditioneringsstrategieën gebaseerd op de observatie dat de Gramiaan een sterke spectrale afname vertoont (snelle afname van singuliere waarden):

   • NyströmNGD: Gebruikt de gerandomiseerde Nyström-benadering om een preconditioner met lage rang te construeren. Deze methode benadert de Gramiaan met een random testmatrix en berekent een preconditioner die gebruikmaakt van de dominante eigenruimte.
   • RPCholNGD: Past de Randomly Pivoted Partial Cholesky (RPCholesky)-decompositie toe. Deze adaptieve methode selecteert pivots op basis van de residudiagonaal en construeert een benadering met lage rang $FF^\top$ die dient als basis voor de preconditioner.

3. Algoritmisch Kader: De auteurs presenteren een meta-algoritme dat de implementatie van NGD ontkoppelt van de constructie van de preconditioner. Dit maakt flexibele combinaties mogelijk van directe oplossers, matrixvrije iteratieve oplossers en diverse preconditioneringsstrategieën. Het kader bevat adaptieve heuristieken voor het selecteren van de regularisatieparameter $\mu$ en de rang $\ell$ van de preconditioner, waarmee numerieke stabiliteit wordt gewaarborgd en de computatiekosten worden gebalanceerd met de kwaliteit van de preconditioner.

Belangrijkste Bijdragen

• Generalisatie van Matrixvrije NGD: Het artikel biedt een methodologie om Gramiaan-matvecs via automatische differentiatie te berekenen voor een breed scala aan PDE-gerelateerde metrieken, en generaliseert hiermee voorbij eerdere specifieke implementaties.
• Nieuwe Gepreconditioneerde Algoritmen: De introductie van NyströmNGD en RPCholNGD, die matrixvrije NGD respectievelijk combineren met gerandomiseerde Nyström- en RPCholesky-preconditionering.
• Schaalbaar Kader: Een modulaire architectuur die het geheugengebruik reduceert van $O(p^2)$ (vereist door geheugenvolle quasi-Newton-methoden) naar $O(p\ell)$, waarbij $\ell$ een instelbare parameter met lage rang is.
• Efficiëntie: De voorgestelde methoden transformeren de computatiechallenge van slechte conditionering in een voordeel door de structuur met lage rang van de Gramiaan te exploiteren, waardoor efficiënte preconditionering mogelijk wordt zonder expliciete matrixassemblage.

Resultaten
Uitgebreide numerieke experimenten werden uitgevoerd op vijf PDE-problemen, waaronder de 3D Poisson-vergelijking, de 3+1D Warmtevergelijking, Kovasznay-stroming en 2D diffusie-reactie-transportproblemen, waarbij zowel PINNs als FEINNs werden gebruikt.

• Prestaties versus Baselines: De gepreconditioneerde varianten van matrixvrije NGD (NyströmNGD en RPCholNGD) bereikten uiteindelijke nauwkeurigheden die vergelijkbaar waren met NGD met een directe oplosser en de state-of-the-art SSBroyden-optimalisator.
• Computatie-efficiëntie: De voorgestelde methoden presteerden aanzienlijk beter dan ongepreconditioneerde matrixvrije NGD, die vanwege slechte CG-convergentie geen hoge nauwkeurigheid bereikte. In vergelijking met SSBroyden en andere optimalisatoren (Adam, BFGS, L-BFGS) toonden NyströmNGD en RPCholNGD op de meeste benchmarks superieure of vergelijkbare nauwkeurigheid met aanzienlijk minder wandkloktijd.
• Geheugen en Schaalbaarheid: De methoden behielden een geheugengebruik van $O(p\ell)$, waardoor ze konden schalen naar grotere netwerken waar volledige-matrixmethoden of geheugenvolle quasi-Newton-methoden onhaalbaar zouden zijn.
• Rang-Adaptiviteit: De experimenten toonden aan dat de effectieve rang die voor preconditionering vereist is ($\ell$) vaak veel kleiner was dan het totale aantal parameters, wat de sterke spectrale afname van de Gramiaan in deze settings bevestigt.

Betekenis en Beweringen
Het artikel beweert dat door Natuurlijke Gradiëntafstijging te verbinden met Geautomatiseerde Numerieke Lineaire Algebra, de computatiebottlenecks die het praktische gebruik van NGD in PDE-oplossers hebben beperkt, kunnen worden overwonnen. De auteurs stellen dat hun gepreconditioneerde matrixvrije aanpak een superieure afweging biedt tussen nauwkeurigheid en computatietijd in vergelijking met bestaande state-of-the-art optimalisatoren voor PINNs.

Specifiek benadrukt het werk dat:

• Preconditionering cruciaal is voor matrixvrije NGD; zonder deze is de methode inefficiënt.
• De structuur met lage rang van de Gramiaan, die vaak empirisch wordt waargenomen, systematisch kan worden benut via RandNLA om efficiënte preconditioners te creëren.
• Het voorgestelde kader bijzonder goed geschikt is voor op neurale netwerken gebaseerde oplossers die gericht zijn op zwakke formuleringen van PDE's (zoals FEINNs), waarbij de variatiestructuur op natuurlijke wijze de selectie van de metriek leidt.
• Hoewel SSBroyden concurrerend blijft (en in specifieke gevallen zoals Kovasznay-stroming iets sneller is), bereiken de voorgestelde methoden over het algemeen vergelijkbare of betere nauwkeurigheid in minder tijd en met lagere geheugeneisen, waardoor ze een veelbelovende richting vormen voor grootschalige PDE-beperkte optimalisatie.