NVIDIA / transformer-ls

Official PyTorch Implementation of Long-Short Transformer (NeurIPS 2021).

Long-Short Transformer (Transformer-LS)

This repository hosts the code and models for the paper:

Long-Short Transformer: Efficient Transformers for Language and Vision

Updates

• December 6, 2021: Release the code for autoregressive language modeling
• July 23, 2021: Release the code and models for ImageNet classification and Long-Range Arena

Architecture

Long-short Transformer substitutes the full self attention of the original Transformer models with an efficient attention that considers both long-range and short-term correlations. Each query attends to tokens from the segment-wise sliding window to capture short-term correlations, and the dynamically projected features to capture long-range correlations. To align the norms of the original and projected feature vectors and improve the efficacy of the aggregation, we normalize the original and project feature vectors with two sets of Layer Normalizations.

Tasks

• >>> Transformer-LS for ImageNet classification
• >>> Transformer-LS for Long Range Areana
• >>> Transformer-LS for autoregressive language modeling