TRecViT : Un Transformateur Vidéo Récurent

TRecViT: A Recurrent Video Transformer

December 18, 2024

Auteurs: Viorica Pătrăucean, Xu Owen He, Joseph Heyward, Chuhan Zhang, Mehdi S. M. Sajjadi, George-Cristian Muraru, Artem Zholus, Mahdi Karami, Ross Goroshin, Yutian Chen, Simon Osindero, João Carreira, Razvan Pascanu

cs.AI

Résumé

Nous proposons un nouveau bloc pour la modélisation vidéo. Il repose sur une factorisation temps-espace-canal avec des blocs dédiés pour chaque dimension : des unités récurrentes linéaires à portes (LRUs) effectuent un mélange d'informations dans le temps, des couches d'auto-attention effectuent un mélange dans l'espace, et des MLPs dans les canaux. L'architecture résultante TRecViT donne de bons résultats sur des tâches clairsemées et denses, formées en régimes supervisés ou auto-supervisés. Notamment, notre modèle est causal et surpasse ou est au niveau d'un modèle d'attention pure ViViT-L sur des ensembles de données vidéo à grande échelle (SSv2, Kinetics400), tout en ayant 3 fois moins de paramètres, une empreinte mémoire 12 fois plus petite et un compte FLOPs 5 fois inférieur. Le code et les points de contrôle seront disponibles en ligne sur https://github.com/google-deepmind/trecvit.

English

We propose a novel block for video modelling. It relies on a time-space-channel factorisation with dedicated blocks for each dimension: gated linear recurrent units (LRUs) perform information mixing over time, self-attention layers perform mixing over space, and MLPs over channels. The resulting architecture TRecViT performs well on sparse and dense tasks, trained in supervised or self-supervised regimes. Notably, our model is causal and outperforms or is on par with a pure attention model ViViT-L on large scale video datasets (SSv2, Kinetics400), while having 3times less parameters, 12times smaller memory footprint, and 5times lower FLOPs count. Code and checkpoints will be made available online at https://github.com/google-deepmind/trecvit.