Implicit differentiation of tensor network algorithms

Technische Samenvatting: Impliciete Differentiatie van Tensor Netwerk Algoritmen

Probleemstelling
De variationele optimalisatie van Projected Entangled-Pair States (PEPS) voor grondtoestandberekeningen is sterk afhankelijk van de nauwkeurige berekening van energiegradiënten met betrekking tot lokale variationele parameters. De huidige state-of-the-art aanpak maakt gebruik van Automatische Differentiatie (AD) om deze gradiënten te berekenen. Deze methode staat echter voor drie belangrijke uitdagingen:

1. Efficiëntie: Standaard reverse-mode AD vereist het terugpropageren door elke iteratie van het contractiealgoritme (bijv. CTMRG of Boundary MPS). Zelfs met fixed-point formuleringen die het opslaan van de volledige computationele graaf vermijden, creëren de aanwezigheid van iteratieve subroutines (zoals ijlere eigensolvers of SVD's) een geneste iteratieve procedure voor gradiëntberekening, wat leidt tot hoge computationele kosten.
2. Stabiliteit: Het differentiëren van de primitieve subroutines die worden gebruikt in contractiealgoritmen (bijv. getrunceerde eigenwaardedecompositie of SVD's) introduceert numerieke instabiliteiten. Deze instabiliteiten ontstaan vaak bij bijna-degeneratieve eigenwaarden of singuliere waarden, waarbij de afgeleiden van deze subroutines divergeren. Cruciaal is dat deze instabiliteiten artefacten zijn van de differentiatieprocedure; de fysieke energiegradiënt zelf blijft stabiel ten opzichte van de onderliggende variaties.
3. Implementatiecomplexiteit: Het toepassen van AD op de optimalisatie van PEPS vereist dat elke primitieve operatie in het contractiealgoritme compatibel is met de AD-engine. Het definiëren van aangepaste differentieregels voor niet-native subroutines is arbeidsintensief en beperkt de flexibiliteit.

Methodologie
De auteurs stellen een raamwerk voor op basis van impliciete differentiatie om deze uitdagingen aan te pakken. In plaats van door de iteratieve stappen van het contractiealgoritme te differentiëren, herformuleert de methode de geconvergeerde contractie-omgeving als de wortel van een enkele set algebraïsche karakteristieke vergelijkingen.

De kernstappen van de methodologie zijn:

1. Algebraïsche Karakterisering: De geconvergeerde omgeving (bijv. corner en edge tensoren in CTMRG, of boundary MPS tensoren) wordt niet gekarakteriseerd door de iteratieve update-regel $x = f(x, p)$, maar door een set vergelijkingen $F(y, p) = 0$. Hierbij is $y$ een set van omgeving variabelen (inclusief intermediaire variabelen zoals isometrieën of singuliere waarden) die impliciet afhangen van de variationele parameters $p$.
2. Impliciete Differentiatie: Door de impliciete functietheorema toe te passen op $F(y, p) = 0$, reduceert de gradiëntberekening tot het oplossen van een enkele, niet-geneste systeem van lineaire vergelijkingen:
   $$\tilde{y} \partial_y F = -\tilde{y} \partial_p F$$
   waarbij $\tilde{y}$ de adjoint-variabele is. Dit elimineert de noodzaak om door de iteratieve subroutines te terugpropageren.
3. Stabiele Parametrisatie: Een belangrijke innovatie is de keuze van de variabele set $y$ en de parametrisatie van de karakteristieke vergelijkingen. Door gebruik te maken van de intrinsieke gauge-symmetrieën van de contractie-omgeving, construeren de auteurs parametrisaties die expliciet de vrijheidsgraden verwijderen die verantwoordelijk zijn voor numerieke instabiliteiten (bijv. variaties binnen degeneratieve subruimten). Bijvoorbeeld, in CTMRG worden de variaties van de isometrieën beperkt tot hun nulruimten, waardoor de divergerende termen geassocieerd met het differentiëren van getrunceerde eigensolvers worden omzeild.
4. Modulaire Implementatie: De aanpak ontkoppelt de contractie-algoritme van de gradiëntberekening. De forward pass blijft ongewijzigd (gebruikmakend van standaard contractie-routines), terwijl de reverse pass alleen de algebraïsche karakteristieke vergelijkingen vereist te differentiëren, die bestaan uit eenvoudige tensorcontracties.

Belangrijkste Bijdragen
Het artikel biedt expliciete formuleringen van karakteristieke vergelijkingen voor drie belangrijke PEPS-contractieschema's:

• $C_{4v}$ Symmetrische CTMRG: De auteurs leiden een set karakteristieke vergelijkingen af die corner tensoren, edge tensoren en een geparametriseerde isometrie bevatten. Deze formulering elimineert de noodzaak om de getrunceerde eigenwaardedecompositie (eigh) te differentiëren en lost instabiliteiten op die geassocieerd worden met degeneratieve eigenwaarden.
• Boundary MPS (VUMPS): Een set karakteristieke vergelijkingen wordt geformuleerd voor de top en bottom boundary MPS eigenvectoren en gemengde omgevingen. Deze aanpak voorkomt de noodzaak om de resulterende boundary MPS tensoren tijdens de gradiëntberekening te gauge-fixen en elimineert geneste lineaire problemen die voortkomen uit het differentiëren van de dominante eigensolvers.
• Asymmetrische CTMRG (Zonder Symmetrieën): De methode wordt uitgebreid naar algemene systemen zonder ruimtelijke symmetrieën. De auteurs introduceren gemodificeerde corner en edge tensoren en aangepaste "wortelkwadraat"-variabelen voor singuliere waarden om gauge-invariantie en stabiliteit te waarborgen bij het differentiëren van getrunceerde SVD's.

Resultaten
De auteurs hebben de impliciete differentiatie-aanpak getest tegenover fixed-point differentiatie en directe reverse-mode AD met behulp van de PEPSKit.jl en MPSKit.jl pakketten:

• Nauwkeurigheid: De via impliciete differentiatie berekende gradiënten komen overeen met referentie eindige-verschil waarden en black-box reverse-mode AD binnen de tolerantie van de lineaire solver, zelfs voor matig geconvergeerde omgevingen.
• Prestaties:
  • Voor $C_{4v}$ CTMRG en Boundary MPS presteert de impliciete aanpak consistent beter dan fixed-point differentiatie, waarbij de versnelling toeneemt naarmate de bond dimensie ($D$) en de omgeving dimensie ($\chi$) groeien.
  • Voor Asymmetrische CTMRG op niet-triviale eenheidscellen vertoont de impliciete methode een superieure asymptotische schaling. Hoewel fixed-point differentiatie bij lage dimensies sneller kan zijn vanwege implementatie-overhead, wordt de impliciete methode aanzienlijk efficiënter bij hoge dimensies waar ijlere matrix-decompositie vereist is.
• Schaling: In veel gevallen wordt de kosten van de gradiëntberekening met impliciete differentiatie een sub-leidende stap vergeleken met de contractiekosten zelf, terwijl fixed-point differentiatie vaak een dominante bottleneck blijft.

Betekenis en Claims
Het artikel beweert dat dit raamwerk een paradigmaverschuiving vertegenwoordigt in tensor netwerk programmeren. Door algebraïsche optimaliteitscondities te differentiëren in plaats van de algoritmische subroutines, maakt de methode het mogelijk om:

• Implementatie te vereenvoudigen: Het verwijdert de vereiste dat elke subroutine in het contractiealgoritme differentieerbaar moet zijn, waardoor standaard, hoog geoptimaliseerde lineaire algebra-routines in de forward pass kunnen worden gebruikt.
• Stabiliteit te verbeteren: Het verwijdert op natuurlijke wijze numerieke instabiliteiten die voortkomen uit de differentiatie van subroutines door gebruik te maken van gauge-symmetrieën om stabiele parametrisaties te construeren.
• Efficiëntie te verhogen: Het vermindert de computationele complexiteit van gradiëntevaluatie door de lineaire problemen te ont-nesten, wat leidt tot betere schaling met systeemgrootte.

De auteurs benadrukken dat hoewel zij de aanpak op drie specifieke contractie-schema's hebben gedemonstreerd, het kernalgoritme (Algoritme 1 en 2) algemeen is en toepasbaar op elk optimalisatieprobleem waarbij de oplossing gedefinieerd kan worden door differentieerbare optimaliteitscondities. Zij beweren niet dat ze alle PEPS-uitdagingen hebben opgelost, maar stellen dat deze aanpak stabiele, gradiënt-gebaseerde PEPS-optimalisatie aanzienlijk praktischer en robuuster maakt.