TrajTok: Learning Trajectory Tokens enables better Video Understanding

summary

The gist

The gist: TrajTok, an end-to-end differentiable trajectory tokenizer, learns to implicitly propose trajectory tokens by integrating a unified segmenter and a trajectory encoder that co-train with

In short

TrajTok is an end-to-end video tokenizer that learns to create compact trajectory tokens by integrating a unified segmenter and a trajectory encoder. It dynamically adapts token granularity based on semantic complexity, decoupling token count from video duration. This method results in efficient models like TrajViT2, achieving state-of-the-art accuracy at scale and showing strong versatility for various downstream tasks.

Key concepts

Universal Segmenter: This component partitions the video into semantic groups by extracting frame features and using learnable queries processed through Perceiver layers. It generates soft segmentation masks that group pixels in both space and time, prioritizing robust semantic grouping over pixel-perfect accuracy.
Trajectory Encoder: The encoder aggregates the feature maps from the segmenter into compact latent tokens. It refines these representations using hard attention masks derived from the soft masks to recover fine-grained motion details specific to each trajectory region, ensuring disentangled and meaningful token generation.
Adaptive Token Mechanism: Inspired by Matryoshka representations, this mechanism allows the encoder to emit a variable number of tokens (n) per trajectory. This balances efficiency and expressivity by setting the token count based on a predefined compute budget while using distinct learnable queries for diversity.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Retrieval-Augmented Generation (RAG) — explained in 20 episodes
Sample Complexity — explained in 20 episodes
Quantum Fisher Information (QFI) — explained in 19 episodes
Flow Matching — explained in 18 episodes
Federated Learning — explained in 17 episodes

This episode discusses

The paper

TrajTok: Learning Trajectory Tokens enables better Video Understanding · Read on arXiv

Chenhao Zheng, Jieyu Zhang, Jianing Zhang, Weikai Huang, Ashutosh Kumar, Quan Kong

University of Washington · Allen Institute for Artificial Intelligence · Apple

Tokenization in video models, typically through patchification, generates an excessive and redundant number of tokens. This severely limits video efficiency and scalability. While recent trajectory-based tokenizers offer a promising solution by decoupling video duration from token count, they rely on complex external segmentation and tracking pipelines that are slow and task-agnostic. We propose TrajTok, an end-to-end video tokenizer module that is fully integrated and co-trained with video models for a downstream objective, dynamically adapting its token granularity to semantic complexity, independent of video duration. TrajTok contains a unified segmenter that performs implicit clustering over pixels in both space and time to directly produce object trajectories in a single forward pass. By prioritizing downstream adaptability over pixel-perfect segmentation fidelity, TrajTok is lightweight and efficient, yet empirically improves video understanding performance. With TrajTok, we implement a video CLIP model trained from scratch (TrajViT2). It achieves the best accuracy at scale across both classification and retrieval benchmarks, while maintaining efficiency comparable to the best token-merging methods. TrajTok also proves to be a versatile component beyond its role as a tokenizer. We show that it can be seamlessly integrated as either a probing head for pretrained visual features (TrajAdapter) or an alignment connector in vision-language models (TrajVLM) with especially strong performance in long-video reasoning.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "TrajTok: Learning Trajectory Tokens enables better Video Understanding".

Jane: Here's the gist of the paper. The gist: TrajTok, an end-to-end differentiable trajectory tokenizer, learns to implicitly propose trajectory tokens by integrating a unified segmenter and a trajectory encoder that co-train with downstream objectives, offering low token counts, efficiency and adaptability to downstream objectives.

Tom: How it works 1.

Jane: The system consists of two differentiable components: a universal segmenter that partitions the input into semantic groups, and a trajectory encoder that aggregates these groups into compact latent tokens (Page 3).

Tom: 2.

Paper summary: Jane: They are arguing that by organizing visual tokens around object trajectories, it aligns better with how humans actually perceive things in video. But the paper points out a big limitation with previous trajectory-based tokenizers.

Lu: That external pipeline approach, where you use general-purpose segmentation models to generate trajectories, is what they criticize because it makes the tokenizer slow and independent of the actual task you want to perform later on.

Tom: So they want something that’s not a slow preprocessing step but something that learns these groupings directly through training with the model itself.

Meng: And they do this by using a unified segmenter and a trajectory encoder, where the segmenter proposes masks for all objects in a video in one forward pass, and the encoder then aggregates features based on those masks to make those tokens.

Lalam: The system uses a Perceiver layer stack for the segmenter to generate soft segmentation maps, which are then used by the trajectory encoder to produce initial proposal embeddings z init k through weighted aggregation of features using those soft masks.

Jane: And they refine these initial embeddings by using a second Perceiver module and enforcing masked cross-attention with hard segmentation maps, which sharpens those representations.

Tom: A key part is that the adaptive token mechanism is inspired by Matryoshka representations, allowing the encoder to emit n tokens per trajectory based on a compute budget, balancing efficiency and expressivity.

Lu: They specifically mention that they prioritize downstream adaptability over pixel-perfect segmentation fidelity, which makes it lightweight and efficient while still improving video understanding performance.

Meng: What I want to know from that is how they handle the training side because if it's truly end-to-end, the training objective has to be very clever.

Lalam: They achieve this by using a combination of Dice loss and Focal loss for segmentation, which prioritizes finding all object regions over strict pixel-level class accuracy.

Tom: So they train the segmenter jointly with the transformer under an added segmentation supervision during training on TrajViT2. That’s where the real learning happens, right?

Jane: Right, and that training setup allows for pre-training the segmenter and then reusing it for other things later, like probing or vision-language model training by annotating videos with pseudo ground truth masks generated through that TrajViT pipeline.

Lu: It suggests a pathway where you can leverage this learned segmentation structure without needing to re-segment everything from scratch every time.

Conclusion: Tom: So, looking at the paper "TrajTok: Learning Trajectory Tokens enables better Video Understanding", the authors are essentially saying that traditional patchification is inefficient because it creates too much noise.

Jane: They propose TrajTok as a way around that by learning to group visual information into trajectories directly from the video input. This is achieved through an end-to-end differentiable framework with a segmenter and an encoder.

Lu: The implication for research is that we can move away from slow, external pipelines for segmentation and tracking in video understanding if we can design a system that learns the trajectory tokens organically alongside the main model training.

Meng: From an engineering standpoint, it’s exciting because they show performance comparable to existing token-merging methods while maintaining efficiency similar to them. It seems like a practical way to get better representation with less computational overhead at scale.

Lalam: And its versatility is notable; they show it works as a video transformer from scratch for classification, and it can be used as an adapter for probing features or as an alignment connector in vision-language models.

Tom: So, the main point is that trajectory-based tokenization offers a more efficient and semantically aligned alternative to the old patchification methods when you’re dealing with video understanding tasks.

Jane: It's a lot of work to make it end-to-end, but they demonstrate that learning these trajectory tokens can lead to much better performance on long-video reasoning benchmarks, which is where this approach seems strongest.

More episodes

← Home