TrajTok: Learning Trajectory Tokens enables better Video Understanding

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.

How it works

1. 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).

2. The universal segmenter performs frame-wise feature extraction using a lightweight patch encoder, extracts multi-scale maps, and uses learnable queries processed through a stack of Perceiver layers to generate soft segmentation masks (Page 3).

3. Soft segmentation is achieved by computing the similarity between processed queries and patch features via softmax over the query dimension to obtain a soft segmentation map Msoft (Page 4).

4. The trajectory encoder aggregates patch-level feature maps into compact tokens, initially generating proposal embeddings z init k through a weighted aggregation of features using the soft masks (Page 4).

5. Trajectory embedding refinement sharpens these representations by employing a second Perceiver module and enforcing masked cross-attention using hard segmentation maps Mhard (Page 4).

6. An adaptive token mechanism inspired by Matryoshka representations allows the encoder to emit n ∈ three tokens per trajectory, balancing efficiency and expressivity based on a predefined compute budget (Page 4).

Tokenization and Segmentation Details

The segmenter is designed to decouple video duration from the final token count, prioritizing robust semantic grouping over pixel-perfect segmentation masks (Page 3). The training of the segmenter can be done either independently or jointly with other objectives (Page 4). The combination of Dice loss and Focal loss yields the best downstream understanding, as it prioritizes the discovery of all object regions over strict pixellevel class accuracy (Page 4).

When training TrajViT2, the segmenter is trained jointly with the transformer under an added segmentation supervision (Page 4). The segmenter can be pre-trained and reused for downstream probing or VLM training by annotating videos with pseudo ground truth masks generated via the TrajViT pipeline (Page 10).

Trajectory Encoder Refinements

The trajectory encoder accepts both soft and hard segmentations to ensure differentiability while maintaining disentangled representations (Page 4). Initial trajectory embeddings z init k are computed as z init k = X t,i,j Msoft k,t,i,j · Ft,i,j (Page 4). To recover fine-grained motion and texture details specific to the trajectory’s region, the encoder uses hard attention masks Mhard obtained by applying an argmax to Msoft and converting it to one-hot binary assignments (Page 4).

The adaptive token mechanism allows for n ∈ three tokens per trajectory, where each copy is associated with a distinct learnable query vector (Page 4). To encourage diversity among sub-tokens, these queries are initialized with Fourier positional embeddings with angular offsets that maximally separate them in feature space (Page 4).

Versatility and Evaluation Scenarios

TrajTok is versatile and can be implemented in three scenarios:

1. TrajViT2, which trains a video transformer encoder from scratch under the CLIP objective, achieving improvements over baselines (Page 6).

2. TrajAdapter, a plug-in feature adapter that aggregates dense feature maps from any pretrained video encoder to enhance probing performance without full fine-tuning (Page 7).

3. TrajVLM, a vision-language model where TrajTok serves as an alignment module between a ViT and an LLM, improving video question-answering performance especially for long-video questioning benchmarks (Page 7).

Performance Results

TrajViT2 achieves the best accuracy at scale across both classification and retrieval benchmarks, showing a large-margin improvement of +4.8% on Kinetics-400 and +4.1% on SSv2 over a standard video ViT (Page 6). TrajViT2 exhibits stronger scaling behavior than TrajViT as the training dataset size increases, possibly because of its segmenter’s flexibility in adapting to downstream tasks (Page 6).

TrajAdapter consistently improves the probing accuracy of pretrained backbone features, and performance further increases with more tokens per trajectory (Page 7). TrajVLM notably outperforms patch pooling baselines on long-video benchmarks, including a notable +8.8% on LongVideoBench and +5.4% on LVBench over PatchVLM with default poolsize=3 (Page 8).

Overall, these results validate TrajTok as an effective connector for VLMs, particularly in long-video understanding (Page 8). The segmenter's lightweight design and low output resolution introduce several expected limitations, such as occasionally missing very small objects (Page 8).

The paper concludes by introducing TrajTok as an end-to-end and efficient tokenizer that learns to group visual trajectories and produce trajectory-level tokens directly from video inputs (Page 9). The results highlight the potential of trajectory-based tokenization as a more efficient and semantically aligned alternative to traditional patchification (Page 9). The project was funded by DSO National Laboratories in Singapore and by Toyota Motor Inc. (Page 10).

1. TrajTok as an end-to-end differentiable video tokenizer allows for low token counts, efficiency and adaptability to downstream objectives, which enables training models like TrajViT2 that achieve the best accuracy at scale across both classification and retrieval benchmarks. This means AI systems can process long videos with significantly fewer tokens than traditional patchification, leading to reduced memory bottlenecks and better scalability.

2. The unified segmenter in TrajTok performs implicit clustering over pixels in both space and time to directly produce object trajectories in a single forward pass, which is much more efficient than prior work [94] because it trades pixel-perfect accuracy with end-task performance. This allows for dynamic token granularity that adapts its tokenization to the downstream tasks, meaning the system can generate fine-grained tokens for complex movements (like individual body parts) or coarser tokens for group formations, as demonstrated in Figure 3.

3. TrajAdapter enables a cost-effective way to enhance the probing performance of pretrained encoders in common video classification benchmarks without full fine-tuning. This allows researchers to efficiently adapt large, pre-trained vision encoders to new tasks by reorganizing their dense feature maps into trajectory tokens, proving that naively inserting a Perceiver module does not yield any gain over the attentive probing baseline.

4. TrajVLM can serve as a connector in vision-language models, where it improves video question-answering performance, especially for long-video questioning benchmarks, by grouping ViT features along trajectories and passing the grouped trajectory tokens as visual inputs to the language model, leading to superior performance on long-video reasoning tasks like LongVideoBench.

5. The adaptive token mechanism inspired by Matryoshka representations allows for a flexible trade-off between efficiency and expressivity, enabling the encoder to emit n ∈ 1, 2, 4 tokens per trajectory, which can be adjusted during inference based on available computational resources. This ensures that complex or articulated motion is represented without resorting to overly compressed or rigid token structures.

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.