| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # VidCom2 (Video Compression Commander) video token pruning. |
| 5 | # Liu et al., EMNLP 2025 — https://arxiv.org/abs/2505.14454 |
| 6 | # Adapted from the reference implementation: |
| 7 | # https://github.com/xuyang-liu16/VidCom2 (Apache-2.0, |
| 8 | # Copyright (c) 2025 the VidCom2 authors). |
| 9 | |
| 10 | import torch |
| 11 | import torch.nn.functional as F |
| 12 | |
| 13 | # Multi-scale Gaussian bandwidths from the reference implementation. |
| 14 | _ALPHAS: tuple[float, ...] = tuple(2.0**k for k in range(-3, 2)) |
| 15 | _LOW_VAR_CHANNEL_RATIO: float = 0.5 |
| 16 | _SOFTMAX_TEMPERATURE: float = 0.01 |
| 17 | |
| 18 | |
| 19 | def compute_retained_tokens_count( |
| 20 | tokens_per_frame: int, num_frames: int, q: float |
| 21 | ) -> int: |
| 22 | """Number of video tokens retained after VidCom2 pruning. |
| 23 | |
| 24 | The target is `(1 - q) * total_tokens`, i.e. a retention ratio of |
| 25 | `1 - q` averaged across frames. Because the per-frame budget is floored |
| 26 | at one token, the global minimum is `num_frames` (one token per frame). |
| 27 | """ |
| 28 | total_tokens = tokens_per_frame * num_frames |
| 29 | base_num = int(total_tokens * (1.0 - q)) |
| 30 | return max(num_frames, min(base_num, total_tokens)) |
| 31 | |
| 32 | |
| 33 | def compute_retention_mask( |
| 34 | video_embeds: torch.Tensor, |
| 35 | video_size_thw: torch.LongTensor | tuple[int, int, int], |
| 36 | spatial_merge_size: int, |
| 37 | q: float, |
| 38 | ) -> torch.Tensor: |
| 39 | """Compute the VidCom2 retention mask for a single video. |
| 40 | |
| 41 | Args: |
| 42 | video_embeds: `(T*H*W/merge^2, hidden_size)` post-ViT token features. |
| 43 | video_size_thw: `(T, H, W)` grid dimensions. |
| 44 | spatial_merge_size: ViT spatial merge factor (e.g. 2). |
| 45 | q: Pruning fraction in `[0, 1)`; retention ratio is `1 - q`. |
| 46 | |
| 47 | Returns: |
| 48 | Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained |
| 49 | tokens. The True count equals `compute_retained_tokens_count` so |
| 50 | placeholders sized at prompt-processing time match exactly. |
| 51 | |
| 52 | """ |
| 53 | T, H, W = map(int, video_size_thw) |
| 54 | rows = H // spatial_merge_size |
| 55 | cols = W // spatial_merge_size |
| 56 | tokens_per_frame = rows * cols |
| 57 | total_tokens = T * tokens_per_frame |
| 58 | |
| 59 | device = video_embeds.device |
| 60 | if tokens_per_frame == 0 or total_tokens == 0: |
| 61 | return torch.ones(0, dtype=torch.bool, device=device) |
| 62 | |
| 63 | target_retained = compute_retained_tokens_count( |
| 64 | tokens_per_frame=tokens_per_frame, num_frames=T, q=q |
| 65 | ) |
| 66 | target_retained = min(target_retained, total_tokens) |
| 67 | |
| 68 | # 1. Score in the lowest-variance half of channels. |
| 69 | variances = video_embeds.var(dim=0, unbiased=False) |
| 70 | k_channels = max(1, int(video_embeds.size(-1) * _LOW_VAR_CHANNEL_RATIO)) |
| 71 | _, low_var_idx = torch.topk(variances, k=k_channels, largest=False) |
| 72 | sel = video_embeds.index_select(-1, low_var_idx) |
| 73 | |
| 74 | # 2. Multi-scale Gaussian similarity to video and per-frame centers. |
| 75 | frames = sel.view(T, tokens_per_frame, sel.size(-1)) |
| 76 | frames = F.normalize(frames, dim=-1) |
| 77 | vid_center = frames.mean(dim=(0, 1), keepdim=True) # (1, 1, C) |
| 78 | frame_center = frames.mean(dim=1, keepdim=True) # (T, 1, C) |
| 79 | v_score = _multi_scale_gaussian(frames, vid_center) |
| 80 | f_score = _multi_scale_gaussian(frames, frame_center) |
| 81 | # Higher similarity = more redundant; lowest-similarity tokens are kept. |
| 82 | similarity = v_score + f_score # (T, tpf) |
| 83 | |
| 84 | # 3. Per-frame dynamic budget: distinctive frames get a larger share. |
| 85 | base = 1.0 - q |
| 86 | frame_scores = -v_score.mean(dim=-1) # (T,) |
| 87 | probs = F.softmax((frame_scores - frame_scores.max()) / _SOFTMAX_TEMPERATURE, dim=0) |
| 88 | scales = (base * (1.0 + probs - probs.mean())).clamp(max=1.0) |
| 89 | ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame) |
| 90 | |
| 91 | # 4. Retain the smallest-similarity tokens per frame. |
| 92 | budgets = ks.tolist() |
| 93 | mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device) |
| 94 | for i, k_i in enumerate(budgets): |
| 95 | _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False) |
| 96 | mask_2d[i].scatter_(0, idx, True) |
| 97 | |
| 98 | # 5. Reconcile rounding/clamp drift to the exact target count by score. |
| 99 | flat_mask = mask_2d.view(-1) |
| 100 | flat_sim = similarity.view(-1) |
| 101 | current = sum(budgets) |
| 102 | if current > target_retained: |
| 103 | drop_n = current - target_retained |
| 104 | retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1) |
| 105 | retained_sim = flat_sim[retained_idx] |
| 106 | _, worst = torch.topk(retained_sim, k=drop_n, largest=True, sorted=False) |
| 107 | flat_mask[retained_idx[worst]] = False |
| 108 | elif current < target_retained: |
| 109 | add_n = target_retained - current |
| 110 | available_idx = (~flat_mask).nonzero(as_tuple=False).squeeze(-1) |
| 111 | if available_idx.numel() > 0: |
| 112 | available_sim = flat_sim[available_idx] |
| 113 | add_n = min(add_n, available_idx.numel()) |
| 114 | _, best = torch.topk(available_sim, k=add_n, largest=False, sorted=False) |
| 115 | flat_mask[available_idx[best]] = True |
| 116 | |
| 117 | return flat_mask |
| 118 | |
| 119 | |
| 120 | def _multi_scale_gaussian(x: torch.Tensor, center: torch.Tensor) -> torch.Tensor: |
| 121 | """Sum Gaussian kernels over `_ALPHAS`; `(T, N, C) -> (T, N)` scores.""" |
| 122 | dist_sq = ((x - center) ** 2).sum(dim=-1) |
| 123 | return sum(torch.exp(-dist_sq / (2.0 * a)) for a in _ALPHAS) |