Congratulations to the Berkeley Artificial Intelligence Research (BAIR) Lab class of 2026! This year, BAIR celebrates another remarkable group of Ph.D. graduates whose curiosity, creativity, and perseverance have pushed the frontiers of artificial intelligence and machine learning. Their work spans the breadth of modern AI — robotics and embodied intelligence, large language models and reasoning, computer vision, generative modeling, AI safety, human-AI interaction, AI for science and healthcare, and much more. Along the way, they have published influential research, built systems with real-world impact, mentored their peers, and shaped the BAIR community for the better. Now they are headed everywhere ideas travel: to faculty and postdoctoral positions, to industry research labs, and to startups of their own founding — and several are still exploring what comes next and would love to hear from you. Please join us in celebrating the achievements of these wonderful graduates. We are proud of everything they have accomplished at Berkeley, and we can’t wait to see what they do next! Thank you to our friends at the Stanford AI Lab for this idea! Baifeng Shi Email: [email protected] Website: https://bfshi.github.io/ Advisor(s): Trevor Darrell Research Blurb: I work on building generalist vision and robotic models. What's next: Member of Technical Staff at Physical Intelligence Charlie Snell Email: [email protected] Website: https://sea-snell.github.io Advisor(s): Dan Klein Research Blurb: My work aims to understand when and how the different LLM scaling paradigms can be traded off and interchanged. In particular, test-time scaling treats each prompt independently, drawing long chains of inferences and then forgetting them entirely between prompts. This differs critically from pretraining, which instead learns a compressed representation from a large dataset. I believe bridging the gap between these methods of scaling computation, presents a key open challenge in the field: how can we develop methods which turn the inferences drawn at test-time back into learned representations that the model can hold onto across interactions. Devin Guillory Email: [email protected] Website: https://devinguillory.com Advisor(s): Trevor Darrell Research Blurb: Accounting for data shifts in computer vision models What's next: Building collaborative AI systems, looking for conspirators. Eve Fleisig Email: [email protected] Website: https://efleisig.com Advisor(s): Dan Klein Research Blurb: I design language models to work reliably and fairly for the broad range of real LLM users. First, my research leverages disagreement among user preferences as signal, in order to train and evaluate LLMs for entire populations of users. Second, I work on designing rigorous evaluations to extricate challenging LLM harms that diverse users face. Finally, I work on core technical failures of LLMs, like miscalibrated confidence, to reduce downstream risks when models are deployed to users with different needs. Combined, these interventions facilitate building LLMs that minimize societal harms, and maximize benefits to a wider range of real-world users. What's next: Postdoctoral fellow at Princeton CITP Grace Luo Email: [email protected] Website: https://graceluo.net Advisor(s): Trevor Darrell Research Blurb: My research is on interpreting and controlling generative models. For example, I've worked on re-purposing image generators for computer vision tasks, and meta-modeling language activations for better LLM probing and steering. What's next: Research scientist in industry Hanlin Zhu Email: [email protected] Website: https://hanlinzhu.com/ Advisor(s): Stuart Russell, Jiantao Jiao Research Blurb: My research centers on understanding and improving the reasoning capabilities of large language models (LLMs). What's next: Member of Technical Staff at OpenAI Haozhi Qi Email: [email protected] Website: https://haozhi.io/ Advisor(s): Jitendra Malik, Yi Ma Research Blurb: Dexterous Manipulation and Robot Learning What's next: Research scientist at Amazon; Faculty at University of Chicago J.D. Zamfirescu-Pereira Email: [email protected] Website: https://zamfi.net Advisor(s): Bjoern Hartmann Research Blurb: My research focuses on effective human-AI co-design. I study the boundaries of language interfaces as a medium for interacting with AI, creating systems that blend language-focused interactions with structured user interfaces that draw on different levels of abstraction. I focus on language-oriented technologies, like LLMs and text-to-image models, that are powerful mediators of design processes. These technologies enable humans to describe their desires at almost any level of abstraction, from high-level goals vaguely specified (“I’d like a game to help my kid learn to read”) to low-level corrections of undesired outputs (“Don’t say ‘I know because I’ve tasted it’ when about a recipe substitution's taste”). What's next: Assistant Professor, Computer Science, UCLA Jiachen Lian Email: [email protected] Website: https://jlian2.github.io Advisor(s): Gopala Anumanchipalli Research Blurb: My research focuses on human-centered AI across speech, healthcare, and systems. Looking for: Look for AI talents to join our startup Josh Kang Email: [email protected] Website: https://joshuaminwookang.github.io/ Advisor(s): John Canny Research Blurb: I study language modeling and related topics in NLP; specific interests are human user simulation and building conversational, collaborative AI agents. What's next: AI Scientist at Mistral AI Junhao (Bear) Xiong Email: [email protected] Website: https://www.linkedin.com/in/junhao-bear-xiong Advisor(s): Jennifer Listgarten, Yun Song Research Blurb: Junhao (Bear) Xiong is a PhD candidate at UC Berkeley, advised by Jennifer Listgarten and Yun S. Song. His work focuses on machine learning methods for biology, with an emphasis on generative modeling for proteins. Previously, he studied Applied Math and Computer Science at Johns Hopkins. Looking for: Research scientist Kaylo Littlejohn Email: [email protected] Website: https://kaylolittlejohn.com Advisor(s): Gopala Anumanchipalli Research Blurb: My research is focused on speech modeling and natural language processing. I co-led the development of multimodal AI tools to accurately translate brain activity into text, audible personalized speech, and a high-fidelity "digital talking avatar" (Nature 2023, Nature Neuroscience 2025). I am also tech lead for voice modeling at Roblox. Looking for: Research Scientist / Engineer Kent Chang Email: [email protected] Website: https://kentkc.org Advisor(s): David Bamman Research Blurb: I work on NLP and multimodal machine learning, with a focus on evaluating large language models and building multimodal systems for understanding dialogue, narrative, and social interaction. My research includes benchmarks for LLM memorization, multimodal datasets sourced from feature films and television, and studies of model behavior. I'm interested in bridging computational methods with questions from the humanities and social sciences about whose voices get represented in AI systems, and about AI's broader impact. My work has appeared at EMNLP and ACL, among others. Looking for: (teaching) faculty, Research Scientist, ML/AI SWE Kevin Black Email: [email protected] Website: https://kevin.black Advisor(s): Sergey Levine Research Blurb: I work on large-scale robot learning: including imitation learning, reinforcement learning, generative modeling, real-time control, and whatever else it takes to make robots work in the real world! What's next: Research Scientist of Physical Intelligence Kunhe Yang Email: [email protected] Website: https://www.kunheyang.com/ Advisor(s): Nika Haghtalab Research Blurb: My research focuses on the theoretical foundations of designing and evaluating AI algorithms in environments shaped by human incentives and AI agency. My work spans human-centric policy learning, incentive-aware evaluation, and multi-agent collaboration and information transmission, drawing on tools from machine learning theory and computational economics. What's next: Postdoc Research at Stanford Lisa Dunlap Email: [email protected] Website: https://lisabdunlap.com Advisor(s): Joseph Gonzalez, Trevor Darrell Research Blurb: Auditing generative models. What's next: Research Engineer at Anthropic Long (Tony) Lian Email: [email protected] Website: https://tonylian.com/ Advisor(s): Trevor Darrell, Adam Yala Research Blurb: My research primarily focuses on developing real-time multi-modal multi-agent systems and parallel reasoning systems through end-to-end RL. What's next: Member of Technical Staff at Thinking Machines Lab Maulik Bhatt Email: [email protected] Website: https://maulikb.com Advisor(s): Negar Mehr Research Blurb: My research develops autonomous robots that can safely coordinate with humans and other robots in shared environments. I build scalable algorithms grounded in game theory and diffusion models that let agents reason about the intent and behavior of others around them. My work spans real-time multi-agent trajectory planning and imitation learning in the presence of multi-modality. I've validated these methods on hardware platforms ranging from quadrotors to manipulators, with the goal of making multi-agent coordination robust, interpretable, and deployable in the real world. What's next: Joining Toyota Woven's end-to-end autonomous driving team. Michael Psenka Email: [email protected] Website: https://www.michaelpsenka.io/ Advisor(s): Aditi Krishnapriyan Research Blurb: Work in various domains (reinforcement learning, world models, AI+bio/chem), generally working on longer-horizon and out-of-distribution problems in planning and interpolation (e.g. robot manipulation from start state to goal, molecular dynamics of proteins between ground states). My thesis took a variational approach (think calculus of variations) directly from deep generative models of the environment, framing path-finding as minimizing a functional induced by the learned model itself (its score, its critic, or its dynamics). Through my research I've gained insight on how to properly handle dynamics in deep learning systems, and I plan to continue developing systems that are dynamic and adaptive. What's next: Lead Research Scientist at Baseten Nathan Lichtlé Email: [email protected] Website: https://nathanlichtle.com Advisor(s): Alexandre M. Bayen Research Blurb: RL for autonomous driving. What's next: Chief Scientist & Co-founder at Yumi Health Neerja Thakkar Email: [email protected] Website: https://neerja.me/ Advisor(s): Jitendra Malik Research Blurb: My research focuses on scaling predictive world models to handle the complexity of in-the-wild motion. Using autoregressive and diffusion frameworks, I develop better representations for real-world prediction and propose methods to efficiently adapt these models to new domains. Looking for: Research scientist Nikita Mehandru Email: [email protected] Website: https://n-mehandru.github.io/ Advisor(s): Ahmed Alaa and David Bamman Research Blurb: My research develops and applies machine learning methods for clinical reasoning and disease progression modeling using unstructured text and time series data from electronic health records. In collaboration with physicians at UCSF, I bridge method development and clinical validation with the intention to build reliable, interpretable AI systems in medicine. Looking for: Research Scientist Niklas Lauffer Email: [email protected] Website: https://niklaslauffer.github.io/ Advisor(s): Stuart Russell and Sanjit Seshia Research Blurb: Niklas's research is focused on AI safety and reinforcement learning, particularly in the area of multi-agent interaction and LM agents. He's worked on enabling adversarial learning in cooperative and mixed-motive settings, solving issues of covariate shift in training LM agents on long-horizon tasks, as well as evaluating safety risks posed by LM agents in multi-agent settings. What's next: Research Scientist at Google Deepmind Qiyang Li Email: [email protected] Website: https://colinqiyangli.github.io/ Advisor(s): Sergey Levine Research Blurb: Recent progress in robotic manipulation policy learning has been largely driven by (1) the increasing availability of large-scale prior datasets and (2) the success of action chunking, where the policy predicts a short sequence of future actions rather than a single one. However, most action chunking policies are trained via supervised imitation learning, because efficient online self-improvement with reinforcement learning (RL) remains challenging—limiting real-world applicability. My PhD research studied how we could leverage prior data to optimize action-chunking policies with RL, combining empirical results with theoretical insights. Looking for: Post-doc/research scientist for RL in robotics and LLMs! Sampada Deglurkar Email: [email protected] Website: https://sdeglurkar.github.io/ Advisor(s): Prof Claire Tomlin Research Blurb: My research is in providing safety assurances for AI-enabled autonomous systems, ranging from robots to autonomous vehicles to aviation systems. For this, I have worked with uncertainty quantification for machine learning models, decision-making under uncertainty algorithms, and tools for producing probabilistic guarantees on system operation. Looking for: Research scientist, Research engineer Vinamra Benara Email: [email protected] Website: https://cs.berkeley.edu/~vbenara Advisor(s): Ion Stoica Research Blurb: My research focuses on LLM post-training, including data curation, RLHF, RLVR with VLMs, evaluations, reasoning, agentic workflows, and interpretability. I also have strong expertise in systems infrastructure for distributed computing. Looking for: Research scientist / Research Engineer Vongani Maluleke Email: [email protected] Website: https://people.eecs.berkeley.edu/~vongani_maluleke/ Advisor(s): Jitendra Malik and Angjoo Kanazawa Research Blurb: Vongani Maluleke is a PhD candidate at UC Berkeley (BAIR, advised by Jitendra Malik and Angjoo Kanazawa), where she led the development of MAGNet, a unified multi-agent motion generation framework that supports a wide range of motion generation tasks without retraining or architectural changes, outperforming task-specialized state-of-the-art baselines. She is currently extending this work by deploying it on a Unitree G1 humanoid to make it embody social intelligence. Before her PhD, she was a Senior AI Consultant at Deloitte, awarded Exceptional Performer two consecutive years, leading AI system development across media, telecommunications, retail, and financial services. Looking for: Research scientist Wei-Jer Chang Email: [email protected] Website: https://weijer-chang.github.io/ Advisor(s): Masayoshi Tomizuka Research Blurb: My research focuses on developing safe and intelligent autonomous systems for complex, human-centered environments. I work at the intersection of machine learning, generative models, and reinforcement learning, with applications in autonomy. My work addresses challenges in multi-agent interaction, interactive human behavior, and long-tail safety-critical scenarios at scale. Looking for: Research Scientist, Applied Scientist, Roboticist Xiuyu Li Email: [email protected] Website: https://xiuyuli.com/ Advisor(s): Kurt Keutzer Research Blurb: My research focuses on developing scalable and self-improving large language model agents, with emphasis on coding agents for complex, long-horizon tasks. This direction builds on my work in parallel reasoning, and on broader expertise in making generative models more efficient in training and inference across language and vision. What's next: Member of Technical Staff at xAI Yichen Xie Email: [email protected] Website: https://yichen928.github.io/ Advisor(s): Masayoshi Tomizuka Research Blurb: My research focuses on building multimodal foundation models and world models that understand and interact with complex physical environments. I aim to develop unified representations across modalities, enabling AI systems to reason over space, time, and dynamics toward general-purpose embodied intelligence. What's next: Research Scientist at Luma AI Yigit Efe Erginbas Email: [email protected] Website: https://www.linkedin.com/in/erginbas/ Advisor(s): Kannan Ramchandran, Thomas A. Courtade Research Blurb: My PhD research spans two threads: online learning in large-scale markets, and interpretability of large machine learning models. In the first, I work on sequential decision-making with applications to recommendation, pricing, and assortment selection. My focus is on designing algorithms with provable guarantees for welfare maximization, revenue maximization, and stability. In the second, I develop scalable attribution methods that exploit the sparse, low-degree structure of real-world interactions, using tools from signal processing and information theory. More recently, I have been exploring principled ways to evaluate the faithfulness of model self-explanations. What's next: Researcher at Hudson River Trading's AI Labs (HAIL) Yiheng Li Email: [email protected] Website: https://Yihengli.com Advisor(s): Masayoshi Tomizuka Research Blurb: I am working on vision world modeling, with prior experience in diffusion model's efficiency as well as in autonomous driving. What's next: Research Scientist at Waymo Zhe Fu Email: [email protected] Website: https://fu-zhe.com/ Advisor(s): Alexandre Bayen Research Blurb: My research focuses on physics-informed learning and control for mixed-autonomy systems, with applications in transportation. I design physics-informed neural networks to learn solutions of nonlinear partial differential equations, enabling accurate and data-efficient prediction of traffic dynamics. Building on these models, I develop both model-based and learning-based control strategies that coordinate automated vehicles to improve system-level performance. My work bridges machine learning, control, and real-world deployment, and has been validated in large-scale field experiments. More broadly, I aim to advance trustworthy, interpretable AI for decision-making in complex, real-world systems. What's next: I will be an Energy Fellow at Stanford after graduation. Also looking for Faculty, or research scientist positions in AI, control, and autonomy.

.grasp-results-table table { font-size: 0.875rem; line-height: 1.35; width: 100%; } .grasp-results-table th, .grasp-results-table td { padding: 0.35rem 0.5rem; } /* Consistent whitespace between major sections (this post is long and hr-heavy) */ article.post-content h2 { margin-top: 2.75rem; margin-bottom: 0.75rem; } article.post-content h2:first-of-type { margin-top: 2.25rem; } article.post-content h3 { margin-top: 1.65rem; margin-bottom: 0.5rem; } article.post-content hr { margin-top: 2.5rem; margin-bottom: 2.5rem; } GRASP is a new gradient-based planner for learned dynamics (a “world model”) that makes long-horizon planning practical by (1) lifting the trajectory into virtual states so optimization is parallel across time, (2) adding stochasticity directly to the state iterates for exploration, and (3) reshaping gradients so actions get clean signals while we avoid brittle “state-input” gradients through high-dimensional vision models. Large, learned world models are becoming increasingly capable. They can predict long sequences of future observations in high-dimensional visual spaces and generalize across tasks in ways that were difficult to imagine a few years ago. As these models scale, they start to look less like task-specific predictors and more like general-purpose simulators. But having a powerful predictive model is not the same as being able to use it effectively for control/learning/planning. In practice, long-horizon planning with modern world models remains fragile: optimization becomes ill-conditioned, non-greedy structure creates bad local minima, and high-dimensional latent spaces introduce subtle failure modes. In this blog post, I describe the problems that motivated this project and our approach to address them: why planning with modern world models can be surprisingly fragile, why long horizons are the real stress test, and what we changed to make gradient-based planning much more robust. This blog post discusses work done with Mike Rabbat, Aditi Krishnapriyan, Yann LeCun, and Amir Bar (* denotes equal advisorship), where we propose GRASP. What is a world model? These days, the term “world model” is quite overloaded, and depending on the context can either mean an explicit dynamics model or some implicit, reliable internal state that a generative model relies on (e.g. when an LLM generates chess moves, whether there is some internal representation of the board). We give our loose working definition below. Suppose you take actions $a_t \in \mathcal{A}$ and observe states $s_t \in \mathcal{S}$ (images, latent vectors, proprioception). A world model is a learned model that, given the current state and a sequence of future actions, predicts what will happen next. Formally, it defines a predictive distribution on a sequence of observed states $s_{t-h:t}$ and current action $a_t$: that approximates the environment’s true conditional $P(s_{t+1} \mid s_{t-h:t},\; a_t)$. For this blog post, we’ll assume a Markovian model $P(s_{t+1} \mid s_{t-h:t},\; a_t)$ for simplicity (all results here can be extended to the more general case), and when the model is deterministic it reduces to a map over states: In practice the state $s_t$ is often a learned latent representation (e.g., encoded from pixels), so the model operates in a (theoretically) compact, differentiable space. The key point is that a world model gives you a differentiable simulator; you can roll it forward under hypothetical action sequences and backpropagate through the predictions. Planning: choosing actions by optimizing through the model Given a start $s_0$ and a goal $g$, the simplest planner chooses an action sequence $\mathbf{a}=(a_0,\dots,a_{T-1})$ by rolling out the model and minimizing terminal error: Here we use $\mathcal{F}^T$ as shorthand for the full rollout through the world model (dependence on model parameters $\theta$ is implicit): In short horizons and low-dimensional systems, this can work reasonably well. But as horizons grow and models become larger and more expressive, its weaknesses become amplified. So why doesn’t this just work at scale? Why long-horizon planning is hard (even when everything is differentiable) There are two separate pain points for the more general world model, plus a third that is specific to learned, deep learning-based models. 1) Long-horizon rollouts create deep, ill-conditioned computation graphs Those familiar with backprop through time (BPTT) may notice that we’re differentiating through a model applied to itself repeatedly, which will lead to the exploding/vanishing gradients problem. Namely, if we take derivatives (note we’re differentiating vector-valued functions, resulting in Jacobians that we denote with $D_x (\cdots)$) with respect to earlier actions (e.g. $a_0$): We see that the Jacobian’s conditioning scales exponentially with time $T$: leading to exploding or vanishing gradients. 2) The landscape is non-greedy and full of traps At short horizons, the greedy solution, where we move straight toward the goal at every step, is often good enough. If you only need to plan a few steps ahead, the optimal trajectory usually doesn’t deviate much from “head toward $g$” at each step. As horizons grow, two things happen. First, longer tasks are more likely to require non-greedy behavior: going around a wall, repositioning before pushing, backing up to take a better path. And as horizons grow, more of these non-greedy steps are typically needed. Second, the optimization space itself scales with horizon: $\mathrm{dim}(\mathcal{A} \times \cdots \times \mathcal{A}) = T\mathrm{dim}(\mathcal{A})$, further expanding the space of local minima for the optimization problem. Distance to goal along the optimal path is non-monotonic, and the resulting loss landscape can be rough. A long-horizon fix: lifting the dynamics constraint Suppose we treat the dynamics constraint $s_{t+1} = F_{\theta}(s_t, a_t)$ as a soft constraint, and we instead optimize the following penalty function over both actions $(a_0,\ldots,a_{T-1})$ and states $(s_0,\ldots,s_T)$: This is also sometimes called collocation in planning/robotics literature. Note the lifted formulation shares the same global minimizers as the original rollout objective (both are zero exactly when the trajectory is dynamically feasible). But the optimization landscapes are very different, and we get two immediate benefits: Each world model evaluation $F_{\theta}(s_t,a_t)$ depends only on local variables, so all $T$ terms can be computed in parallel across time, resulting in a huge speed-up for longer horizons, and You no longer backpropagate through a single deep $T$-step composition to get a learning signal, since the previous product of Jacobians now splits into a sum, e.g.: Being able to optimize states directly also helps with exploration, as we can temporarily navigate through unphysical domains to find the optimal plan: Collocation-based planning allows us to directly perturb states and explore midpoints more effectively. However, lunch is never free. And indeed, especially for deep learning-based world models, there is a critical issue that makes the above optimization quite difficult in practice. An issue for deep learning-based world models: sensitivity of state-input gradients The tl;dr of this section is: directly optimizing states through a deep learning-based $F_{\theta}$ is incredibly brittle, à la adversarial robustness. Even if you train your world model in a lower-dimensional state space, the training process for the world model makes unseen state landscapes very sharp, whether it be an unseen state itself or simply a normal/orthogonal direction to the data manifold. Adversarial robustness and the “dimpled manifold” model Adversarial robustness originally looked at classification models $f_\theta : \mathbb{R}^{w\times h \times c} \to \mathbb{R}^K$, and showed that by following the gradient of a particular logit $\nabla f_\theta^k$ from a base image $x$ (not of class $k$), you did not have to move far along $x’ = x + \epsilon\nabla f_\theta^k$ to make $f_\theta$ classify $x’$ as $k$ (Szegedy et al., 2014; Goodfellow et al., 2015): Depiction of the classic example from (Goodfellow et al., 2015). Later work has painted a geometric picture for what’s going on: for data near a low-dimensional manifold $\mathcal{M}$, the training process controls behavior in tangential directions, but does not regularize behavior in orthogonal directions, thus leading to sensitive behavior (Stutz et al., 2019). Another way stated: $f_\theta$ has a reasonable Lipschitz constant when considering only tangential directions to the data manifold $\mathcal{M}$, but can have very high Lipschitz constants in normal directions. In fact, it often benefits the model to be sharper in these normal directions, so it can fit more complicated functions more precisely. As a result, such adversarial examples are incredibly common even for a single given model. Further, this is not just a computer vision phenomenon; adversarial examples also appear in LLMs (Wallace et al., 2019) and in RL (Gleave et al., 2019). While there are methods to train for more adversarially robust models, there is a known trade-off between model performance and adversarial robustness (Tsipras et al., 2019): especially in the presence of many weakly-correlated variables, the model must be sharper to achieve higher performance. Indeed, most modern training algorithms, whether in computer vision or LLMs, do not train adversarial robustness out. Thus, at least until deep learning sees a major regime change, this is a problem we’re stuck with. Why is adversarial robustness an issue for world model planning? Consider a single component of the dynamics loss we’re optimizing in the lifted state approach: Let’s further focus on just the base state: Since world models are typically trained on state/action trajectories $(s_1, a_1, s_2, a_2, \ldots)$, the state-data manifold for $F_{\theta}$ has dimensionality bounded by the action space: where $\mathcal{R}$ is some optional space of augmentations (e.g. translations/rotations). Thus, we can typically expect $\mathrm{dim}(\mathcal{M}_s)$ to be much lower than $\mathrm{dim}(\mathcal{S})$, and thus: it is very easy to find adversarial examples that hack any state to any other desired state. As a result, the dynamics optimization feels incredibly “sticky,” as the base points $s_t$ can easily trick $F_{\theta}$ into thinking it’s already made its local goal.1 1. This adversarial robustness issue, while particularly bad for lifted-state approaches, is not unique to them. Even for serial optimization methods that optimize through the full rollout map $\mathcal{F}^T$, it is possible to get into unseen states, where it is very easy to have a normal component fed into the sensitive normal components of $D_s F_{\theta}$. The action Jacobian’s chain rule expansion is See what happens if any stage of the product has any component normal to the data manifold. ↩ Our fix This is where our new planner GRASP comes in. The main observation: while $D_s F_{\theta}$ is untrustworthy and adversarial, the action space is usually low-dimensional and exhaustively trained, so $D_a F_{\theta}$ is actually reasonable to optimize through and doesn’t suffer from the adversarial robustness issue! The action input is usually lower-dimensional and densely trained (the model has seen every action direction), so action gradients are much better behaved. At its core, GRASP builds a first-order lifted state / collocation-based planner that is only dependent on action Jacobians through the world model. We thus exploit the differentiability of learned world models $F_{\theta}$, while not falling victim to the inherent sensitivity of the state Jacobians $D_s F_{\theta}$. GRASP: Gradient RelAxed Stochastic Planner As noted before, we start with the collocation planning objective, where we lift the states and relax dynamics into a penalty: We then make two key additions. Ingredient 1: Exploration by noising the state iterates Even with a smoother objective, planning is nonconvex. We introduce exploration by injecting Gaussian noise into the virtual state updates during optimization. A simple version: Actions are still updated by non-stochastic descent: The state noise helps you “hop” between basins in the lifted space, while the actions remain guided by gradients. We found that specifically noising states here (as opposed to actions) finds a good balance of exploration and the ability to find sharper minima.2 2. Because we only noise the states (and not the actions), the corresponding dynamics are not truly Langevin dynamics. ↩ Ingredient 2: Reshape gradients: stop brittle state-input gradients, keep action gradients As discussed, the fragile pathway is the gradient that flows into the state input of the world model, \(D_s F_{\theta}\). The most straightforward way to do this initially is to just stop state gradients into \(F_{\theta}\) directly: Let $\bar{s}_t$ be the same value as $s_t$, but with gradients stopped. Define the stop-gradient dynamics loss: This alone does not work. Notice now states only follow the previous state’s step, without anything forcing the base states to chase the next ones. As a result, there are trivial minima for just stopping at the origin, then only for the final action trying to get to the goal in one step. Dense goal shaping We can view the above issue as the goal’s signal being cut off entirely from previous states. One way to fix this is to simply add a dense goal term throughout prediction: In normal settings this would over-bias towards the greedy solution of straight chasing the goal, but this is balanced in our setting by the stop-gradient dynamics loss’s bias towards feasible dynamics. The final objective is then as follows: The result is a planning optimization objective that does not have dependence on state gradients. Periodic “sync”: briefly return to true rollout gradients The lifted stop-gradient objective is great for fast, guided exploration, but it’s still an approximation of the original serial rollout objective. So every $K_{\text{sync}}$ iterations, GRASP does a short refinement phase: Roll out from $s_0$ using current actions $\mathbf{a}$, and take a few small gradient steps on the original serial loss: The lifted-state optimization still provides the core of the optimization, while this refinement step adds some assistance to keep states and actions grounded towards real trajectories. This refinement step can of course be replaced with a serial planner of your choice (e.g. CEM); the core idea is to still get some of the benefit of the full-path synchronization of serial planners, while still mostly using the benefits of the lifted-state planning. How GRASP addresses long-range planning Collocation-based planners offer a natural fix for long-horizon planning, but this optimization is quite difficult through modern world models due to adversarial robustness issues. GRASP proposes a simple solution for a smoother collocation-based planner, alongside stable stochasticity for exploration. As a result, longer-horizon planning ends up not only succeeding more, but also finding such successes faster: Push-T demo: longer-horizon planning with GRASP. Horizon CEM GD LatCo GRASP H=40 61.4% / 35.3s 51.0% / 18.0s 15.0% / 598.0s 59.0% / 8.5s H=50 30.2% / 96.2s 37.6% / 76.3s 4.2% / 1114.7s 43.4% / 15.2s H=60 7.2% / 83.1s 16.4% / 146.5s 2.0% / 231.5s 26.2% / 49.1s H=70 7.8% / 156.1s 12.0% / 103.1s 0.0% / — 16.0% / 79.9s H=80 2.8% / 132.2s 6.4% / 161.3s 0.0% / — 10.4% / 58.9s Push-T results. Success rate (%) / median time to success. Bold = best in row. Note the median success time will bias higher with higher success rate; GRASP manages to be faster despite higher success rate. What’s next? There is still plenty of work to be done for modern world model planners. We want to exploit the gradient structure of learned world models, and collocation (lifted-state optimization) is a natural approach for long-horizon planning, but it’s crucial to understand typical gradient structure here: smooth and informative action gradients and brittle state gradients. We view GRASP as an initial iteration for such planners. Extension to diffusion-based world models (deeper latent timesteps can be viewed as smoothed versions of the world model itself), more sophisticated optimizers and noising strategies, and integrating GRASP into either a closed-loop system or RL policy learning for adaptive long-horizon planning are all natural and interesting next steps. I do genuinely think it’s an exciting time to be working on world model planners. It’s a funny sweet spot where the background literature (planning and control overall) is incredibly mature and well-developed, but the current setting (pure planning optimization over modern, large-scale world models) is still heavily underexplored. But, once we figure out all the right ideas, world model planners will likely become as commonplace as RL. For more details, read the full paper or visit the project website. Citation @article{psenka2026grasp, title={Parallel Stochastic Gradient-Based Planning for World Models}, author={Michael Psenka and Michael Rabbat and Aditi Krishnapriyan and Yann LeCun and Amir Bar}, year={2026}, eprint={2602.00475}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2602.00475} }

--> Understanding the behavior of complex machine learning systems, particularly Large Language Models (LLMs), is a critical challenge in modern artificial intelligence. Interpretability research aims to make the decision-making process more transparent to model builders and impacted humans, a step toward safer and more trustworthy AI. To gain a comprehensive understanding, we can analyze these systems through different lenses: feature attribution, which isolates the specific input features driving a prediction (Lundberg & Lee, 2017; Ribeiro et al., 2022); data attribution, which links model behaviors to influential training examples (Koh & Liang, 2017; Ilyas et al., 2022); and mechanistic interpretability, which dissects the functions of internal components (Conmy et al., 2023; Sharkey et al., 2025). Across these perspectives, the same fundamental hurdle persists: complexity at scale. Model behavior is rarely the result of isolated components; rather, it emerges from complex dependencies and patterns. To achieve state-of-the-art performance, models synthesize complex feature relationships, find shared patterns from diverse training examples, and process information through highly interconnected internal components. Therefore, grounded or reality-checked interpretability methods must also be able to capture these influential interactions. As the number of features, training data points, and model components grow, the number of potential interactions grows exponentially, making exhaustive analysis computationally infeasible. In this blog post, we describe the fundamental ideas behind SPEX and ProxySPEX, algorithms capable of identifying these critical interactions at scale. Attribution through Ablation Central to our approach is the concept of ablation, measuring influence by observing what changes when a component is removed. Feature Attribution: We mask or remove specific segments of the input prompt and measure the resulting shift in the predictions. Data Attribution: We train models on different subsets of the training set, assessing how the model’s output on a test point shifts in the absence of specific training data. Model Component Attribution (Mechanistic Interpretability): We intervene on the model’s forward pass by removing the influence of specific internal components, determining which internal structures are responsible for the model’s prediction. In each case, the goal is the same: to isolate the drivers of a decision by systematically perturbing the system, in hopes of discovering influential interactions. Since each ablation incurs a significant cost, whether through expensive inference calls or retrainings, we aim to compute attributions with the fewest possible ablations. --> SPEX and ProxySPEX Framework To discover influential interactions with a tractable number of ablations, we have developed SPEX (Spectral Explainer). This framework draws on signal processing and coding theory to advance interaction discovery to scales orders of magnitude greater than prior methods. SPEX circumvents this by exploiting a key structural observation: while the number of total interactions is prohibitively large, the number of influential interactions is actually quite small. We formalize this through two observations: sparsity (relatively few interactions truly drive the output) and low-degreeness (influential interactions typically involve only a small subset of features). These properties allow us to reframe the difficult search problem into a solvable sparse recovery problem. Drawing on powerful tools from signal processing and coding theory, SPEX uses strategically selected ablations to combine many candidate interactions together. Then, using efficient decoding algorithms, we disentangle these combined signals to isolate the specific interactions responsible for the model’s behavior. --> In a subsequent algorithm, ProxySPEX, we identified another structural property common in complex machine learning models: hierarchy. This means that where a higher-order interaction is important, its lower-order subsets are likely to be important as well. This additional structural observation yields a dramatic improvement in computational cost: it matches the performance of SPEX with around 10x fewer ablations. Collectively, these frameworks enable efficient interaction discovery, unlocking new applications in feature, data, and model component attribution. Feature Attribution Feature attribution techniques assign importance scores to input features based on their influence on the model’s output. For example, if an LLM were used to make a medical diagnosis, this approach could identify exactly which symptoms led the model to its conclusion. While attributing importance to individual features can be valuable, the true power of sophisticated models lies in their ability to capture complex relationships between features. The figure below illustrates examples of these influential interactions: from a double negative changing sentiment (left) to the necessary synthesis of multiple documents in a RAG task (right). --> The figure below illustrates the feature attribution performance of SPEX on a sentiment analysis task. We evaluate performance using faithfulness: a measure of how accurately the recovered attributions can predict the model’s output on unseen test ablations. We find that SPEX matches the high faithfulness of existing interaction techniques (Faith-Shap, Faith-Banzhaf) on short inputs, but uniquely retains this performance as the context scales to thousands of features. In contrast, while marginal approaches (LIME, Banzhaf) can also operate at this scale, they exhibit significantly lower faithfulness because they fail to capture the complex interactions driving the model’s output. --> SPEX was also applied to a modified version of the trolley problem, where the moral ambiguity of the problem is removed, making “True” the clear correct answer. Given the modification below, GPT-4o mini answered correctly only 8% of the time. When we applied standard feature attribution (SHAP), it identified individual instances of the word trolley as the primary factors driving the incorrect response. However, replacing trolley with synonyms such as tram or streetcar had little impact on the prediction of the model. SPEX revealed a much richer story, identifying a dominant high-order synergy between the two instances of trolley, as well as the words pulling and lever, a finding that aligns with human intuition about the core components of the dilemma. When these four words were replaced with synonyms, the model’s failure rate dropped to near zero. --> Data Attribution Data attribution identifies which training data points are most responsible for a model’s prediction on a new test point. Identifying influential interactions between these data points is key to explaining unexpected model behaviors. Redundant interactions, such as semantic duplicates, often reinforce specific (and possibly incorrect) concepts, while synergistic interactions are essential for defining decision boundaries that no single sample could form alone. To demonstrate this, we applied ProxySPEX to a ResNet model trained on CIFAR-10, identifying the most significant examples of both interaction types for a variety of difficult test points, as shown in the figure below. --> As illustrated, synergistic interactions (left) often involve semantically distinct classes working together to define a decision boundary. For example, grounding the synergy in human perception, the automobile (bottom left) shares visual traits with the provided training images, including the low-profile chassis of the sports car, the boxy shape of the yellow truck, and the horizontal stripe of the red delivery vehicle. On the other hand, redundant interactions (right) tend to capture visual duplicates that reinforce a specific concept. For instance, the horse prediction (middle right) is heavily influenced by a cluster of dog images with similar silhouettes. This fine-grained analysis allows for the development of new data selection techniques that preserve necessary synergies while safely removing redundancies. Attention Head Attribution (Mechanistic Interpretability) The goal of model component attribution is to identify which internal parts of the model, such as specific layers or attention heads, are most responsible for a particular behavior. Here too, ProxySPEX uncovers the responsible interactions between different parts of the architecture. Understanding these structural dependencies is vital for architectural interventions, such as task-specific attention head pruning. On an MMLU dataset (highschool‐us‐history), we demonstrate that a ProxySPEX-informed pruning strategy not only outperforms competing methods, but can actually improve model performance on the target task. --> On this task, we also analyzed the interaction structure across the model’s depth. We observe that early layers function in a predominantly linear regime, where heads contribute largely independently to the target task. In later layers, the role of interactions between attention heads becomes more pronounced, with most of the contribution coming from interactions among heads in the same layer. --> What’s Next? The SPEX framework represents a significant step forward for interpretability, extending interaction discovery from dozens to thousands of components. We have demonstrated the versatility of the framework across the entire model lifecycle: exploring feature attribution on long-context inputs, identifying synergies and redundancies among training data points, and discovering interactions between internal model components. Moving forwards, many interesting research questions remain around unifying these different perspectives, providing a more holistic understanding of a machine learning system. It is also of great interest to systematically evaluate interaction discovery methods against existing scientific knowledge in fields such as genomics and materials science, serving to both ground model findings and generate new, testable hypotheses. We invite the research community to join us in this effort: the code for both SPEX and ProxySPEX is fully integrated and available within the popular SHAP-IQ repository (link). https://github.com/mmschlk/shapiq (SHAP-IQ Github) https://openreview.net/forum?id=KI8qan2EA7 (ProxySPEX NeurIPS 2025) https://openreview.net/forum?id=pRlKbAwczl (SPEX ICML 2025) https://openreview.net/forum?id=glGeXu1zG4 (Learning to Understand NeurIPS 2024)

An encoder (optical system) maps objects to noiseless images, which noise corrupts into measurements. Our information estimator uses only these noisy measurements and a noise model to quantify how well measurements distinguish objects. Many imaging systems produce measurements that humans never see or cannot interpret directly. Your smartphone processes raw sensor data through algorithms before producing the final photo. MRI scanners collect frequency-space measurements that require reconstruction before doctors can view them. Self-driving cars process camera and LiDAR data directly with neural networks. What matters in these systems is not how measurements look, but how much useful information they contain. AI can extract this information even when it is encoded in ways that humans cannot interpret. And yet we rarely evaluate information content directly. Traditional metrics like resolution and signal-to-noise ratio assess individual aspects of quality separately, making it difficult to compare systems that trade off between these factors. The common alternative, training neural networks to reconstruct or classify images, conflates the quality of the imaging hardware with the quality of the algorithm. We developed a framework that enables direct evaluation and optimization of imaging systems based on their information content. In our NeurIPS 2025 paper, we show that this information metric predicts system performance across four imaging domains, and that optimizing it produces designs that match state-of-the-art end-to-end methods while requiring less memory, less compute, and no task-specific decoder design. Why mutual information? Mutual information quantifies how much a measurement reduces uncertainty about the object that produced it. Two systems with the same mutual information are equivalent in their ability to distinguish objects, even if their measurements look completely different. This single number captures the combined effect of resolution, noise, sampling, and all other factors that affect measurement quality. A blurry, noisy image that preserves the features needed to distinguish objects can contain more information than a sharp, clean image that loses those features. Information unifies traditionally separate quality metrics. It accounts for noise, resolution, and spectral sensitivity together rather than treating them as independent factors. Previous attempts to apply information theory to imaging faced two problems. The first approach treated imaging systems as unconstrained communication channels, ignoring the physical limitations of lenses and sensors. This produced wildly inaccurate estimates. The second approach required explicit models of the objects being imaged, limiting generality. Our method avoids both problems by estimating information directly from measurements. Estimating information from measurements Estimating mutual information between high-dimensional variables is notoriously difficult. Sample requirements grow exponentially with dimensionality, and estimates suffer from high bias and variance. However, imaging systems have properties that enable decomposing this hard problem into simpler subproblems. Mutual information can be written as: The first term, $H(Y)$, measures total variation in measurements from both object differences and noise. The second term, $H(Y \mid X)$, measures variation from noise alone. Mutual information equals the difference between total measurement variation and noise-only variation. Imaging systems have well-characterized noise. Photon shot noise follows a Poisson distribution. Electronic readout noise is Gaussian. This known noise physics means we can compute $H(Y \mid X)$ directly, leaving only $H(Y)$ to be learned from data. For $H(Y)$, we fit a probabilistic model (e.g. a transformer or other autoregressive model) to a dataset of measurements. The model learns the distribution of all possible measurements. We tested three models spanning efficiency-accuracy tradeoffs: a stationary Gaussian process (fastest), a full Gaussian (intermediate), and an autoregressive PixelCNN (most accurate). The approach provides an upper bound on true information; any modeling error can only overestimate, never underestimate. Validation across four imaging domains Information estimates should predict decoder performance if they capture what limits real systems. We tested this relationship across four imaging applications. Information estimates predict decoder performance across color photography, radio astronomy, lensless imaging, and microscopy. Higher information consistently produces better results on downstream tasks. Color photography. Digital cameras encode color using filter arrays that restrict each pixel to detect only certain wavelengths. We compared three filter designs: the traditional Bayer pattern, a random arrangement, and a learned arrangement. Information estimates correctly ranked which designs would produce better color reconstructions, matching the rankings from neural network demosaicing without requiring any reconstruction algorithm. Radio astronomy. Telescope arrays achieve high angular resolution by combining signals from sites across the globe. Selecting optimal telescope locations is computationally intractable because each site’s value depends on all others. Information estimates predicted reconstruction quality across telescope configurations, enabling site selection without expensive image reconstruction. Lensless imaging. Lensless cameras replace traditional optics with light-modulating masks. Their measurements bear no visual resemblance to scenes. Information estimates predicted reconstruction accuracy across a lens, microlens array, and diffuser design at various noise levels. Microscopy. LED array microscopes use programmable illumination to generate different contrast modes. Information estimates correlated with neural network accuracy at predicting protein expression from cell images, enabling evaluation without expensive protein labeling experiments. In all cases, higher information meant better downstream performance. Designing systems with IDEAL Information estimates can do more than evaluate existing systems. Our Information-Driven Encoder Analysis Learning (IDEAL) method uses gradient ascent on information estimates to optimize imaging system parameters. IDEAL optimizes imaging system parameters through gradient feedback on information estimates, without requiring a decoder network. The standard approach to computational imaging design, end-to-end optimization, jointly trains the imaging hardware and a neural network decoder. This requires backpropagating through the entire decoder, creating memory constraints and potential optimization difficulties. IDEAL avoids these problems by optimizing the encoder alone. We tested it on color filter design. Starting from a random filter arrangement, IDEAL progressively improved the design. The final result matched end-to-end optimization in both information content and reconstruction quality. IDEAL matches end-to-end optimization performance while avoiding decoder complexity during training. Implications Information-based evaluation creates new possibilities for rigorous assessment of imaging systems in real-world conditions. Current approaches require either subjective visual assessment, ground truth data that is unavailable in deployment, or isolated metrics that miss overall capability. Our method provides an objective, unified metric from measurements alone. The computational efficiency of IDEAL suggests possibilities for designing imaging systems that were previously intractable. By avoiding decoder backpropagation, the approach reduces memory requirements and training complexity. We explore these capabilities more extensively in follow-on work. The framework may extend beyond imaging to other sensing domains. Any system that can be modeled as deterministic encoding with known noise characteristics could benefit from information-based evaluation and design, including electronic, biological, and chemical sensors. This post is based on our NeurIPS 2025 paper “Information-driven design of imaging systems”. Code is available on GitHub. A video summary is available on the project website.