TRecViT: Un Transformador de Video Recurrente

TRecViT: A Recurrent Video Transformer

December 18, 2024

Autores: 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

Resumen

Proponemos un nuevo bloque para el modelado de video. Se basa en una factorización de tiempo-espacio-canal con bloques dedicados para cada dimensión: las unidades recurrentes lineales con compuertas (LRUs) realizan mezcla de información en el tiempo, las capas de autoatención realizan mezcla en el espacio y las MLPs en los canales. La arquitectura resultante, TRecViT, funciona bien en tareas dispersas y densas, entrenadas en regímenes supervisados o auto-supervisados. Notablemente, nuestro modelo es causal y supera o se iguala a un modelo de atención pura, ViViT-L, en conjuntos de datos de video a gran escala (SSv2, Kinetics400), mientras tiene 3 veces menos parámetros, una huella de memoria 12 veces más pequeña y un conteo de FLOPs 5 veces menor. El código y los puntos de control estarán disponibles en línea en 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.