VideoChat2-TPO: Visual Task Optimization

• VideoChat2-TPO is an enhanced multimodal model that integrates a differentiable task-preference mechanism for task-specific inference.
• It leverages learnable task tokens and dedicated task heads (region, temporal, mask) to address fine-grained challenges in visual tasks.
• A multi-task co-training strategy with optimized training configurations yields significant performance gains across multiple benchmarks.

VideoChat2-TPO refers to the augmentation of the VideoChat2 multimodal LLM (MLLM) with Task Preference Optimization (TPO), a framework designed to enable fine-grained visual task alignment and task-specific inference by leveraging differentiable task preferences, learnable task tokens, and task-specific heads. This instantiation targets the critical challenge of improving task precision and controllability in MLLMs, particularly for video and vision-language grounding, segmentation, and temporal localization. VideoChat2-TPO incorporates architectural, algorithmic, and training modifications that enable scalable, high-accuracy, fine-grained visual reasoning and interaction, with empirical gains substantiated across a comprehensive suite of benchmarks (Yan et al., 2024).

1. Differentiable Task Preference Mechanism

A central innovation in VideoChat2-TPO is its differentiable task-preference mechanism, adapted from Direct Preference Optimization (DPO) for vision-language tasks. For each input $(\mathbf{X}, \mathbf{T}_q)$—where $\mathbf{X}$ denotes a sequence of video frames and $\mathbf{T}_q$ is a text query—the model is required: (a) to predict which fine-grained task should be activated (task assignment), and (b) to produce a task-specific output via a dedicated task head $H_t$.

The model stack is expressed as $M=(E, C, G)$, where $E$ is a vision encoder, $C$ a vision-language connector (Q-former), and $G$ the LLM (Mistral-7B with LoRA adapters). A set of $n$ learnable task tokens $\{\mathbf{v}_i\}_{i=1}^n$ ($\mathbf{X}$0) is associated one-to-one with task-specific heads $\mathbf{X}$1. The LLM $\mathbf{X}$2 first produces a "task-type" token, which is subsequently replaced with $\mathbf{X}$3 (the corresponding learnable token), whose last-layer hidden representation $\mathbf{X}$4 is provided as input to the selected head $\mathbf{X}$5. Prediction proceeds as $\mathbf{X}$6, where $\mathbf{X}$7.

The total loss is a composite of next-token log-likelihood for the conversational objective, cross-entropy for the task-type prediction, and per-task supervised loss: $\mathbf{X}$8 with all $\mathbf{X}$9 weights empirically set to 1 after tuning (Yan et al., 2024).

2. Learnable Task Tokens and Task Heads

Each fine-grained vision task is assigned a distinct learnable token $\mathbf{T}_q$0, appended to the token sequence of $\mathbf{T}_q$1 after emission of the corresponding "task-type" marker. During training, the gradients propagate through $\mathbf{T}_q$2 and $\mathbf{T}_q$3, enabling decoupled, task-specific adaptation via the learnable token and task head elements, while leaving the core LLM weights largely unchanged except for LoRA adapters.

The principal heads are:

• Region head ($\mathbf{T}_q$4): A 2-layer MLP that performs bounding-box regression $\mathbf{T}_q$5 on the LLM-extracted task embedding.
• Temporal head ($\mathbf{T}_q$6): Utilizes a CG-DETR architecture, employing a frozen InternVideo2 pre-extractor for video features and predicts $\mathbf{T}_q$7 with frame-level highlight scores.
• Mask head ($\mathbf{T}_q$8): Deploys SAM2’s image encoder and mask decoder, replacing the prompt encoder with a two-layer "mask adapter" MLP processing $\mathbf{T}_q$9.

This modular construction allows scalable addition of further task heads and learnable tokens as new task types are introduced (Yan et al., 2024).

3. Architectural Modifications from VideoChat2 to VideoChat2-TPO

The base VideoChat2 architecture comprises:

• Vision encoder ($H_t$0): UMT-L, a video-optimized transformer.
• Connector ($H_t$1): BERT$H_t$2-style Q-former, condensing dense visual features to cross-modal queries.
• LLM ($H_t$3): Mistral-7B with LoRA adapters.

TPO integration introduces:

• "Task-type" markers injected into $H_t$4.
• $H_t$5 learnable task tokens.
• The three dedicated heads (region, temporal, mask) with tightly coupled routing through the LLM hidden state.

Inference proceeds via the following sequence: video frames $H_t$6 visual query tokens (concatenated with text tokens to $H_t$7); $H_t$8 emits task-type marker (replaced by $H_t$9), whose hidden state is routed to $M=(E, C, G)$0 (Yan et al., 2024).

4. Multi-Task Co-Training Strategy

VideoChat2-TPO employs a three-stage local-to-global multi-task learning schedule:

1. Task Assignment Tuning: $M=(E, C, G)$1 frozen; LoRA adapters in $M=(E, C, G)$2 trained to accurately emit the task-type token from a corpus of 50K samples per task across DiDeMo, QuerYD, RefCOCO(+), SAMv2, MeViS.
2. Task Head Pretraining: $M=(E, C, G)$3 remain frozen; each $M=(E, C, G)$4 pair is tuned on task-specific datasets—540K examples for region, 116.5K temporal, 114.6K segmentation—with LoRA adapters updated.
3. Global Multi-Task Co-Training: $M=(E, C, G)$5 are unfrozen; all task-heads, tokens, and LoRA adapters trained jointly with a mixture of 3M multimodal conversations (VideoChat2-IT, ShareGPT-4o) and held-out splits of each fine-grained task. All gradients back-propagate through the shared $M=(E, C, G)$6 weights.

Ablation of this scheme demonstrates clear synergies: co-training improves each task beyond what single-task training provides (e.g., Charades-STA [email protected] improves from 30.2 to 40.2 with full synergy) (Yan et al., 2024).

5. Training Configurations and Hyperparameters

Training is conducted using 32 NVIDIA A100 GPUs under DeepSpeed Stage 2 with bf16 precision and the AdamW optimizer (weight decay 0.02, cosine learning rate schedule). Notable hyperparameters are:

• LoRA rank: 16, $M=(E, C, G)$7
• Assignment head LR: $M=(E, C, G)$8
• Region/temporal/mask heads: $M=(E, C, G)$9
• Global multi-task stage: $E$0 for all modules, batch size 256, 3 epochs, using a mix of conversational data (3M) and fine-grained task data ($E$1K temporal grounding, $E$2K spatial grounding, $E$3K mask, $E$4K reasoning) (Yan et al., 2024).

Inference incurs less than 5% extra latency compared to base VideoChat2.

6. Quantitative Performance and Evaluation

VideoChat2-TPO achieves statistically significant improvements over its predecessor and other baseline MLLMs. Key results include:

Benchmark	VideoChat2	VideoChat2-TPO	Absolute Gain
MVBench AVG	60.4%	66.8%	+6.4
Video-MME	39.5/43.8	48.8/53.8	+9.3/+10.0
MLVU M-AVG	44.5	54.7	+10.2
MMIU Overall	35.0	40.2	+5.2
SEED2 (I/M)	26.5/27.6	67.3/70.0	+40.8/+42.4
Next-GQA Acc@IoP	54.9	77.7	+22.8

For fine-grained vision tasks:

Task	VideoChat2	VideoChat2-TPO	Absolute Gain
Moment retrieval (Charades)	14.3	40.2	+25.9
Highlight detection (mIoU)	—	65.0	—
Spatial grounding (RefCOCO)	90.6	90.8	+0.2
Tracking (LaSOT)	54.4	69.4	+15.0
Referring segmentation	61.7	75.4	+13.7

Applying TPO to alternative MLLMs (e.g., LLaVA-OV) also yields robust improvements (MVBench: 56.7 → 64.8, +8.1). Across more than 10 multimodal benchmarks, the average relative improvement is +14.6% over baseline MLLMs (Yan et al., 2024).

7. Implementation and Broader Context

VideoChat2-TPO is implemented in PyTorch with DeepSpeed Stage 2. Training durations are: Stage 1 ($E$5K examples, 2 hours), Stage 2 ($E$6 hours), and Stage 3 (3 epochs on $E$7M examples, $E$8 hours). All training and inference employ bf16 precision, and LoRA adapters are managed via HuggingFace PEFT. The overhead introduced by task selection and specialized head invocation is negligible in practice ($E$9 latency added).

By introducing a differentiable task-preference layer—learnable tokens paired with specialized heads—VideoChat2-TPO enables MLLMs to handle complex, fine-grained visual tasks while preserving conversational fluency and generalization capacity. This architecture demonstrates scalability and extensibility to new vision tasks and models (Yan et al., 2024).

---

Primary references: (Yan et al., 2024, Oquab et al., 2020).