TS2-Net: Token Shift and Selection Transformer for Text-Video Retrieval

In plain words

Text-video search means finding the videos that match a sentence. A video model has to notice both small objects, such as a hat, and small movements, such as someone talking. Many current models keep the image backbone's fixed structure and only combine per-frame summaries, so these details get lost.

TS2-Net changes how the image transformer handles frames. In its last two layers, each frame borrows whole patch tokens from its neighbouring frames, so the model can see how things move over time without losing the content of any patch. Each frame also keeps only a few of its most informative patches, which cuts the repeated background that would otherwise dominate the result. The sentence is compared with each frame, and the frame scores are combined into one match score. The model is trained so that matching pairs score higher than mismatched ones, and the authors report new best results on five common benchmarks.

TL;DR

• Video encoder: CLIP ViT-B/32. Token shift in blocks 11 and 12 moves whole tokens between adjacent frames; a learned top-k keeps the [CLS] token plus K patch tokens per frame. §3.1 §3.2
• Matching: per-frame cosine similarities, combined with softmax weights (λ = 4). Trained with a symmetric contrastive loss. §3.3
• Author-reported: R@1 47.0 on MSR-VTT 1k-A (ViT-B/32), 49.4 with ViT-B/16. Five benchmarks, with claimed records on MSR-VTT, VATEX, LSMDC, ActivityNet and DiDeMo. tab. 4 tab. 5
• The released code differs from the paper in several places (selection score, K on LSMDC, shift ratio, frame normalisation, frame aggregator, and the initialisation of that aggregator). See Paper vs code.

Problem

Text-video retrieval ranks candidate videos for a text query. The paper argues that video encoders built on pre-trained vision backbones keep a fixed structure, so they cannot be adapted to produce finer spatial-temporal features §1. Encoders that use only the [CLS] token miss small objects, and encoders that use every token of every frame suffer from redundancy, where background tokens dominate the video representation §3.2. The paper targets subtle movements and small objects, which the examples in fig. 1 illustrate §1.

Method

The model has three parts: a text encoder, a video encoder and a matching step §3. The text encoder is a 12-layer transformer that the paper calls GPT and that is initialised from CLIP §3 §4.1. The query vector q is the output at an [EOS] token appended to the query §3.

Video encoder, token shift. The backbone is a 12-layer ViT-B/32 §4.1. The shift is applied in layers 11 and 12 §4.2 tab. 1. It is parameter-free: in the chosen layers, some tokens of each frame are replaced by the same-position tokens of the previous frame, and others by those of the next frame, so the shift runs in both directions §3.1. Shifting whole tokens rather than channel slices is the paper's main departure from earlier shift variants §3.1 fig. 3 tab. 2. The paper places a residual connection between the original and the shifted features §3.1. The shift ratio that works best is 25% of tokens tab. 1.

Video encoder, token selection. Per frame, an MLP reduces the channels from C to C/2. The reduced [CLS] feature is concatenated to every token, and a second MLP followed by softmax gives an importance score to each of the N+1 tokens (eq. 1) §3.2. The K most informative tokens are taken with an indicator matrix M, so the selected tokens are M transposed times the token matrix (eq. 2) §3.2. Top-k is not differentiable, so the paper uses the perturbed maximum with Gaussian noise: ε = 0.05 and m = 500 samples. Eq. 3 is the optimisation form of the top-k. The forward and backward estimators are not in the compiled paper: they sit in a commented-out block §3.2. The selected tokens of all frames go to a joint spatial-temporal transformer. The paper says each frame is represented by its most informative token §3.2.

Matching. Each frame gets a cosine similarity s_i to the query (eq. 4) §3.3. The score of a video is the sum of the frame similarities, weighted by softmax(λ s_i) with λ = 4 (eq. 5) §3.3.

Training. For a batch of B pairs, the loss is the average of a text-to-video and a video-to-text cross-entropy over the score matrix, scaled by a trainable temperature τ (eq. 6 to eq. 8) §3.3. At test time, videos are ranked by the score §3.3.

Pseudocode for the three components. It follows the paper's text; where the released code differs, the line is marked and listed under Paper vs code.

# Token shift on the attention input of ViT blocks 11 and 12  [§3.1] [tab. 1]
# tokens: (T, N+1, C) for one video; index 0 is the [CLS] token
# The shifted tensor feeds the attention call only; the block adds back its
# unshifted input  [code: modules/module_clip.py#L251-L253]
def token_shift(tokens):
    # Stride between shifted positions, derived from the token count
    # (no ratio parameter)  [code: modules/until_module.py#L360]
    stride = (N + 1) // 7
    shifted = copy(tokens)
    # Each frame takes whole tokens from the previous frame; [CLS] (index
    # 0) is never chosen  [code: modules/until_module.py#L382-L383]
    left = range(1, N + 1, stride)
    shifted[1:, left] = tokens[:-1, left]
    # Each frame takes whole tokens from the next frame, starting at
    # position 4  [code: modules/until_module.py#L384-L385]
    right = range(4, N + 1, stride)
    shifted[:-1, right] = tokens[1:, right]
    # Edge frames keep their own tokens; nothing is zero-padded
    # [code: modules/until_module.py#L362-L363]
    return shifted

# Per-frame token selection with [CLS] always kept  [§3.2] [eq. 1] [eq. 2]
# tokens: (T, N+1, C); index 0 is [CLS]; k = 3 patches in the released code
def select_tokens(tokens, k):
    # Channel reduction C -> C/2 (the code adds a LayerNorm first)
    # [code: modules/differential_topk.py#L66-L70]
    reduced = reduce_channels(tokens)
    # Append the reduced [CLS] feature to every token  [eq. 1]
    # [code: modules/differential_topk.py#L93]
    joined = concat(reduced, reduced[:, :1])
    # One score per token; the paper uses softmax, the code uses tanh
    # [code: modules/differential_topk.py#L78]
    scores = tanh(score_mlp(joined))
    # Soft selection of k patch tokens; [CLS] gets no score
    # [code: modules/differential_topk.py#L116-L117]
    indicator = perturbed_top_k(scores[:, 1:], k)  # (T, k, N)
    # Output [CLS] followed by the selected patch tokens  [eq. 2]
    # [code: modules/differential_topk.py#L120-L127]
    picked = indicator @ tokens[:, 1:]
    return concat(tokens[:, :1], picked)  # (T, k+1, C)

# Perturbed top-k: soft one-hot selection under Gaussian noise  [§3.2]
# Backward (Jacobian estimate): mean over draws of (one-hot times noise),
# divided by sigma  [§3.2] [code: modules/differential_topk.py#L44-L58]
# scores: (T, N) patch scores; defaults are the paper's values
def perturbed_top_k(scores, k, draws=500, sigma=0.05):
    # m = 500 noise vectors per frame  [code: modules/differential_topk.py#L24]
    noise = normal(0, 1, shape=(T, draws, N))
    # Perturb the scores and keep the k largest in each draw
    # [code: modules/differential_topk.py#L25-L28]
    noisy = scores[:, None, :] + sigma * noise
    indices = sort(topk(noisy, k))
    # Average the one-hot picks over the draws
    # [code: modules/differential_topk.py#L30-L31]
    return mean(one_hot(indices, N), over_draws)  # (T, k, N)

# Frame similarities and the text-video score  [§3.3] [eq. 4] [eq. 5]
# query: (C,) text vector; frames: (T, C), T <= 12 on the MSR-VTT setting
def video_score(query, frames, lam=4.0):
    # Cosine similarity of the query with every frame  [eq. 4]
    sims = cosine(frames, query)  # (T,)
    # Released code divides by the L2 norm over frames before the softmax
    # [code: modules/modeling.py#L561]  (the paper's weights use sims directly)
    sims_scaled = sims / norm(sims)
    # Padded frames get weight zero  [code: modules/modeling.py#L562]
    # Frame weights: softmax of lam times the scaled similarities
    # [code: modules/modeling.py#L563]
    weights = softmax(lam * sims_scaled)
    # Score: weighted sum of frame similarities  [eq. 5]
    return sum(weights * sims)

# Symmetric contrastive loss over a batch of B pairs  [§3.3] [eq. 6-8]
# sims: (B, B) video scores; diagonal entries are the matching pairs
def retrieval_loss(sims, scale):
    # scale = exp(learned logit_scale), the paper's trainable tau
    # [code: modules/modeling.py#L569]
    logits = scale * sims
    # Text-to-video: row-wise softmax, cross-entropy on the diagonal  [eq. 6]
    # [code: modules/until_module.py#L187-L191]
    loss_t2v = mean(-log_softmax(logits, dim=1).diagonal())
    # Video-to-text: the same on the transposed matrix  [eq. 7]
    loss_v2t = mean(-log_softmax(logits.T, dim=1).diagonal())
    # Average of the two directions  [eq. 8]
    return (loss_t2v + loss_v2t) / 2

Hyperparameters

Results

Author-reported. Metrics: R@1, R@5, R@10, median rank (MdR) and mean rank (MnR); rsum sums the six R@K values. Unless noted, the paper's own model uses CLIP ViT-B/32 and no inverted softmax tab. 4 tab. 5.

MSR-VTT 1k-A test set, tab. 4. Columns: text-to-video R@1, R@5, R@10, MdR, MnR; video-to-text R@1, R@5, R@10, MdR, MnR.

Note: the paper says CLIP2TV uses 16×16 patches, so ViT-B/16 TS2-Net is the fair comparison tab. 4.

Text-to-video on VATEX, ActivityNet, LSMDC and DiDeMo, tab. 5. The paper's caption says QB-Norm uses inverted softmax at inference, while the other rows do not.

Summary rsum values tab. 7: MSR-VTT 205.3 (CLIP2Video 199.9); VATEX 244.3 (CLIP2Video 242.8); LSMDC 116.6 (CAMoE 116.0); DiDeMo 195.4 (CLIP4Clip 193.3); ActivityNet 199.1 (CLIP4Clip rsum not given). Author-reported.

Inverted softmax (appendix A), tab. 6. Text-to-video only, reported separately from the main tables.

Ablations on MSR-VTT 1k-A (ViT-B/32), tab. 1 to tab. 3. Columns: T→V R@1, R@5, R@10, MnR; V→T R@1, R@5, R@10, MnR; rsum. The baseline removes both modules §4.2.

Token shift layers and ratio, tab. 1:

Shift variants, tab. 2 (Channel, [VIS] and [CLS] Channel Shift come from Shift-ViT):

Token selection, tab. 3 (the first row is labelled top-K 1 in the paper):

Other reported results: on a 103-query subset of MSR-VTT with small objects or slight movements, R@1 is 79.6 against 39.8 for CLIP4Clip §4.4. Among the failure cases fig. 9, the paper attributes some errors to vague or empty annotations, which it estimates at 1 to 2% of the dataset app. C.

Paper vs code

Code: https://github.com/yuqi657/ts2_net, commit 1d9d4d53d27057f29fbdba9c4c1586dbda294519. The repository has no LICENSE file, so its code cannot be copied; only pseudocode is used here. The repository's README reports the same ViT-B/32 and ViT-B/16 numbers as the paper; we did not re-run training or evaluation [README.md#L29-L47]. Differences found:

1. Selection score. The paper scores all N+1 tokens with softmax eq. 1 §3.2. The code scores patch tokens with tanh, and [CLS] is not scored differential_topk.py differential_topk.py.
2. [CLS] is always kept. The paper selects the top-K of N+1 tokens, [CLS] included §3.2. The code always keeps [CLS] and adds the top 3 patch tokens modeling.py differential_topk.py. The README says this on purpose [README.md#L17].
3. K on LSMDC. The paper uses K=1 for LSMDC and K=4 elsewhere §4.1. The code fixes topk=3 for every dataset with no flag modeling.py.
4. Frame representation. The paper uses the most informative token of each frame §3.2. The code uses the [CLS] token modeling.py.
5. Selection at inference. The paper does not say. The code always runs the 500-draw noisy top-k, since the training argument of the selector is never read differential_topk.py. Scores therefore vary with the random draws unless the seed is fixed (seed 42 by default) main_task_retrieval.py.
6. Frame weights. The paper uses softmax(λ s_i) eq. 5 §3.3. The code divides the similarity vector by its L2 norm over frames, masks padded frames, then applies softmax(4 ·) modeling.py.
7. Global branch. The paper has one loss §3.3. The code computes two similarity branches and weights the global one by zero (frame_match_weight = 1.0) modeling.py modeling.py. By our reading, for the single [EOS] text vector the token-level softmax in the frame branch has one element and returns 1, so the two branches compute the same score modeling.py.
8. Frame aggregator. The paper names no module for the 4 layers on the selected tokens beyond a joint spatial-temporal transformer §3.2 §4.1. The scripts set --sim_header seqTransf with 4 layers, which adds frame position embeddings and a transformer over all selected tokens [scripts/run_msrvtt.sh] modeling.py. The code's default is meanP, which skips this step main_task_retrieval.py modeling.py.
9. Initialisation of that aggregator. The paper says CLIP weights of similar dimension are used §4.1. The code copies the first 4 blocks of CLIP's text transformer into it, and the frame position embeddings from CLIP's text positional embedding modeling.py modeling.py.
10. Shift ratio. The paper finds 25% best tab. 1. The code has no ratio: the stride is the token count divided by 7, so about 7 positions per direction are shifted until_module.py until_module.py. For the 7×7 grid of ViT-B/32 (not stated in the paper), that is 14 of 49 patches, about 29%, across both directions. For ViT-B/16 it is 14 of 196, about 7%. The 12.5% and 50% sweeps in tab. 1 cannot be reproduced with the released code.
11. Edge frames. fig. 3 shows padding for shifted-out positions. The code starts the shifted tensor as a copy of the input, so edge frames keep their own tokens until_module.py.
12. Shifted positions. The code never shifts [CLS] (indices start at 1) until_module.py. The paper does not say.
13. Shift location. Blocks 11 and 12 (1-indexed) match §4.2 and tab. 1 until_module.py.
14. Learning rates. Consistent with §4.1: parameters whose names contain "clip." get 1e-7 and the rest get 1e-4 main_task_retrieval.py. The frame aggregator is named transformerClip, so it falls in the 1e-4 group.
15. Loss. Consistent with eq. 6 to eq. 8 until_module.py. Only the frame branch is weighted, as in item 7.
16. Perturbed top-k constants. Consistent: ε = 0.05 and 500 samples §3.2 differential_topk.py. The code's forward and backward estimators follow the cited perturbed-optimizer method. The paper's own forward and backward formulas are in a commented-out block and absent from the compiled paper, so the code is the only source for them §3.2 differential_topk.py.
17. Optimizer. Not stated in the paper. The code uses BertAdam with 10% warmup, cosine decay, weight decay 0.2, betas (0.9, 0.98), eps 1e-6 and gradient clipping at 1.0 main_task_retrieval.py optimization.py.
18. Unused settings. --lr_decay (default 0.9) is defined but not used in the training code main_task_retrieval.py. The text selector is defined but its call is commented out modeling.py modeling.py.
19. Missing scripts. The repository gives runs for MSR-VTT and VATEX only [README.md#L54-L62]. There are no ActivityNet, DiDeMo or LSMDC scripts, so the 64-frame and 64-word settings of §4.1 cannot be checked against code.

Pitfalls

• The 25% shift ratio is not a knob. Changing the ratio requires changing the stride rule itself until_module.py. Other ViT patch grids change the effective fraction, so results will not transfer directly.
• Edge handling. The paper's illustration uses padding for shifted-out positions fig. 3. The code keeps each edge frame's own tokens. Use the code's behaviour if you want to match the reported numbers.
• Stochastic selection at inference. The released code samples 500 noisy draws at test time too, so ranking scores can change between runs with different seeds or batch orders differential_topk.py. Fix the seed, or verify run-to-run variance before comparing with table numbers.
• Memory of the perturbed top-k. The one-hot tensor has one entry per draw, per frame, per selected token, and per patch. For 128 videos, 12 frames, 3 selected patches and 49 patches (ViT-B/32 grid, not stated in the paper), that is about 113 million float entries, or about 450 MB in fp32, kept for backward on one device differential_topk.py differential_topk.py. Larger frame counts, as in the 64-frame ActivityNet and DiDeMo settings, multiply this cost.
• λ on normalised scores. The code divides the frame similarities by their L2 norm before the temperature-4 softmax modeling.py. The effective sharpness therefore depends on the spread of the scores and on the number of valid frames. The paper's formula has no such step.
• Frame sampling is unstated. The scripts take 1 fps and sample uniformly across the video to reach 12 frames [scripts/run_msrvtt.sh] dataloader_msrvtt_retrieval.py. The paper gives only the maximum frame count §4.1.
• Initialisation source. The frame aggregator starts from CLIP's text-encoder blocks, not its visual blocks modeling.py. The paper's "similar dimension" wording does not tell you this.
• Warmup and cosine. The schedule is cos(π x) over the whole run with x the fraction of training done, not a cosine after warmup. At the end of warmup the multiplier drops from 1.0 to about 0.976 optimization.py.
• Learning-rate split. The 1e-7 rate applies only to parameters whose names contain "clip." main_task_retrieval.py. Renaming modules changes the schedule without any warning.
• Scores are author-reported. The headline records are on the paper's own benchmarks and splits. Inverted softmax is reported separately and changes the numbers tab. 6.
• Repository licence. No LICENSE file means the code cannot be copied or redistributed; use it as a reference only.

Lineage

• Builds on: CLIP ViT-B/32 and ViT-B/16 weights for initialisation §4.1 tab. 4; ViT §3.1; Shift-Transformer shift variants, which tab. 2 compares against §3.1; the temporal shift module TSM §1; CLIP4Clip, used as a mean-pooling baseline §3.2; the perturbed maximum of Berthet et al. and the backward estimators of Abernethy et al. and Cordonnier et al. §3.2.
• Compared against: CE, TACo, MMT, Frozen, MDMMT, CLIP, CLIP4Clip, CAMoE, CLIP2Video, CLIP2TV, QB-Norm and others tab. 4 tab. 5 tab. 6.
• Code basis: the repository's README says it is based on CLIP4Clip [README.md#L85].