Linear transformers for solving parametric partial differential equations

Abstract

The simulation of physical phenomena relies on solving Partial Differential Equations (PDEs), and Machine Learning models have increasingly addressed this task in recent years. PDEs often involve parameters influencing their evolution, prompting the development of models that consider these parameters as additional input. These parameter-conditioned models aim to generalize across different PDE parameters, replacing the need for multiple models trained on specific ones. Transformer models have been achieving great success in Natural Language Processing (NLP), Speech Processing, and even in domains such as Computer Vision. Due to their ability to effectively model long-range dependencies in sequential data, their field of application is steadily increasing. Calculating attention via Scaled Dot-Product Attention in Vanilla Transformers is computationally expensive and scales quadratically with the input length. This leads to a bottleneck for very long sequences. To address this challenge, Linear Transformers have been introduced, substituting the Scaled Dot-Product Attention to achieve linear time and space complexity. Consequently, Linear Transformers have shown promising potential for processing very long sequences efficiently. We investigate two approaches of utilizing Linear Transformers for solving PDEs and their associated problems. Moreover, we conduct a comprehensive comparison between our proposed transformer-based models and state-of-the-art models for solving parametric PDEs. The evaluation criteria include accuracy for short and long rollouts, memory consumption, and inference times. The results demonstrate that our proposed models perform competitively with the current state-of-the-art models, providing an efficient solution for PDE solving. Die Simulation physikalischer Phänomene beruht auf dem Lösen von partiellen Differentialgleichungen (engl. Partial Differential Equations, kurz PDEs). In den letzten Jahren wurden zunehmend Machine Learning Modelle entwickelt, die partielle Differentialgleichungen lösen. PDEs beinhalten oft Parameter, die ihre Lösungen beeinflussen, was die Entwicklung von Machine Learning Modellen, die diese Parameter als zusätzlichen Input berücksichtigen, vorangetrieben hat. Diese parameterkonditionierten Modelle zielen darauf ab, über verschiedene PDE-Parameter hinweg zu generalisieren, sodass nicht mehr mehrere Modelle benötigt werden, die für bestimmte Parameter trainiert wurden. Transformer-Modelle haben sich in der Verarbeitung natürlicher Sprache (engl. Natural Language Processing, kurz NLP), der Verarbeitung von Audiosequenzen und sogar in Bereichen wie der Computer Vision als sehr erfolgreich erwiesen. Aufgrund ihrer Fähigkeit, Abhängigkeiten über große Distanzen hinweg in sequenziellen Daten effektiv zu modellieren, wird das Anwendungsgebiet von Transformer zunehmend größer. Die Berechnung von Attention mittels Scaled Dot-Product Attention in Vanilla Transformer ist rechenintensiv und skaliert quadratisch mit der Eingabelänge. Dies führt bei sehr langen Sequenzen zu einem Engpass. Um dieses Problem zu lösen, wurden Linear Transformer eingeführt, die die Berechnung von Scaled Dot-Product Attention ersetzen, um eine lineare Zeit- und Speicherkomplexität zu erreichen. Folglich haben Linear Transformer ein großes Potential für die effiziente Verarbeitung sehr langer Sequenzen. Wir untersuchen zwei Ansätze zur Verwendung von Linear Transformer für das Lösen von PDEs und den damit verbundenen Problemen. Darüber hinaus führen wir einen umfassenden Vergleich zwischen den von uns entwickelten Transformer-Modellen und anderen Modellen zur Lösung parametrischer PDEs durch. Zu den Vergleichskriterien gehören die Genauigkeit für kurze und lange Rollouts, der Speicherverbrauch und die Inferenzzeiten. Die Ergebnisse zeigen, dass die von uns entwickelten Modelle mit aktuellen Modellen konkurrieren können und eine effiziente Lösung für das Lösen von PDEs bieten.