Zak Mhammedi
Zak is a research scientist at the Learning Theory Team at Google, specializing in the design of efficient and practical reinforcement learning (RL) algorithms and optimization methods for large neural networks. His research focuses on two main areas: on the one hand, developing RL algorithms that are not only implementable and deployable but also provably statistically efficient; on the other hand, designing fast algorithms for both convex and non-convex optimization problems, with applications to neural network training.
Research Areas
Authored Publications
Sort By
Preview abstract
We study online linear optimization with matrix variables constrained by the operator norm, a setting where the geometry renders designing data-dependent and efficient adaptive algorithms challenging. The best-known adaptive regret bounds are achieved by Shampoo-like methods, but they require solving a costly quadratic projection subproblem. To address this, we extend the gradient-based prediction scheme to adaptive matrix online learning and cast algorithm design as constructing a family of smoothed potentials for the nuclear norm. We define a notion of admissibility for such smoothings and prove any admissible smoothing yields a regret bound matching the best-known guarantees of one-sided Shampoo. We instantiate this framework with two efficient methods that avoid quadratic projections. The first is an adaptive Follow-the-Perturbed-Leader (FTPL) method using Gaussian stochastic smoothing. The second is Follow-the-Augmented-Matrix-Leader (FAML), which uses a deterministic hyperbolic smoothing in an augmented matrix space. By analyzing the admissibility of these smoothings, we show both methods admit closed-form updates and match one-sided Shampoo’s regret up to a constant factor, while significantly reducing computational cost. Lastly, using the online-to nonconvex conversion, we derive two matrix-based optimizers, Pion (from FTPL) and Leon (from FAML). We prove convergence guarantees for these methods in nonsmooth nonconvex settings, a guarantee that the popular Muon optimizer lacks.
View details
Preview abstract
Simulators are a pervasive tool in reinforcement learning, but most existing algorithms cannot efficiently exploit simulator access---particularly in high-dimensional domains that require general function approximation. We explore the power of simulators through online reinforcement learning with local simulator access (also known as local planning), an RL protocol where the agent is allowed to reset to previously observed states and follow their dynamics during training. We use local simulator access to unlock new statistical guarantees that were previously out of reach:
I) We show that MDPs with low coverability Xie et. al 2023---a general structural condition that subsumes Block MDPs and Low-Rank MDPs---can be learned in a sample-efficient fashion with only Vstar-realizability (realizability of the optimal state-action value function); existing online RL algorithms require significantly stronger representation conditions; and II) As a consequence, we show that the notorious Exogenous Block MDP problem Efroni et al. 2022 is tractable under local simulator access.
The results above are achieved through a computationally-inefficient algorithm. We complement them with an oracle-efficient algorithm, RVFS (``Recursive Value Function Search''), which achieves provable sample complexity guarantees under strengthened statistical assumption known as pushforward coverability. RVFS can be viewed as a principled, provable counterpart to a successful empirical paradigm that combines recursive search (e.g. MCTS) with value function approximation.
View details
