Transformers meet Scientific Computing

The transformer architecture forms the backbone of modern machine learning and drives key decisions behind scaling up the broader AI economy. Transformers have been a focus of my group's work, seeking efficiency in inference and training, and looking to elements of that architecture as tools to use in other problems.
I will begin with the architecture and where its cost actually sits: attention quadratic in sequence length against feed-forward blocks quadratic in width, and how that balance shifts between training and inference. I then describe two results that attack these costs from opposite ends. First, in joint work with Armin Gerami, a linear attention mechanism we propose matches quadratic softmax attention in CLIP-style multimodal pretraining and obeys a similar scaling law, at far lower cost for long sequences. In a second work, with Ryan Synk and Monte Hoover, we look at speeding up inference, storing keys and values in a nearest-neighbor index and retrieving only the few that carry most of the attention mass. This allows us to bring million-token contexts to a single commodity GPU.
I next move to joint work with Meenakshi Krishnan and Pranav Pulijala. Here we look at the solution of partial-differential-equation based forward and inverse problems, using an architecture that borrows elements of the transformer in novel ways. Our approach, GAIA (the Geometry-Adaptive Integral Autoencoder), conditions integral-operator kernels, via cross-attention, on tokens that encode the domain boundary and the interior field distribution, and can be trained on the results of many forward solves. At inference time a single architecture solves forward, boundary value and inverse PDE problems on arbitrary domains in one pass, with no retraining per instance and no iterative optimization. It gives the best reported accuracy on every inverse and boundary value task we tried, including electrical impedance tomography, optical tomography, and reconstruction of transonic airfoil flow from sparse noisy measurements, and it holds that accuracy under changes of discretization where transformer baselines lose an order of magnitude. Along the way we release new varying-geometry benchmarks for inverse problems and boundary value problems.