Collaborative Research: MFAI: Mathematical Foundations of Structure-Aware Adaptive Machine Learning Algorithms

Recent years have witnessed rapid progress in computational algorithms for large-scale machine learning, but many advances remain segmented and disconnected. This project develops a unified mathematical framework for modern machine learning algorithms with emphasis on scalability, adaptivity, and structure-aware optimization and generalization. This framework will deepen understanding of the limitations of current machine learning algorithms and guide the design of more advanced methods with better performance, scalability, and generalization power. The outcomes of this project will broaden access to large-scale artificial intelligence (AI) by developing theory-guided methods for training under resource constraints, lowering barriers for researchers with limited computational budgets. The proposed research will be integrated into graduate-level courses on optimization and machine learning, providing new materials for educating the next generation of AI researchers. Software implementations will be released as open-source tools to facilitate adoption in academia, industry, and government labs. By linking efficiency, robustness, and generalization, the project will help make large-scale AI training more reliable, reproducible, and widely accessible. Large-scale deep learning relies heavily on first-order algorithms such as Stochastic Gradient Descent (SGD), Adaptive Moment Estimation (Adam), and their variants, yet their remarkable empirical success is still only partially understood. We aim to bridge this gap by systematically studying how model structures (e.g., sparsity, low-rankness, manifold constraints), data structures (e.g., anisotropic gradient noise), and neural architectural structures (e.g., scaling laws, wide-network limits) affect optimization dynamics and generalization. Our analysis will leverage adaptive gradient methods viewed as extensions of mirror descent, nonlinear preconditioning, and Riemannian optimization. The project will pursue three thrusts: (i) structure-aware adaptive optimization techniques, (ii) adaptive learning rate schedules, and (iii) structure-dependent generalization. These thrusts will advance the mathematical foundations of machine learning by (i) establishing unified, structure-dependent convergence guarantees for adaptive algorithms such as Adam, Shampoo, and Muon, (ii) deriving scaling laws that remain valid under structural constraints and adaptive updates, and (iii) developing generalization bounds that depend explicitly on spectral properties of gradient noise. Novel contributions include a framework for analyzing model structures with time-varying potentials, effective learning rate schedules in stochastic optimization beyond isotropic noise models, and dimension-free generalization bounds under anisotropic gradients. Together, these results will provide principled mathematical theory and provably efficient algorithms that achieve both computational scalability and robust generalization, with direct implications for training large language models and other modern AI systems. 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.