← Search

Jimmy Ba

47 accepted papers

2024

Identifying the Risks of LM Agents with an LM-Emulated Sandbox

ICLR 2024spotlight

Recent advances in Language Model (LM) agents and tool use, exemplified by applications like ChatGPT Plugins, enable a rich set of capabilities but also amplify potential risks—such as leaking private data or causing financial losses. Identifying these risks is labor-intensive, necessitating impleme…

2024

OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text

ICLR 2024poster

There is growing evidence that pretraining on high quality, carefully thought-out tokens such as code or mathematics plays an important role in improving the reasoning abilities of large language models. For example, Minerva, a PaLM model finetuned on billions of tokens of mathematical documents fro…

2024

The WMDP Benchmark: Measuring and Reducing Malicious Use with Unlearning

ICML 2024poster

The White House Executive Order on Artificial Intelligence highlights the risks of large language models (LLMs) empowering malicious actors in developing biological, cyber, and chemical weapons. To measure these risks, government institutions and major AI labs are developing evaluations for hazardou…

Cited by 145SourcePDFScholar
2023

AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback

NeurIPS 2023spotlight

Large language models (LLMs) such as ChatGPT have seen widespread adoption due to their ability to follow user instructions well. Developing these LLMs involves a complex yet poorly understood workflow requiring training with human feedback. Replicating and understanding this instruction-following p…

Cited by 523SourcePDFScholar
2023

Large Language Models are Human-Level Prompt Engineers

ICLR 2023poster

By conditioning on natural language instructions, large language models (LLMs) have displayed impressive capabilities as general-purpose computers. However, task performance depends significantly on the quality of the prompt used to steer the model, and most effective prompts have been handcrafted b…

2023

Learning in the Presence of Low-dimensional Structure: A Spiked Random Matrix Perspective

NeurIPS 2023poster

We consider the learning of a single-index target function $f_*: \mathbb{R}^d\to\mathbb{R}$ under spiked covariance data: $$f_*(\boldsymbol{x}) = \textstyle\sigma_*(\frac{1}{\sqrt{1+\theta}}\langle\boldsymbol{x},\boldsymbol{\mu}\rangle), ~~ \boldsymbol{x}\overset{\small\mathrm{i.i.d.}}{\sim}\mathca…

Cited by 39SourcePDFScholar
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

Residual Prompt Tuning: improving prompt tuning with residual reparameterization

ACL 2023findings

Prompt tuning is one of the successful approaches for parameter-efficient tuning of pre-trained language models. Despite being arguably the most parameter-efficient (tuned soft prompts constitute <0.1% of total parameters), it typically performs worse than other efficient tuning methods and is quite…

2023

STEVE-1: A Generative Model for Text-to-Behavior in Minecraft

NeurIPS 2023spotlight

Constructing AI models that respond to text instructions is challenging, especially for sequential decision-making tasks. This work introduces a methodology, inspired by unCLIP, for instruction-tuning generative models of behavior without relying on a large dataset of instruction-labeled trajectorie…

Cited by 64SourcePDFScholar
2023

TR0N: Translator Networks for 0-Shot Plug-and-Play Conditional Generation

ICML 2023poster

We propose TR0N, a highly general framework to turn pre-trained unconditional generative models, such as GANs and VAEs, into conditional models. The conditioning can be highly arbitrary, and requires only a pre-trained auxiliary model. For example, we show how to turn unconditional models into class…

2022

High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

NeurIPS 2022accept

We study the first gradient descent step on the first-layer parameters $\boldsymbol{W}$ in a two-layer neural network: $f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^\top\sigma(\boldsymbol{W}^\top\boldsymbol{x})$, where $\boldsymbol{W}\in\mathbb{R}^{d\times N}, \boldsymbol{a}\in\mathbb{R}^{N}…

Cited by 179SourcePDFScholar
2022

Understanding the Variance Collapse of SVGD in High Dimensions

ICLR 2022poster

Stein variational gradient descent (SVGD) is a deterministic inference algorithm that evolves a set of particles to fit a target distribution. Despite its computational efficiency, SVGD often underestimates the variance of the target distribution in high dimensions. In this work we attempt to explai…

Cited by 31SourcePDFScholar
2022

You Can’t Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments

NeurIPS 2022accept

Recently, methods such as Decision Transformer that reduce reinforcement learning to a prediction task and solve it via supervised learning (RvS) have become popular due to their simplicity, robustness to hyperparameters, and strong overall performance on offline RL tasks. However, simply conditioni…

Cited by 73SourcePDFScholar
2021

Efficient Statistical Tests: A Neural Tangent Kernel Approach

ICML 2021spotlight

For machine learning models to make reliable predictions in deployment, one needs to ensure the previously unknown test samples need to be sufficiently similar to the training data. The commonly used shift-invariant kernels do not have the compositionality and fail to capture invariances in high-dim…

2021

How does a Neural Network's Architecture Impact its Robustness to Noisy Labels?

NeurIPS 2021poster

Noisy labels are inevitable in large real-world datasets. In this work, we explore an area understudied by previous works --- how the network's architecture impacts its robustness to noisy labels. We provide a formal framework connecting the robustness of a network to the alignments between its arch…

Cited by 27SourcePDFScholar
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

LIME: Learning Inductive Bias for Primitives of Mathematical Reasoning

ICML 2021spotlight

While designing inductive bias in neural architectures has been widely studied, we hypothesize that transformer networks are flexible enough to learn inductive bias from suitable generic tasks. Here, we replace architecture engineering by encoding inductive bias in the form of datasets. Inspired by…

2021

Learning Domain Invariant Representations in Goal-conditioned Block MDPs

NeurIPS 2021poster

Deep Reinforcement Learning (RL) is successful in solving many complex Markov Decision Processes (MDPs) problems. However, agents often face unanticipated environmental changes after deployment in the real world. These changes are often spurious and unrelated to the underlying problem, such as backg…

2021

Mastering Atari with Discrete World Models

ICLR 2021poster

Intelligent agents need to generalize from past experience to achieve goals in complex environments. World models facilitate such generalization and allow learning behaviors from imagined outcomes to increase sample-efficiency. While learning world models from image inputs has recently become feasib…

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
2020

An Empirical Study of Stochastic Gradient Descent with Structured Covariance Noise

AISTATS 2020poster

The choice of batch-size in a stochastic optimization algorithm plays a substantial role for both optimization and generalization. Increasing the batch-size used typically improves optimization but degrades generalization. To address the problem of improving generalization while maintaining optimal…

Cited by 15SourcePDFScholar
2020

An Inductive Bias for Distances: Neural Nets that Respect the Triangle Inequality

ICLR 2020poster

Distances are pervasive in machine learning. They serve as similarity measures, loss functions, and learning targets; it is said that a good distance measure solves a task. When defining distances, the triangle inequality has proven to be a useful constraint, both theoretically---to prove convergenc…

Cited by 28SourcecodeScholar
2020

BatchEnsemble: an Alternative Approach to Efficient Ensemble and Lifelong Learning

ICLR 2020poster

Ensembles, where multiple neural networks are trained individually and their predictions are averaged, have been shown to be widely successful for improving both the accuracy and predictive uncertainty of single neural networks. However, an ensemble’s cost for both training and testing increases lin…

Cited by 550SourcecodeScholar
2020

Dream to Control: Learning Behaviors by Latent Imagination

ICLR 2020spotlight

Learned world models summarize an agent's experience to facilitate learning complex behaviors. While learning world models from high-dimensional sensory inputs is becoming feasible through deep learning, there are many potential ways for deriving behaviors from them. We present Dreamer, a reinforcem…

Cited by 1635SourcecodeScholar
2020

Generalization of Two-layer Neural Networks: An Asymptotic Viewpoint

ICLR 2020spotlight

This paper investigates the generalization properties of two-layer neural networks in high-dimensions, i.e. when the number of samples $n$, features $d$, and neurons $h$ tend to infinity at the same rate. Specifically, we derive the exact population risk of the unregularized least squares regression…

Cited by 97SourceScholar
2020

Improving Transformer Optimization Through Better Initialization

ICML 2020poster

The Transformer architecture has achieved considerable success recently; the key component of the Transformer is the attention layer that enables the model to focus on important regions within an input sequence. Gradient optimization with attention layers can be notoriously difficult requiring trick…

2020

Maximum Entropy Gain Exploration for Long Horizon Multi-goal Reinforcement Learning

ICML 2020poster

What goals should a multi-goal reinforcement learning agent pursue during training in long-horizon tasks? When the desired (test time) goal distribution is too distant to offer a useful learning signal, we argue that the agent should not pursue unobtainable goals. Instead, it should set its own intr…

2019

Lookahead Optimizer: k steps forward, 1 step back

NeurIPS 2019poster

The vast majority of successful deep neural networks are trained using variants of stochastic gradient descent (SGD) algorithms. Recent attempts to improve SGD can be broadly categorized into two approaches: (1) adaptive learning rate schemes, such as AdaGrad and Adam and (2) accelerated schemes, su…

2019

Neural Graph Evolution: Towards Efficient Automatic Robot Design

ICLR 2019poster

Despite the recent successes in robotic locomotion control, the design of robot relies heavily on human engineering. Automatic robot design has been a long studied subject, but the recent progress has been slowed due to the large combinatorial search space and the difficulty in evaluating the found…

2018

Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches

ICLR 2018poster

Stochastic neural net weights are used in a variety of contexts, including regularization, Bayesian neural nets, exploration in reinforcement learning, and evolution strategies. Unfortunately, due to the large number of weights, all the examples in a mini-batch typically share the same weight pertur…

Cited by 418SourcePDFScholar
2018

NerveNet: Learning Structured Policy with Graph Neural Networks

ICLR 2018poster

We address the problem of learning structured policies for continuous control. In traditional reinforcement learning, policies of agents are learned by MLPs which take the concatenation of all observations from the environment as input for predicting actions. In this work, we propose NerveNet to exp…

Cited by 326SourcePDFScholar
2018

On the Convergence and Robustness of Training GANs with Regularized Optimal Transport

NeurIPS 2018poster

Generative Adversarial Networks (GANs) are one of the most practical methods for learning data distributions. A popular GAN formulation is based on the use of Wasserstein distance as a metric between probability distributions. Unfortunately, minimizing the Wasserstein distance between the data distr…

Cited by 173SourcePDFScholar
2017

Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation

NeurIPS 2017spotlight

In this work, we propose to apply trust region optimization to deep reinforcement learning using a recently proposed Kronecker-factored approximation to the curvature. We extend the framework of natural policy gradient and propose to optimize both the actor and the critic using Kronecker-factored ap…

2016

Using Fast Weights to Attend to the Recent Past

NeurIPS 2016oral

Until recently, research on artificial neural networks was largely restricted to systems with only two types of variable: Neural activities that represent the current or recent input and weights that learn to capture regularities among inputs, outputs and payoffs. There is no good reason for this re…

Cited by 317SourcePDFScholar
2015

Learning Wake-Sleep Recurrent Attention Models

NeurIPS 2015spotlight

Despite their success, convolutional neural networks are computationally expensive because they must examine all image locations. Stochastic attention-based models have been shown to improve computational efficiency at test time, but they remain difficult to train because of intractable posterior in…

Cited by 77SourcePDFScholar
2015

Show, Attend and Tell: Neural Image Caption Generation with Visual Attention

ICML 2015poster

Inspired by recent work in machine translation and object detection, we introduce an attention based model that automatically learns to describe the content of images. We describe how we can train this model in a deterministic manner using standard backpropagation techniques and stochastically by ma…

Cited by 13573SourcePDFScholar