Overlapping Schwarz Attention: Hierarchical Attention via Domain Decomposition

We propose a hierarchical attention mechanism based on two-level overlapping Schwarz domain decomposition. The method is motivated by domain decomposition methods in partial differential equations which combine local subdomain corrections with a coarse level that communicates global, long-range information. We test its usefulness in the context of finite-dimensional operator learning using a simple, one-dimensional diffusion problem. Although elementary, this problem provides a controlled sequence-to-sequence setting in which the exact nonlocal solution operator is known. After discretization, learning the solution operator amounts to approximating the inverse of a symmetric positive definite matrix. As a baseline, we use a global softmax-free low-rank attention operator of the form (QKT). The proposed construction replaces this global factorization by a two-level additive structure: local low-rank attention blocks on overlapping subdomains are combined with a coarse attention block. The resulting operator has the form Mθ-1 = Phi Q0 K0T PhiT + sumi=1N RiT Di1/2 Qi KiT Di1/2 Ri . Here, Ri restricts to an overlapping subdomain, Di is a partition-of-unity weight, and Phi is a coarse interpolation %(or prolongation) matrix. Numerical experiments for synthetic Fourier right-hand sides indicate that the domain-decomposition attention operator can converge faster and can give more accurate approximations than a global low-rank attention baseline while using significantly fewer parameters.