Mondrian: Transformer Operators via Domain Decomposition
This addresses the problem of high computational cost in operator learning for PDEs, offering a scalable solution for researchers in scientific computing, though it is incremental as it builds on existing transformer and domain decomposition methods.
The paper tackles the challenge of scaling transformer-based operator models for partial differential equations (PDEs) by introducing Mondrian, which decomposes domains into subdomains to decouple attention from discretization, achieving strong performance on Allen-Cahn and Navier-Stokes PDEs with resolution scaling without retraining.
Operator learning enables data-driven modeling of partial differential equations (PDEs) by learning mappings between function spaces. However, scaling transformer-based operator models to high-resolution, multiscale domains remains a challenge due to the quadratic cost of attention and its coupling to discretization. We introduce \textbf{Mondrian}, transformer operators that decompose a domain into non-overlapping subdomains and apply attention over sequences of subdomain-restricted functions. Leveraging principles from domain decomposition, Mondrian decouples attention from discretization. Within each subdomain, it replaces standard layers with expressive neural operators, and attention across subdomains is computed via softmax-based inner products over functions. The formulation naturally extends to hierarchical windowed and neighborhood attention, supporting both local and global interactions. Mondrian achieves strong performance on Allen-Cahn and Navier-Stokes PDEs, demonstrating resolution scaling without retraining. These results highlight the promise of domain-decomposed attention for scalable and general-purpose neural operators.