PaperlayerView exact revision on arXiv

Source-linked AI summary

TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?

Michael S. Ryoo, AJ Piergiovanni, Anurag Arnab, Mostafa Dehghani, Anelia Angelova

arXiv:2106.11297v4cs.CVcs.LG

TL;DR

TokenLearner addresses the cost of processing many densely sampled visual tokens by adaptively extracting a small set of important tokens from images and video frames. It improves or maintains recognition performance across image and video benchmarks while reducing computation, though the evaluated video dataset uses Creative Commons-licensed videos.

  • Problem

    Vision Transformers often require too many tokens, while image and video understanding still requires efficient selection of important spatial and temporal information.

  • Method

    TokenLearner is a learnable module that adaptively generates 8-16 tokens from image-like tensors and can be inserted into Vision Transformer architectures.

  • Results

    TokenLearner maintains or improves classification accuracy while reducing computation across image classification and challenging video understanding tasks, including new state-of-the-art results on Charades and AViD.

  • Takeaways & Limitations

    Adaptive tokenization provides a more computationally efficient visual representation while preserving or improving recognition performance across image and video models.

  • Takeaways & Limitations

    The video results use videos with Creative Commons licenses.

Abstract

from arXiv · show

In this paper, we introduce a novel visual representation learning which relies on a handful of adaptively learned tokens, and which is applicable to both image and video understanding tasks. Instead of relying on hand-designed splitting strategies to obtain visual tokens and processing a large number of densely sampled patches for attention, our approach learns to mine important tokens in visual data. This results in efficiently and effectively finding a few important visual tokens and enables modeling of pairwise attention between such tokens, over a longer temporal horizon for videos, or the spatial content in images. Our experiments demonstrate strong performance on several challenging benchmarks for both image and video recognition tasks. Importantly, due to our tokens being adaptive, we accomplish competitive results at significantly reduced compute amount. We obtain comparable results to the state-of-the-arts on ImageNet while being computationally more efficient. We also confirm the effectiveness of the approach on multiple video datasets, including Kinetics-400, Kinetics-600, Charades, and AViD. The code is available at: https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner

1 INTRODUCTION

TokenLearner addresses the computational burden of densely tokenized Vision Transformers by adaptively generating a small set of visual tokens for images and videos. It reduces computation while maintaining or improving recognition performance across image and video tasks.

  • The tokenization challenge: Vision Transformers can require 1,024 tokens for a 512x512 image and tens of thousands for videos, making quadratic attention costly.Tokens are processed at every layer, so the computational and memory burden grows rapidly with token count.
  • Adaptive token learning: TokenLearner adaptively generates a smaller number of tokens from image-like tensors instead of relying on uniformly split patches.It uses learned spatial attention to emphasize important regions, weight the input, and spatially pool the result into tokens.
  • Efficiency and recognition: Using 8–16 intermediate tokens instead of 200∼500 significantly reduces Vision Transformer computation while maintaining or increasing classification accuracy.The approach reduces total FLOPS and can be inserted at different locations in the model.
  • Video and image evaluation: Video experiments report improved performance over state-of-the-art methods on three challenging datasets using only 8–16 tokens per frame.The paper also reports comparable ImageNet and ImageNet ReaL performance with meaningfully reduced computation.

2 TOKENLEARNER MODULES FOR ADAPTIVE TO-

The paper replaces fixed spatial or spatio-temporal patch tokenization with adaptive token selection and adds TokenFuser for token mixing and spatial remapping. Together, these modules provide compact representations that can be processed efficiently and restored to spatial resolution when needed.

  • Fixed tokenization: Vision and video Transformers usually tokenize inputs by cutting images or videos into regular spatial or spatio-temporal patches.Examples include 16x16 image patches and 16x16x2 video cubes.
  • TokenLearner: For an input tensor with spatial dimensions H × W, temporal dimension T, and channels C, TokenLearner generates S token vectors per frame.For images, T = 1; for videos, tokens from all frames form an output of size ST × C.
  • TokenLearner: TokenLearner dynamically selects informative spatial locations from each frame and produces a small set of adaptive tokens rather than fixed input splits.Different tokens can be mined per frame, allowing video token interactions to be modeled across time.
  • TokenLearner: Each TokenLearner token is formed by an input-conditioned spatial weight map, element-wise multiplication with the input, and spatial global average pooling.The resulting tokens are gathered into an S × C representation for each frame.
  • Integration with Transformers: The learned tokens replace large token sets in subsequent Transformer layers, such as MHSA, reducing computation because attention cost is quadratic in token count.The module can be inserted at different locations within the network.
  • TokenFuser: TokenFuser first applies a learned linear transformation over tokens rather than channels, preserving the ST × C tensor size.This token-wise operation is intended to capture patterns formed by the tokens.
  • TokenFuser: TokenFuser then processes each temporal slice and remaps its S × C tokens to H × W × C by learning location-specific combinations.This supports fusing token information and recovering the original spatial representation when required.

3 EXPERIMENTS WITH IMAGES

The image experiments insert TokenLearner into ViT architectures, optionally with TokenFuser, and evaluate accuracy, few-shot performance, computation, and efficiency across model configurations. TokenLearner generally preserves or improves accuracy while substantially reducing computation, including on larger models.

  • Experimental setup: The experiments compare standard ViT models with TokenLearner inserted at different locations, with or without TokenFuser, using 8 or 16 learned tokens.The evaluated backbones include ViT-B/16, ViT-L/16, and variants with different input resolutions and patch sizes.
  • TokenLearner location: Placing TokenLearner in the middle of the network achieves almost identical ImageNet few-shot accuracy while cutting computation by almost half.Placing it after three-quarters of the network achieves superior performance to the model without TokenLearner while remaining faster.
  • ImageNet fine-tuning: On ImageNet fine-tuning, TokenLearner maintains ViT accuracy while reducing computation by almost half, and additional layers can use the saved computation to improve performance.The efficiency-accuracy trade-off improves across the evaluated smaller ViT-S and ViT-B models.
  • Larger models: TokenLearner saves about half the computation without sacrificing accuracy in larger models, can outperform base models when saved compute is redeployed, and performs comparably to much larger ViT models.For ViT-L/16 versus L/16 + TokenLearner, measured throughput is 1400 versus 2000 images per second; an early TokenLearner placement can outperform ViT B/16 at around half its runtime.
  • Ablations: Compared with alternative token-reduction designs, pooling loses accuracy relative to the base model, while TokenLearner performs slightly better than the base model at lower computation.Unpooling performs worse, whereas reprojection is comparable to TokenFuser but requires more FLOPS.

4 TOKENLEARNER FOR VIDEOS

For videos, TokenLearner generates a small set of adaptive tokens per frame, processes their joint space-time relations with Transformers, and can remap them with TokenFuser. This reduces the tokens handled by later layers while supporting efficient video representation learning.

  • Video TokenLearner framework: The video implementation applies the same TokenLearner and TokenFuser ideas as the image model, differing by generating multiple per-frame token sets before stacking them.
  • Video TokenLearner framework: TokenLearner generates S adaptive tokens per frame, which are stacked into ST tokens before joint Transformer processing.TokenFuser can optionally remap the learned tokens back to the representation tensor shape.
  • Video TokenLearner framework: The combined video module repeats TokenLearner, Transformer, and optionally TokenFuser, while architectures without TokenFuser repeat only Transformer layers after TokenLearner.

5 EXPERIMENTS WITH VIDEOS: TOKENLEARNER WITH VIDEO VISION TRANSFORMER

The authors insert TokenLearner into ViViT, a video Vision Transformer based on spatio-temporal patch tokens, and evaluate it on Kinetics-400 and Kinetics-600. The resulting models improve the accuracy–compute trade-off relative to comparable ViViT systems and extend state-of-the-art results on Kinetics.

  • Architecture and setup: TokenLearner is inserted directly into ViViT, which uses spatio-temporal video patches as Transformer tokens.The models are evaluated under the same general ViViT architecture and initialization framework.
  • Architecture and setup: The experiments use Kinetics-400 and Kinetics-600, with standard prior-work settings and validation accuracy as the evaluation metric.The datasets contain approximately 240k and 390k training samples, respectively, with clips of about 10 seconds.
  • Architecture and setup: The models are initialized from JFT-pretrained weights, adapted for 16x16x2 video patches, and then fine-tuned on Kinetics.
  • Efficiency: TokenLearner ViViT models use 8 or 16 tokens in added layers, keeping their computation increase minimal while reporting both classification accuracy and FLOPS.
  • Results: On Kinetics-400, TokenLearner improves classification accuracy while reducing compute compared with directly comparable ViViT baselines.
  • Results: On Kinetics-600, TokenLearner extends the state-of-the-art while remaining computationally efficient.

6 EXPERIMENTS WITH VIDEOS: TOKENLEARNER WITH BOTTLENECK TRANSFORMER

The authors combine TokenLearner with a bottleneck Transformer and an X(2+1)D backbone for longer video actions. On Charades and AViD, the approach achieves strong recognition results, while comparisons indicate that adaptive tokenization can outperform heavier fixed space-time tokenization.

  • Motivation and architecture: The bottleneck-Transformer experiment targets Charades videos, whose longer actions require modeling longer-term temporal information across multiple frames.TokenLearner is used to compute multiple temporal tokens efficiently over many frames.
  • Motivation and architecture: The X(2+1)D backbone replaces 3D convolutions with separate 2D spatial and 1D temporal convolutions, while removing squeeze-and-excitation layers and swish activations.
  • Motivation and architecture: In the final residual design, TokenLearner replaces the 1D temporal convolution while the 2D 3 × 3 convolution is retained.The Transformer layer uses a Vector Transformer rather than standard multi-head self-attention.
  • Datasets and setup: The Charades setup uses 64-frame inputs and S = 8 learned tokens, matching the goal of capturing longer-term temporal information efficiently.Charades contains 157 activity classes and videos averaging 30 seconds.
  • Results: 66.3% mAP on Charades establishes a new state-of-the-art result.
  • Results: On AViD, TokenLearner outperforms prior work and a same-backbone TimeSformer reimplementation while also being more computationally efficient.AViD results are averaged over 887 classes, and the dataset contains geographically diverse video content.
  • Tokenization comparisons: The full joint space-time Transformer uses about eight times more tokens than TokenLearner in the bottleneck setting, yet performs slightly worse.This comparison uses the X(2+1)D backbone and measures computation for the modules rather than the entire network.
  • Tokenization comparisons: On the Charades ablation, spatial-attention token learning reaches 59.6 accuracy versus 58.8 for fixed splitting, 56.6 for direct generation, and 58.6 for pooling and generation.These results support the importance of spatial attention for TokenLearner token formation.

7 RELATED WORK

Related vision Transformers typically tokenize images and videos using fixed spatial or spatio-temporal grids. TokenLearner instead learns fewer tokens from the data and applies the approach across both image and video recognition.

  • Video representation learning: Video understanding requires attention across spatial and temporal domains, increasing the challenge of tokenizing and processing video inputs.
  • Fixed tokenization: Vision Transformers commonly split images into regular-grid patches, while video Transformers use spatial or spatio-temporal cubes as tokens.
  • TokenLearner: TokenLearner learns tokens from data instead of relying on fixed tokenization, producing fewer tokens for more efficient image and video recognition.
  • TokenLearner: The paper reports that even 8x fewer learned tokens, such as 512 versus 4096, can capture the information needed for video representation learning.
  • TokenLearner: Spatial attention maps visualize four of the eight learned tokens and show their input-dependent focus across visual regions.

8 CONCLUSIONS

TokenLearner adaptively tokenizes images and video frames to extract important tokens for recognition, improving efficiency and accuracy across image and video tasks.

  • TokenLearner adaptively extracts important tokens from images and video frames for recognition tasks.
  • The approach finds a few important space-time tokens to model visual representations more efficiently than contemporary approaches.
  • TokenLearner improves accuracy across image classification and challenging video understanding tasks, outperforming prior approaches on many datasets.
  • A remaining challenge is learning full spatio-temporal tokens, since the current method focuses on spatial tokens over frame sequences.

APPENDIX A VECTOR TRANSFORMER: PAIRWISE VECTOR ATTEN-

The Vector Transformer computes channel-aware pairwise attention among TokenLearner’s adaptive tokens, enabling efficient modeling of space-time relations over the supplied temporal range.

  • After TokenLearner generates adaptive tokens, vector attention computes pairwise relations between key-query pairs with one attention matrix per channel.This preserves channel-specific information while modeling pairwise space-time relations.
  • Given tokens representing different video space-time aspects, the Transformer models interactions between them.
  • For each token zi, the output yi considers all tokens zj in the token set Z.
  • The projections fq, fk, and fv produce query, key, and value representations, while γ projects channel dimensions and applies softmax over j.
  • The vector formulation produces an attention tensor of shape ST × ST × d that is multiplied by value representations to obtain Transformer outputs.The tensor preserves channel information rather than reducing attention to a standard QK^T matrix.
  • The temporal range of the global representation depends on the supplied tokens, allowing adaptive tokens to cover more frames and focus on temporal structure.

APPENDIX B IMAGE AND VIDEO CLASSIFICATION TRAINING DE-

The image and video experiments follow established training protocols and dataset splits, with dataset-specific optimization settings and reported infrastructure details.

  • Image classification experiments follow the exact training protocols and hyperparameters of [12].
  • Video experiments follow the exact training protocols and hyperparameters of, using the same Scenic code and hardware for training and evaluation.
  • Kinetics models are trained for 30 epochs with a base learning rate of 0.05 using the Momentum optimizer, following ViViT settings.
  • The experiments use the standard training and testing splits of the datasets.
  • Charades uses a cosine-decay schedule with a base learning rate of 0.8 per TPU core, while AViD uses 0.025 per TPU.
  • A single model takes approximately ∼16 hours to train on 32 TPU v3, with training time bottlenecked by the data pipeline.

APPENDIX C COMPARING MORE MODELS ON FEW-SHOT LEARN-

Few-shot experiments extend TokenLearner comparisons across larger models and insertion locations, showing that early placement can reduce computation while improving image classification accuracy over base models.

  • Figure 11 reports 5-shot classification accuracies on ImageNet and the average across multiple datasets.‘TL’ denotes TokenLearner.
  • Table 14 compares different TokenLearner components on the Charades dataset at 6fps.
  • Early TokenLearner insertion saves computation while achieving superior image classification accuracy to the base models.

APPENDIX D ADDITIONAL ABLATIONS

The appendix evaluates model components and their combinations through ablations in the Bottleneck Transformer setting, identifying a best-performing Conv2D-plus-transformer bottleneck module.

  • APPENDIX D ADDITIONAL ABLATIONS: The ablation removes and adds different model components to evaluate their individual effects and combinations.These experiments use the Bottleneck Transformer setting.
  • APPENDIX D ADDITIONAL ABLATIONS: The study compares the Vector Transformer with an alternative multi-head self-attention module.The Vector Transformer is described in the preceding subsection.
  • APPENDIX D ADDITIONAL ABLATIONS: The Conv2D + transformer module within the bottleneck performed best among the tested ablations.
Loading 2106.11297v4…