UC Berkeley

Foundation models operate at unprecedented scale and exhibit empirical behaviors that often depart from classical learning and optimization theory, while their deployment introduces new requirements on safety and computational efficiency. This thesis uses statistical and optimization principles to explain representative non-classical phenomena in modern foundation models and design practical algorithms that improve their efficiency and safety under realistic constraints.On the explanatory side, we first study in-context learning (ICL), the ability to solve new tasks at inference time by conditioning on a small number of demonstrations without parameter updates. In simplified yet representative transformer models trained on in-context linear regression data, we establish global convergence guarantees for gradient flow, characterize the effective function class induced at convergence, and analyze robustness under several forms of distribution shift. We further show that augmenting attention with an MLP component can provably reduce approximation error and achieve nearly Bayes-optimal risk in a Gaussian linear-regression setting, and we formalize an equivalence between such architectures and one-step gradient-based adaptation with learnable initialization.We then investigate optimization dynamics beyond the descent-lemma regime, motivated by edge-of-stability behavior observed in large-scale deep learning. Using logistic regression with linearly separable data as a canonical testbed, we analyze gradient descent with large and adaptive stepsizes and show that it reaches an arbitrarily small risk immediately after at most 1/γ2 burn-in iterations, where γ > 0 is the margin. We also prove a matching lower bound: any batch or online first-order method requires Ω(1/γ2 ) iterations to find a separator, establishing minimax optimality and matching the Perceptron algorithm even in constants. Our analysis extends to a broad class of losses and certain two-layer networks.On the algorithmic-design side, we study machine unlearning for LLMs, motivated by the memorization of sensitive, private, or copyrighted data during pre-training. Traditional unlearning methods, which are based on gradient ascent (GA) on the loss of the undesirable data, either fail to effectively unlearn the target data or suffer from catastrophic collapse—a drastic degradation of the model utility. We propose Negative Preference Optimization (NPO) and theoretically show that its progression toward catastrophic collapse is exponentially slower than GA. Empirically, NPO achieves a better trade-off between forget quality and model utility.Finally, we study efficient alignment at inference time under resource constraints. We focus on best-of-N (BoN) decoding, a simple and effective alignment method whose practical deployment is limited by its high computational cost. Motivated by hardware and resource utilization considerations, we propose Speculative Rejection, a resource-inspired inference-time alignment algorithm that leverages the correlation between partial and final rewards and dynamically allocates computation by early-stopping low-quality generations based on partial reward signals. This approach achieves reward performance comparable to BoN decoding with much larger N, while reducing required GPU resources by approximately 16 to 32 times.