MFAI: Any-Dimensional Equivariant Learning

Modern artificial intelligence (AI) systems provide a versatile toolkit for identifying solution maps for a variety of problems in a data-driven fashion -- some prominent application domains include imaging, signal processing, recommendation systems, and partial differential equations. In this paradigm, one is given a training dataset consisting of input-output -- inputs representing problem instances and outputs representing solutions -- and the objective is to learn a solution map that can be applied to future test inputs. For example, in image classification, the inputs correspond to images, and the outputs correspond to labels indicating the presence or absence of a person. While remarkably successful, a major drawback of this paradigm concerns the size or dimension of the training examples. Specifically, one can expect the size of the input problem instances in the training set to be different from that of the problem instances encountered at test time; indeed, even expecting all the data in the training set to lie in the same dimension is often unreasonable. For example, when learning image classifiers, contemporary methods usually require the training set to consist of images of the same size. Moreover, the model trained cannot handle larger or smaller images. A more natural learning objective is to identify an algorithm instead -- a sequence of solution maps that gracefully handles inputs of increasingly larger size. This project advances a research program to learn such an algorithm from training data consisting of inputs of different sizes, allowing for learned algorithms to be deployed on inputs of any size (including those not present in the training set). To achieve this goal, the project utilizes a framework based on representation theory that facilitates the specification of parameterized families of sequences of solution maps, i.e., parameterized families of algorithms, with the parameters given by invariants in a sequence of group representations. The learning problem identifies the best element from such a parameterized family to minimize training error. At test time, one can simply apply the solution map that takes in the dimension of the test instance. Central to this development is the recently identified phenomenon of representation stability in the algebraic topology literature, which provides the mathematical underpinning of the framework. The project tackles approximation-theoretic, statistical, and computational research questions that arise from our framework, along with high-impact applications to imagining inverse problems and protein folding. The framework will provide a mathematical basis for training AI systems on low-dimensional training data while still being applicable to high-dimensional test data, thus offering the possibility of significant energy savings and making ML systems more sustainable. Furthermore, the research is integrated with a comprehensive educational component. The principal investigators will mentor high school students through the Johns Hopkins Center for Educational Outreach and undergraduates through the Caltech Student-Faculty Programs Office. This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.