Location
MC 6460
Candidate
Yanming Kang| Applied Mathematics, University of Waterloo
Title
Multilevel and Multimodal Transformers for Text, Vision, and Scientific Discovery
Abstract
This thesis presents two projects on Transformer architectures and their use in text, vision, and scientific discovery. The first develops efficient multilevel attention for text and vision. The second develops Transformer-based symbolic regression together with a method for selecting among candidate equations.
The first project focuses on the computational cost of self attention in Transformers. Standard self attention compares all pairs of tokens, which gives a global receptive field but scales quadratically with sequence length or image resolution. We propose an efficient attention mechanism, called Fast Multipole Attention (FMA), that adapts ideas from the Fast Multipole Method, a classical algorithm originally developed for large $n$-body simulations. The main idea of our approach is that nearby tokens interact at full resolution, while distant tokens interact through learned summaries arranged in a multilevel hierarchy. In one spatial dimension, the construction yields an efficient attention mechanism for autoregressive and bidirectional language modeling. In two spatial dimensions, a quadtree variant extends the same idea to image classification and semantic segmentation. The resulting models preserve global interaction while reducing attention cost to nearly linear or linear scaling with respect to the input size. While other multilevel and efficient Transformers have appeared in the literature, our approach is attractive because it learns summaries on multiple levels, is parsimonious and dimension independent in principle, and obtains strong empirical results.
The second project studies symbolic equation recovery and model selection. We propose MOSAIC-SR, a hybrid framework in which a pretrained Transformer proposes multiple symbolic sketches that initialize searches in promising regions of the expression space. Monte Carlo tree search maps decoded variable tokens to observed inputs, after which parallel local searches jointly repair symbolic structure and refit constants using scale-aware optimization. The method targets algebraic recovery rather than predictive fit alone. On SRSD-Feynman data with and without irrelevant dummy variables and on six additional benchmarks, MOSAIC-SR achieves the highest symbolic solution rate on every dataset while ranking among the top two methods in predictive accuracy. As a related model-selection method, we propose RIMS, a post hoc residual-invariance method for selecting among candidate symbolic equations in multi-environment data. RIMS addresses cases in which a low error expression depends on descendants of the target and fails to represent the underlying mechanism. It compares residual distributions across environments using a normalized Cram\'{e}r-von Mises discrepancy. RIMS requires no retraining and applies to candidate pools from different symbolic regression systems. Theoretical results give a finite-candidate recovery guarantee under uniform estimation error and a detectable invariance gap, and characterize invariant downstream channels as an identification boundary. On the $119$ official AI Feynman datasets augmented with a planted measurement variable, default PySR selection recovered no physical mechanisms, whereas RIMS recovered $73$ of the $81$ mechanisms available in the union candidate pools.