← Search

Roger Baker Grosse

22 accepted papers

2026

Understanding and improving Shampoo and SOAP via Kullback-Leibler Minimization

ICLR 2026poster

Shampoo and its efficient, Adam-stabilized variant SOAP, employ structured second-moment estimation and have received growing attention for their effectiveness. In practice, Shampoo requires step-size grafting with Adam to achieve competitive performance. SOAP mitigates this by applying Adam in Sha…

Cited by 0SourcecodeScholar
2025

Better Training Data Attribution via Better Inverse Hessian-Vector Products

NeurIPS 2025poster

Training data attribution (TDA) provides insights into which training data is responsible for a learned model behavior. Gradient-based TDA methods such as influence functions and unrolled differentiation both involve a computation that resembles an inverse Hessian-vector product (iHVP), which is dif…

Cited by 0SourceScholar
2025

Distributional Training Data Attribution: What do Influence Functions Sample?

NeurIPS 2025spotlight

Randomness is an unavoidable part of training deep learning models, yet something that traditional training data attribution algorithms fail to rigorously account for. They ignore the fact that, due to stochasticity in the initialisation and batching, training on the same dataset can yield different…

Cited by 0SourceScholar
2025

Reducing the Probability of Undesirable Outputs in Language Models Using Probabilistic Inference

NeurIPS 2025poster

Reinforcement learning (RL) has become a predominant technique to align language models (LMs) with human preferences or promote outputs which are deemed to be desirable by a given reward function. Standard RL approaches optimize average reward, while methods explicitly focused on reducing the probab…

Cited by 0SourceScholar
2025

What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions

NeurIPS 2025poster

Large language models (LLMs) are trained on a vast amount of human-written data, but data providers often remain uncredited. In response to this issue, data valuation (or data attribution), which quantifies the contribution or value of each data to the model output, has been discussed as a potential…

Cited by 0SourcecodeScholar
2024

Connecting the Dots: LLMs can Infer and Verbalize Latent Structure from Disparate Training Data

NeurIPS 2024poster

One way to address safety risks from large language models (LLMs) is to censor dangerous knowledge from their training data. While this removes the explicit information, implicit information can remain scattered across various training documents. Could an LLM infer the censored knowledge by piecing…

2024

Measuring Stochastic Data Complexity with Boltzmann Influence Functions

ICML 2024poster

Estimating the uncertainty of a model’s prediction on a test point is a crucial part of ensuring reliability and calibration under distribution shifts.A minimum description length approach to this problem uses the predictive normalized maximum likelihood (pNML) distribution, which considers every po…

Cited by 1SourcePDFScholar
2024

Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo

ICML 2024oral

Numerous capability and safety techniques of Large Language Models (LLMs), including RLHF, automated red-teaming, prompt engineering, and infilling, can be cast as sampling from an unnormalized target distribution defined by a given reward or potential function over the full sequence. In this work,…

2024

REFACTOR: Learning to Extract Theorems from Proofs

ICLR 2024poster

Human mathematicians are often good at recognizing modular and reusable theorems that make complex mathematical results within reach. In this paper, we propose a novel method called theoREm-from-prooF extrACTOR (REFACTOR) for training neural networks to mimic this ability in formal mathematical theo…

2023

Efficient Parametric Approximations of Neural Network Function Space Distance

ICML 2023poster

It is often useful to compactly summarize important properties of model parameters and training data so that they can be used later without storing and/or iterating over the entire dataset. As a specific case, we consider estimating the Function Space Distance (FSD) over a training set, i.e. the ave…

Cited by 6SourcePDFScholar
2023

Multi-Rate VAE: Train Once, Get the Full Rate-Distortion Curve

ICLR 2023top-5%

Variational autoencoders (VAEs) are powerful tools for learning latent representations of data used in a wide range of applications. In practice, VAEs usually require multiple training rounds to choose the amount of information the latent variable should retain. This trade-off between the reconstruc…

Cited by 21SourcePDFScholar
2023

Similarity-based cooperative equilibrium

NeurIPS 2023poster

As machine learning agents act more autonomously in the world, they will increasingly interact with each other. Unfortunately, in many social dilemmas like the one-shot Prisoner’s Dilemma, standard game theory predicts that ML agents will fail to cooperate with each other. Prior work has shown that…

Cited by 7SourcePDFScholar
2022

If Influence Functions are the Answer, Then What is the Question?

NeurIPS 2022accept

Influence functions efficiently estimate the effect of removing a single training data point on a model's learned parameters. While influence estimates align well with leave-one-out retraining for linear models, recent works have shown this alignment is often poor in neural networks. In this work, w…

2022

Improving Mutual Information Estimation with Annealed and Energy-Based Bounds

ICLR 2022poster

Mutual information (MI) is a fundamental quantity in information theory and machine learning. However, direct estimation of MI is intractable, even if the true joint probability density for the variables of interest is known, as it involves estimating a potentially high-dimensional log partition fun…

2022

Path Independent Equilibrium Models Can Better Exploit Test-Time Computation

NeurIPS 2022accept

Designing networks capable of attaining better performance with an increased inference budget is important to facilitate generalization to harder problem instances. Recent efforts have shown promising results in this direction by making use of depth-wise recurrent networks. In this work, we reproduc…

Cited by 26SourcePDFScholar
2022

Proximal Learning With Opponent-Learning Awareness

NeurIPS 2022accept

Learning With Opponent-Learning Awareness (LOLA) (Foerster et al. [2018a]) is a multi-agent reinforcement learning algorithm that typically learns reciprocity-based cooperation in partially competitive environments. However, LOLA often fails to learn such behaviour on more complex policy spaces para…

2021

Differentiable Annealed Importance Sampling and the Perils of Gradient Noise

NeurIPS 2021poster

Annealed importance sampling (AIS) and related algorithms are highly effective tools for marginal likelihood estimation, but are not fully differentiable due to the use of Metropolis-Hastings correction steps. Differentiability is a desirable property as it would admit the possibility of optimizing…

Cited by 38SourcePDFScholar
2021

INT: An Inequality Benchmark for Evaluating Generalization in Theorem Proving

ICLR 2021poster

In learning-assisted theorem proving, one of the most critical challenges is to generalize to theorems unlike those seen at training time. In this paper, we introduce INT, an INequality Theorem proving benchmark designed to test agents’ generalization ability. INT is based on a theorem generator, wh…

2021

When does preconditioning help or hurt generalization?

ICLR 2021poster

While second order optimizers such as natural gradient descent (NGD) often speed up optimization, their effect on generalization has been called into question. This work presents a more nuanced view on how the \textit{implicit bias} of optimizers affects the comparison of generalization properties.…

Cited by 50SourcePDFScholar