← Search

Taco Cohen

26 accepted papers

2026

Efficient RL Training for LLMs with Experience Replay

ICML 2026poster

While Experience Replay—the practice of storing rollouts and reusing them multiple times during training—is a foundational technique in general RL, it remains largely unexplored in LLM post-training due to the prevailing belief that fresh, on-policy data is essential for high performance. In this wo…

Cited by 0SourceScholar
2025

RLEF: Grounding Code LLMs in Execution Feedback with Reinforcement Learning

ICML 2025spotlight

Large language models (LLMs) deployed as agents solve user-specified tasks over multiple steps while keeping the required manual engagement to a minimum. Crucially, such LLMs need to ground their generations in any feedback obtained to reliably achieve the desired outcomes. We propose an end-to-end…

Cited by 25SourcePDFScholar
2025

The KoLMogorov Test: Compression by Code Generation

ICLR 2025poster

Compression is at the heart of intelligence. A theoretically optimal way to compress any sequence of data is to find the shortest program that outputs that sequence and then halts. However, such Kolmogorov compression is uncomputable, and code generating LLMs struggle to approximate this theoretical…

Cited by 1SourcePDFScholar
2025

UMA: A Family of Universal Models for Atoms

NeurIPS 2025spotlight

The ability to quickly and accurately compute properties from atomic simulations is critical for advancing a large number of applications in chemistry and materials science including drug discovery, energy storage, and semiconductor manufacturing. To address this need, we present a family of Univers…

Cited by 0SourceScholar
2025

What Makes Large Language Models Reason in (Multi-Turn) Code Generation?

ICLR 2025poster

Prompting techniques such as chain-of-thought have established themselves as a popular vehicle for improving the outputs of large language models (LLMs). For code generation, however, their exact mechanics and efficacy are under-explored using unified metrics and benchmarks. We thus investigate the…

Cited by 7SourcePDFScholar
2024

CodeIt: Self-Improving Language Models with Prioritized Hindsight Replay

ICML 2024poster

Large language models are increasingly solving tasks that are commonly believed to require human-level reasoning ability. However, these models still perform very poorly on benchmarks of general intelligence such as the Abstraction and Reasoning Corpus (ARC). In this paper, we approach the ARC as a…

2024

Euclidean, Projective, Conformal: Choosing a Geometric Algebra for Equivariant Transformers

AISTATS 2024poster

The Geometric Algebra Transformer (GATr) is a versatile architecture for geometric deep learning based on projective geometric algebra. We generalize this architecture into a blueprint that allows one to construct a scalable transformer architecture given any geometric (or Clifford) algebra. We stud…

Cited by 12SourcePDFScholar
2023

BISCUIT: Causal Representation Learning from Binary Interactions

UAI 2023poster

Identifying the causal variables of an environment and how to intervene on them is of core value in applications such as robotics and embodied AI. While an agent can commonly interact with the environment and may implicitly perturb the behavior of some of these causal variables, often the targets it…

2023

Causal Representation Learning for Instantaneous and Temporal Effects in Interactive Systems

ICLR 2023poster

Causal representation learning is the task of identifying the underlying causal variables and their relations from high-dimensional observations, such as images. Recent work has shown that one can reconstruct the causal variables from temporal sequences of observations under the assumption that ther…

2023

EDGI: Equivariant Diffusion for Planning with Embodied Agents

NeurIPS 2023poster

Embodied agents operate in a structured world, often solving tasks with spatial, temporal, and permutation symmetries. Most algorithms for planning and model-based reinforcement learning (MBRL) do not take this rich geometric structure into account, leading to sample inefficiency and poor generaliza…

Cited by 35SourcePDFScholar
2023

On the Expressive Power of Geometric Graph Neural Networks

ICML 2023poster

The expressive power of Graph Neural Networks (GNNs) has been studied extensively through the Weisfeiler-Leman (WL) graph isomorphism test. However, standard GNNs and the WL framework are inapplicable for geometric graphs embedded in Euclidean space, such as biomolecules, materials, and other physic…

2022

CITRIS: Causal Identifiability from Temporal Intervened Sequences

ICML 2022spotlight

Understanding the latent causal factors of a dynamical system from visual observations is considered a crucial step towards agents reasoning in complex environments. In this paper, we propose CITRIS, a variational autoencoder framework that learns causal representations from temporal sequences of im…

2022

On the symmetries of the synchronization problem in Cryo-EM: Multi-Frequency Vector Diffusion Maps on the Projective Plane

NeurIPS 2022accept

Cryo-Electron Microscopy (Cryo-EM) is an important imaging method which allows high-resolution reconstruction of the 3D structures of biomolecules. It produces highly noisy 2D images by projecting a molecule's 3D density from random viewing directions. Because the projection directions are unknown,…

Cited by 3SourcePDFScholar
2021

Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs

ICLR 2021spotlight

A common approach to define convolutions on meshes is to interpret them as a graph and apply graph convolutional networks (GCNs). Such GCNs utilize isotropic kernels and are therefore insensitive to the relative orientation of vertices and thus to the geometry of the mesh as a whole. We propose Gau…

2021

Overfitting for Fun and Profit: Instance-Adaptive Data Compression

ICLR 2021poster

Neural data compression has been shown to outperform classical methods in terms of $RD$ performance, with results still improving rapidly. At a high level, neural compression is based on an autoencoder that tries to reconstruct the input instance from a (quantized) latent representation, coupled wit…

Cited by 47SourcePDFScholar
2020

Low Bias Low Variance Gradient Estimates for Boolean Stochastic Networks

ICML 2020poster

Stochastic neural networks with discrete random variables are an important class of models for their expressiveness and interpretability. Since direct differentiation and backpropagation is not possible, Monte Carlo gradient estimation techniques are a popular alternative. Efficient stochastic gradi…

Cited by 17SourcePDFScholar
2019

Gauge Equivariant Convolutional Networks and the Icosahedral CNN

ICML 2019oral

The principle of equivariance to symmetry transformations enables a theoretically grounded approach to neural network architecture design. Equivariant networks have shown excellent performance and data efficiency on vision and medical imaging problems that exhibit symmetries. Here we show how this p…

Cited by 512SourcePDFScholar