Antoine Dedieu

About me

I am a Senior Research Scientist at Google DeepMind. Over the past year, my research has focused on developing agents capable of rapid task adaptation in novel environments. To achieve this, I explore (a) learning a generative world model describing the environment dynamics and (b) leveraging LLMs as rich priors. In addition, I contribute to Google DeepMind's initiative on applying reinforcement learning to accelerate the path to nuclear fusion.

Prior to joining Google DeepMind, I spent 3.5 years at Vicarious (acquired by Alphabet), where I built an AI layer for robots. There, I designed an object detection pipeline used 1M+ times in production. My research focused on building novel generative probabilistic graphical models to solve challenging object-centric vision and navigation problems.

I hold a Master of Science from the MIT Operations Research Center, advised by Prof. Rahul Mazumder, and an MS in Applied Mathematics from Ecole Polytechnique. I am also a movie-fan, a sports addict, and a world-traveller. As an amateur photographer, I share my journeys on my website.

---

Conference articles

• Code world models for general game playing. [Preprint]
  To appear at ICLR 2026. Wolfgang Lehrach, Daniel Hennes, Miguel Lazaro-Gredilla, Xinghua Lou, Carter Wendelken, Zun Li, Antoine Dedieu, Jordi Grau-Moya, Marc Lanctot, Atil Iscen, John Schultz, Marcus Chiam, Ian Gemp, Piotr Zielinski, Satinder Singh, Kevin P Murphy
• Improving Transformer World Models for Data-Efficient RL. [Preprint] [YouTube]
  Outstanding paper award at ICLR 2025 World Model Workshop
  ICML 2025. Antoine Dedieu*, Joseph Ortiz*, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, Swaroop Guntupalli, Miguel Lázaro-Gredilla, Kevin Murphy
• DMC-VB: A Benchmark for Representation Learning for Control with Visual Distractors. [Preprint] [Code]
  Neurips 2024. Antoine Dedieu*, Joseph Ortiz*, Wolfgang Lehrach, Swaroop Guntupalli, Carter Wendelken, Ahmad Humayun, Sivaramakrishnan Swaminathan, Guangyao Zhou, Miguel Lázaro-Gredilla, Kevin Murphy.
• Diffusion Model Predictive Control. [Preprint]
  TMLR 2025. Guangyao Zhou, Sivaramakrishnan Swaminathan, Rajkumar Vasudeva Raju, Swaroop Guntupalli, Wolfgang Lehrach, Joseph Ortiz, Antoine Dedieu, Miguel Lázaro-Gredilla, Kevin Murphy.
• Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments. [Preprint]
  ICML 2024. Antoine Dedieu, Wolfgang Lehrach, Guangyao Zhou, Dileep George, Miguel Lázaro-Gredilla.
• Schema-learning and rebinding as mechanisms of in-context learning and emergence. [Preprint]
  Neurips 2023, Spotlight. Sivaramakrishnan Swaminathan, Antoine Dedieu, Rajkumar Vasudeva Raju, Murray Shanahan, Miguel Lázaro-Gredilla, Dileep George.
• Learning noisy-OR Bayesian Networks with Max-Product Belief Propagation. [Preprint] [Code]
  ICML 2023. Antoine Dedieu, Guangyao Zhou, Dileep George, Miguel Lázaro-Gredilla.
• Graphical Models with Attention for Context-Specific Independence and an Application to Perceptual Grouping. [Preprint] [Code]
  Guangyao Zhou, Wolfgang Lehrach, Antoine Dedieu, Miguel Lázaro-Gredilla, Dileep George.
• Perturb-and-max-product: Sampling and learning in discrete energy-based models. [Preprint] [Code]
  Neurips 2021. Miguel Lázaro-Gredilla, Antoine Dedieu, Dileep George.
• Sample-Efficient L0-L2 Constrained Structure Learning of Sparse Ising Models. [Preprint] [Code]
  AAAI 2021. Antoine Dedieu, Miguel Lázaro-Gredilla, Dileep George.
• Query Training: Learning a Worse Model to Infer Better Marginals in Undirected Graphical Models with Hidden Variables. [Preprint] [Code]
  AAAI 2021. Miguel Lázaro-Gredilla, Wolfgang Lehrach, Nishad Gothoskar, Guangyao Zhou, Antoine Dedieu, Dileep George.
• Improved error rates for sparse (group) learning with Lipschitz loss functions. [Preprint]
  Antoine Dedieu.
• An error bound for Lasso and Group Lasso in high dimensions. [Preprint]
  Antoine Dedieu.
• Learning higher-order sequential structure with cloned HMMs. [Preprint]
  Antoine Dedieu, Nishad Gothoskar, Scott Swingle, Wolfgang Lehrach, Miguel Lázaro-Gredilla, Dileep George.
• Error bounds for sparse classifiers in high-dimensions. [Preprint]
  AiStats 2019. Antoine Dedieu.
• Hierarchical Modeling and Shrinkage for User Session Length Prediction in Media Streaming. [Preprint] [Code]
  CIKM 2018. Antoine Dedieu, Rahul Mazumder, Zhen Zhu, Hossein Vahabi.

Journal articles

• PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX. [Preprint] [Code]
  Journal of Machine Learning Research, 2025. Guangyao Zhou, Antoine Dedieu, Nishanth Kumar, Miguel Lázaro-Gredilla, Shrinu Kushagra, Dileep George.
• A detailed theory of thalamic and cortical microcircuits for predictive visual inference. [Preprint]
  Science Advances, 2025. Dileep George, Miguel Lázaro-Gredilla, Wolfgang Lehrach, Antoine Dedieu, Guangyao Zhou.
• Subset Selection with Shrinkage: Sparse Linear Modeling when the SNR is low. [Preprint] [Code]
  Operations Research, 2023. Rahul Mazumder, Peter Radchenko, Antoine Dedieu.
• Learning attention-controllable border-ownership for objectness inference and binding. [Preprint]
  Antoine Dedieu, Rajeev V. Rikhye, Miguel Lázaro-Gredilla, Dileep George.
• Solving L1-regularized SVMs and related linear programs: Revisiting the effectiveness of Column and Constraint Generation. [Preprint] [Code]
  Journal of Machine Learning Research, 2022. Antoine Dedieu, Rahul Mazumder, Haoyue Wang
• Learning Sparse Classifiers: Continuous and Mixed Integer Optimization Perspectives. [Preprint] [Code]
  Journal of Machine Learning Research, 2021. Antoine Dedieu, Hussein Hazimeh, Rahul Mazumder.
• Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. [Preprint]
  Nature Communications, 2021. Dileep George, Rajeev Rikhye, Nishad Gothoskar, Swaroop Guntupalli, Antoine Dedieu, Miguel Lázaro-Gredilla.

Other software

• MEQ-JAX: Magnetic EQuilibrium Solver in JAX. [Code]
  David Pfau, Antoine Dedieu, Joel Jennings, Wolfgang Lehrach, Norman Casagrande, Matthias Lochbrunner, Tomek Nanowski, Craig Donner, Federico Felici.

Thesis

• Sparse learning: statistical and optimization perspectives.
  MIT Libraries. Antoine Dedieu.