The Reformer - Pushing the limits of language modeling

Summary

Standard transformer models hit memory bottlenecks on long sequence modeling tasks due to the quadratic asymptotic memory complexity of global self-attention and oversized positional embedding matrices. The Reformer architecture overcomes these constraints to train sequences of up to half a million tokens using under 8GB of RAM. It re-engineers transformer operations using local and Locality Sensitive Hashing self-attention, chunked feed forward layers, reversible residual layers, and axial positional encodings. In empirical benchmarks using google/reformer-crime-and-punishment, axial positional encodings reduce the model parameter count from over 136 million to approximately 2.58 million by factorizing the positional dimensions. This architectural change cuts inference memory consumption from 959 MB down to 447 MB for evaluated benchmark workloads.

Takeaways

1. 1

   Global self-attention generates an O(n^2) memory bottleneck in the inner dot-product tensor, restricting models like bert-base-uncased to a maximum position embedding size of 512 tokens.

2. 2

   Local self-attention partitions inputs into chunks of length l_c and adds overlap from preceding and following chunks to maintain access to immediate neighboring token context.

3. 3

   In benchmarks on half-million token sequences, axial positional encodings cut total model parameters from 136,572,416 to 2,584,064 and reduce inference memory from 959 MB to 447 MB at batch size 8.