Mondrian: Transformer Operators via Domain Decomposition

TL;DR

Mondrian introduces a domain decomposition approach to transformer operators, enabling scalable, high-resolution PDE modeling by decoupling attention from discretization and supporting hierarchical attention mechanisms.

Contribution

The paper proposes Mondrian, a novel transformer operator that uses domain decomposition to improve scalability and resolution in neural PDE models, extending attention mechanisms hierarchically.

Findings

Achieves strong performance on Allen-Cahn and Navier-Stokes PDEs.

Supports resolution scaling without retraining.

Decouples attention from discretization for better scalability.

Abstract

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.