Preconditioning and Numerical Stability in Neural Network Training for Parametric PDEs

Technische Samenvatting: Preconditionering en Numerieke Stabiliteit in de Training van Neurale Netwerken voor Parametrische PDE's

Probleemformulering
Het artikel behandelt de uitdaging om efficiënt oplossingen te benaderen voor parameterafhankelijke partiële differentiaalvergelijkingen (PDE's) met behulp van neurale netwerken. Specifiek richt het zich op "hybride" benaderingen waarbij de afhankelijkheid van ruimtelijke variabelen $x$ wordt afgehandeld door klassieke discretisatiemethoden (bijv. eindige elementen), terwijl de afhankelijkheid van parametrische variabelen $y$ wordt gemodelleerd door neurale netwerken. De oplossing wordt benaderd als $u(x, y) \approx \sum u_i(y; \theta)\phi_i(x)$, waarbij de coëfficiënten $u_i$ de outputs zijn van een neuraal netwerk.

De kern van de moeilijkheid ligt in de optimalisatie van de netwerkparameters $\theta$. Of het nu gaat om het minimaliseren van het empirische risico (het direct benaderen van de oplossing) of het minimaliseren van een residu-loss (variationele correctheid), de objectieve functies bevatten kwadratische vormen afgeleid van differentiaaloperatoren. Standaard eindige elementen-bases leiden vaak tot slecht geconditioneerde Gramiaanse matrices, wat de prestaties van gradiëntgebaseerde optimalisatieschema's ernstig verslechtert. Bovendien identificeren de auteurs een kritiek, vaak over het hoofd gezien probleem: zelfs wanneer preconditionering het conditiegetal van het probleem verbetert, kunnen standaardimplementaties van preconditioners nog steeds een catastrofaal verlies van numerieke precisie veroorzaken, met name bij het gebruik van lage-precisie rekenkunde (single- of half-precision floating-point getallen).

Methodologie
De auteurs stellen een tweeledige aanpak voor om slecht conditiering en numerieke instabiliteit aan te pakken:

1. Goed Geconditioneerde Frame-representaties: In plaats van standaard eindige elementen-bases maken de auteurs gebruik van multilevel frame-representaties (specifiek, hiërarchische bases afgeleid van additive Schwarz-methoden). Door een frame $\Phi$ te gebruiken dat de Riesz-basis eigenschap voldoet (of een overcompleet frame is met begrensde frame-constanten), worden de resulterende discretisatiematrices goed geconditioneerd. Het artikel maakt gebruik van First-Order System Least Squares (FOSLS) formuleringen, die directe foutkwantificering en variationele correctheid mogelijk maken.

2. Stabiele Matrix-representaties: Het artikel demonstreert dat het louter preconditioneren van de systeemmatrix (bijv. het vormen van $H^\top A_y H$) onvoldoende is voor numerieke stabiliteit in lage precisie. Het product van slecht geconditioneerde factoren (de preconditioner $H$ en de stijfheidsmatrix $A_y$) kan nog steeds leiden tot precisieverlies. Om dit op te lossen, stellen de auteurs een "stabiele representatie" van de geconditioneerde operator voor. Ze factoriseren de operator in $D^\top C_y D$, waarbij:

   • $D$ de frame-coëfficiënten naar een ruimte van stuksgewijze polynomen mapt (via differentiaaloperatoren).
   • $C_y$ een diagonaal of blokdiagonaal matrix is die de inwendige product in de beeldruimte vertegenwoordigt.
   • Cruciaal is dat zowel $D$ als $C_y$ zo zijn geconstrueerd dat ze uniform goed geconditioneerd zijn, onafhankelijk van het discretisatieniveau. Deze decompositie vermijdt de vermenigvuldiging van slecht geconditioneerde matrices, waardoor de precisie behouden blijft, zelfs in half-precision (float16) rekenkunde.

De auteurs implementeren deze concepten met behulp van een matrix-vrije numerieke realisatie, waarbij matrix-vector producten worden geëvalueerd via tensorcontracties op lokale elementen in plaats van het assembleren van globale ijle (sparse) matrices. Ze testen drie neurale netwerkarchitecturen: een Full Low-Rank ResNet, een Separate Low-Rank ResNet (één netwerk per frame-niveau), en een Separate Frame Representation (met verlengingsoperatoren tussen niveaus).

Belangrijkste Resultaten
Numerieke experimenten op een 1D parametrische elliptische vergelijking leveren de volgende bevindingen op:

• Impact van Preconditionering: Preconditionering verbetert de convergentie van optimalisatieschema's aanzienlijk. Voor de Adam-optimizer reduceert preconditionering de loss met ongeveer vier ordes van grootte vergeleken met ongeconditioneerde training.
• Gevoeligheid van de Optimizer: De Adam-optimizer profiteert het meest van preconditionering en bereikt een hoge nauwkeurigheid in zowel single- als double-precision. In tegen tegenstelling hiertoe vertoont Stochastic Gradient Descent (SGD) slechts beperkte verbetering door louter preconditionering. Tweede-orde methoden zoals L-BFGS en Natural Gradient Descent (NGD) presteren goed in double precision, maar zijn zeer gevoelig voor numerieke precisie; hun prestaties verslechteren aanzienlijk in half-precision zonder de stabiele representatie.
• Numerieke Stabiliteit: De voorgestelde "stabiele representatie" ($D^\top C_y D$) maakt training mogelijk met half-precision (float16) rekenkunde zonder verlies van nauwkeurigheid. In half-precision bereikt de Adam-optimizer met stabiele preconditionering een loss-waarde die dicht bij de machine-precisie limiet ($10^{-4}$) ligt, wat overeenkomt met de prestaties van double-precision training. Standaard preconditionering ($H^\top A_y H$) faalt in het bereiken van dit niveau in half-precision vanwege precisieverlies tijdens de matrix-vector vermenigvuldiging.
• Architectuurvergelijking: De drie geteste netwerkarchitecturen (Full ResNet, Separate ResNet, Separate Frame Representation) vertonen vergelijkbare prestaties over verschillende precisieniveaus, mits het aantal trainbare parameters vergelijkbaar is. De eenvoudigere Full ResNet-architectuur presteerde competitief met de complexere aparte representaties.

Betekenis en Claims
Het artikel claimt dat hoewel preconditionering een bekende techniek is voor het versnellen van convergentie, de standaard toepassing ervan onvoldoende is om numerieke stabiliteit te garanderen in de training van neurale netwerken, met name bij het benutten van lage-precisie hardware. De primaire bijdrage is de demonstratie dat een specifieke factorisatie van de geconditioneerde operator — het vermijden van het expliciete product van slecht geconditioneerde matrices — noodzakelijk is om de volledige precisie in lage-precisie rekenkunde te handhaven.

De auteurs stellen dat deze aanpak het mogelijk maakt om half-precision floating-point getallen te gebruiken bij de training van neurale netwerken voor parametrische PDE's zonder de nauwkeurigheid van de oplossing op te offeren. Dit is significant voor computationele efficiëntie, aangezien half-precision rekenkunde sneller is en minder geheugen vereist, wat het trainen van grotere modellen of complexere problemen op moderne hardware mogelijk maakt. Het werk overbrugt de kloof tussen operator-preconditioneringstheorie (specifiek Full Operator Preconditioning concepten) en praktische training van neurale netwerken, door aan te tonen dat stabiele representaties bereikt kunnen worden zonder de continue operatorvergelijking te herformuleren, maar eerder door een niet-standaard factorisatie van de discrete matrixrepresentatie.