Blog

The latest from Google Research

Reformer: The Efficient Transformer

Thursday, January 16, 2020

Posted by Nikita Kitaev, Student Researcher, UC Berkeley and Łukasz Kaiser, Research Scientist, Google Researchlong short-term memorytranslate sentence-by-sentenceTransformer modelentire Wikipedia articlesgenerate musicimagesattentionthousands of layersReformerlocality-sensitive-hashingreversible residual layersThe Attention Problemhash functionchunks

Locality-sensitive-hashing: Reformer takes in an input sequence of keys, where each key is a vector representing individual words (or pixels, in the case of images) in the first layer and larger contexts in subsequent layers. LSH is applied to the sequence, after which the keys are sorted by their hash and chunked. Attention is applied only within a single chunk and its immediate neighbors.

The Memory Problemgradient descentactivationsreversible layers

Reversible layers: (A) In a standard residual network, the activations from each layer are used to update the inputs into the next layer. (B) In a reversible network, two sets of activations are maintained, only one of which is updated after each layer. (C) This approach enables running the network in reverse in order to recover all intermediate values.

Applications of Reformerthis colab

Top: Image fragments used as input to Reformer. Bottom: “Completed” full-frame images. Original images are from the Imagenet64 dataset.

Crime and Punishmentthis colabConclusionresearch in the openeven longer sequencesimprove handling of positional encodingsReformer paperour codethis colabchat with usAcknowledgementsThis research was conducted by Nikita Kitaev, Łukasz Kaiser and Anselm Levskaya. Additional thanks go to Afroz Mohiuddin, Jonni Kanerva and Piotr Kozakowski for their work on Trax and to the whole JAX team for their support.