← Search

Jack Rae

12 accepted papers

2022

Improving Language Models by Retrieving from Trillions of Tokens

ICML 2022spotlight

We enhance auto-regressive language models by conditioning on document chunks retrieved from a large corpus, based on local similarity with preceding tokens. With a 2 trillion token database, our Retrieval-Enhanced Transformer (RETRO) obtains comparable performance to GPT-3 and Jurassic-1 on the Pil…

2022

Unified Scaling Laws for Routed Language Models

ICML 2022oral

The performance of a language model has been shown to be effectively modeled as a power-law in its parameter count. Here we study the scaling behaviors of Routing Networks: architectures that conditionally use only a subset of their parameters while processing an input. For these models, parameter c…

2020

Multiplicative Interactions and Where to Find Them

ICLR 2020poster

We explore the role of multiplicative interaction as a unifying framework to describe a range of classical and modern neural network architectural motifs, such as gating, attention layers, hypernetworks, and dynamic convolutions amongst others. Multiplicative interaction layers as primitive operatio…

Cited by 154SourceScholar
2020

Stabilizing Transformers for Reinforcement Learning

ICML 2020poster

Owing to their ability to both effectively integrate information over long time horizons and scale to massive amounts of data, self-attention architectures have recently shown breakthrough success in natural language processing (NLP). Harnessing the transformer’s ability to process long time horizon…

2020

Top-KAST: Top-K Always Sparse Training

NeurIPS 2020poster

Sparse neural networks are becoming increasingly important as the field seeks to improve the performance of existing models by scaling them up, while simultaneously trying to reduce power consumption and computational footprint. Unfortunately, most existing methods for inducing performant sparse mod…

Cited by 114SourcePDFScholar
2019

Training Language GANs from Scratch

NeurIPS 2019poster

Generative Adversarial Networks (GANs) enjoy great success at image generation, but have proven difficult to train in the domain of natural language. Challenges with gradient estimation, optimization instability, and mode collapse have lead practitioners to resort to maximum likelihood pre-training,…

2018

Relational recurrent neural networks

NeurIPS 2018poster

Memory-based neural networks model temporal data by leveraging an ability to remember information for long periods. It is unclear, however, whether they also have an ability to perform complex relational reasoning with the information they remember. Here, we first confirm our intuitions that standar…

2016

Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes

NeurIPS 2016poster

Neural networks augmented with external memory have the ability to learn algorithmic solutions to complex tasks. These models appear promising for applications such as language modeling and machine translation. However, they scale poorly in both space and time as the amount of memory grows --- limit…

Cited by 196SourcePDFScholar