{"id":"4f6e3d3e-d459-4647-a682-e6d2053b77fd","arxiv_id":"2412.10569","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"low","formal_verification":"none","parameter_count":5,"one_line_summary":"A lightweight decoupled embedding, trained through a differentiable approximation of token merging, improves the speed-accuracy trade-off of Vision Transformers.","lead":"This paper introduces DTEM, a small learned module that decides which image patches to merge in a Vision Transformer, making the model faster while keeping accuracy nearly unchanged. The method improves the accuracy-computation trade-off over standard token merging across classification, captioning, and segmentation.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Soft-to-hard transfer at the trained temperature is unverified: Eqs. 6-10 are only asymptotically equivalent to hard BSM, and the stated τ=0.1 plus similarity scaling may put the effective temperature near 1.","rationale":"I considered the reader's weakest assumption and agree it is the most load-bearing, with one correction: since the main tables already evaluate hard BSM at inference, a poor soft-to-hard transfer cannot make the measured gains 'shrink or reverse'—those gains are post-transfer. The real risk is that the training relaxation is optimizing a different objective than the deployed hard merging, so the claimed mechanism ('learned directly through the token merging process') is not established, and the improvements might be an artifact of the extra parameters or of a different implicit objective. This matters because the paper's novelty is precisely the continuous relaxation as a training signal. Other weaknesses (no error bars, small effect sizes, missing learned-merging baseline) are real but secondary: the gains are consistent across three tasks and many backbones, and Table 8 partially addresses a learned baseline. The soft-to-hard gap is unmeasured and could be checked directly without retraining, so it is the appropriate focus. The conditional verdict stands: the paper should either release code and report the soft/hard agreement, or retrain with a hard-aware estimator, before the mechanism claim is fully accepted.","tokens_in":19020,"tokens_out":25547,"duration_ms":239940,"concrete_test":"Use the released DTEM checkpoint for DeiT-S (or retrain with the paper's setup). On a held-out sample of ImageNet validation images, at each transformer block compute (a) the soft adjacency E~ from Eqs. 6-8 at the training temperature and (b) the hard BSM adjacency E' from the same learned embeddings. Measure edge-set agreement (fraction of top-r edges that coincide, or Kendall tau between E~ and the hard scores), and compare the model's top-1 accuracy under the full soft forward pass vs the hard forward pass. If agreement is below about 0.8 or the soft/hard accuracy gap exceeds 1 point, the relaxation is not a faithful surrogate at the trained temperature; the paper should retrain with a straight-through estimator of hard BSM and re-check whether the gains over ToMe persist. If agreement is high and the gap is small, the concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that decoupled embeddings learned via continuously relaxed merging improve hard token merging. The load-bearing step is the transfer from the soft relaxation in Secs. 3.2-3.3 (Eqs. 6-10) to the hard BSM used at inference (Sec. 3.4). The paper justifies this only by an asymptotic argument as τ→0, but training uses τ=0.1 and additionally scales similarities by 0.1, so the effective softmax temperature may be near 1, far from the asymptotic regime. If the continuous adjacency matrix E~ disagrees with the hard BSM adjacency E' at this temperature, the embeddings are trained for a surrogate merging policy different from the deployed one. Because the reported gains already come from hard inference, the concern is not that the gains will shrink at test time; rather, the training mechanism might not be the source of the gains, and a method that directly optimizes hard BSM (e.g., straight-through) might do as well or better without the claimed relaxation. No experiment in the paper measures the soft-to-hard gap.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Decoupled Token Embedding for Merging (DTEM), a method for token reduction in Vision Transformers. Instead of computing token similarity from the ViT's intermediate features, DTEM learns a lightweight per-block linear embedding module whose cosine similarities drive the merging policy. Because the grouping and merging operators are discrete, the module is trained through a continuous relaxation of ToMe's bipartite soft matching: soft grouping via a differentiable top-k operator and soft merging via weighted averaging with size tracking. At inference, the hard BSM operator is used. DTEM is evaluated on ImageNet classification (modular and end-to-end), COCO captioning, and ADE20K segmentation, claiming consistent improvements over ToMe and other baselines at matched FLOPs, typically +0.1 to +0.5 accuracy points with larger gains at higher reduction rates. The paper includes ablations showing the decoupled embedding is essential, and that a single model trained at a high reduction rate transfers to lower rates.","tokens_in":19289,"tokens_out":6965,"duration_ms":59580,"significance":"If the results hold, DTEM is a useful and simple add-on for making frozen ViTs faster with a small trained module, and it demonstrates a practical way to train through a continuous relaxation of token merging. The contribution is incremental but clearly positive: the gains are small in classification but more substantial in captioning and segmentation. Strengths: code is released; the experimental coverage is broad (three tasks, five backbones, both modular and end-to-end training); and the ablation in Table 6 directly shows that naive soft merging on raw keys hurts, which supports the core design. The main risks are the unverified soft-to-hard transfer of the relaxation, the lack of error bars for small classification differences, and the training-status imbalance in headline comparisons to ToMe.","major_comments":[{"comment":"The paper's central novelty is learning the decoupled embedding through the continuously relaxed grouping and merging operators (Eqs. 6-10), but the only justification that this surrogate matches the deployed hard BSM is an asymptotic statement as tau -> 0. Training uses tau = 0.1 and additionally scales the similarity matrix by 0.1 (Sec. 4.1, Implementation details), so the effective softmax temperature is near 1, far from the asymptotic regime. The manuscript never measures the agreement between the soft adjacency matrix E~ and the hard BSM adjacency E' at the trained embeddings. Without such a measurement (or a comparison to a straight-through estimator of BSM), the claim that the soft relaxation is what confers the improvement is not established. Please add an experiment quantifying the soft-to-hard gap, e.g., the proportion of selected edges that coincide, or a training variant that optimizes hard BSM directly.","section":"Secs. 3.2-3.4"},{"comment":"The headline classification gains are small: in Figure 2 the improvements over ToMe are 0.12-0.2% for DeiT-S, and in Table 3, 0.12-0.23% over ToMe at matched FLOPs. No multiple seeds, confidence intervals, or significance tests are reported anywhere in the paper. Given that the reported gains are of the same order as typical run-to-run variance for ImageNet fine-tuning, the central claim of 'consistently improves' is not statistically supported as presented. Please report at least three seeds with standard deviations for the key comparisons (Tables 1, 3, and Figure 2), and state whether the improvements are significant.","section":"Sec. 4.1, Tables 1-3, Fig. 2"},{"comment":"Table 1 compares DTEM against ToMe and EViT in the 'off-the-shelf frozen' setting, but ToMe and EViT as used there are training-free, while DTEM has been trained on the target task (ImageNet) for 30 epochs. The improvement may therefore reflect the benefit of task-specific training of the merging policy rather than the decoupled-embedding design per se. The controls in Table 8 (ToMe + Gumbel-Softmax + soft merging) are the right type of comparison, but they are relegated to an ablation and use a different training setup. Please include a trained ToMe-based alternative (e.g., ToMe+GS or a trained linear similarity head) in the main comparison table, or clearly state the training budgets and frame Table 1 as 'trained vs. training-free'.","section":"Table 1 and Sec. 4.1"}],"minor_comments":[{"comment":"The notation in Eq. (7) is ambiguous: the update log(1 - sum_j a^t_ij) suppresses rows with any soft mass, but the text says it suppresses 'the entire outbounding edges from the softly selected nodes in A'; please clarify that the sum is over j in B and that the update applies to all rows i, not only the selected ones.","section":"Eq. (7)"},{"comment":"The appendix uses 'dTPS' while the main text and reference [33] use 'eTPS'; please unify the abbreviation.","section":"Table 14"},{"comment":"In the appendix, the sentence 'In Table 12, we report Figure 4 results across a broader reduction range' should read 'Table 4 results'.","section":"A.1"},{"comment":"The description of how zero-size tokens are handled differs between Sec. 3.3 ('such tokens will be excluded from the subsequent merging process by Eq. 9') and Appendix A.2, which says the r tokens with minimum effective size are successively excluded each block; please reconcile these statements.","section":"Sec. 3.3 and A.2"},{"comment":"The column header 'Reduction' is described as 'roughly represents the decreases in FLOPs', but the FLOPs differ across methods at the same nominal reduction (e.g., DeiT-S 35%: ToMe 3.02 GFLOPs vs DTEM 2.91 GFLOPs); consider reporting exact FLOPs-reduction percentages per method rather than a single nominal value.","section":"Table 1"},{"comment":"The captioning comparisons are at slightly different reduction levels (e.g., ToMe 32% vs DTEM 31%); since CIDEr is sensitive to the exact token count, please match reduction levels more closely or interpolate the baseline curve.","section":"Sec. 4.2"}],"recommendation":"major_revision","confidential_remarks":"The paper is within scope for a vision/ML venue. The main technical uncertainty is whether the continuously relaxed merging is actually the source of the gains; an empirical soft-to-hard gap measurement would resolve this. The comparison fairness issue with ToMe (trained vs. training-free) should be addressed in revision. No concerns about citation patterns or novelty disclosure."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: this is a solid, modest empirical contribution to token merging for ViTs. The core idea is to learn a small per-block embedding module that produces features dedicated to merging, trained through a continuous relaxation of bipartite soft matching, then discretized to ToMe-style hard merging at inference. It consistently beats ToMe by 0.1 to 1.6 accuracy points across ImageNet, COCO captioning, and ADE20K segmentation, and it works modularly with frozen backbones, which is genuinely useful.\n\nThe decoupled embedding is new, and the relaxation of greedy top-k matching is a clever way to make the grouping operator differentiable. The ablation in Table 6 is the right one: soft merging on attention keys hurts, and adding the decoupled embedding recovers and improves beyond ToMe. The data/train efficiency results (Fig. 4) are a nice practical bonus.\n\nSoft spots: effect sizes are small, especially in end-to-end classification, and there are no error bars anywhere. Table 1 compares ToMe and DTEM at slightly different FLOPs (DTEM is lower at the same nominal reduction), so the modular gains are a bit overstated; the FLOPs-matched curves in Fig. 2 show more honest gains around 0.1\\u20130.2 points. The paper never compares against PatchMerger or other learned-merging baselines, which weakens the novelty claim. The load-bearing assumption \\u2014 that the continuous relaxation at tau=0.1 transfers to hard BSM \\u2014 is supported only by an asymptotic argument. Since the implementation scales similarity by 0.1 with tau=0.1, the effective softmax temperature may be near 1, not near zero. There is no experiment measuring the soft-to-hard gap, so it is unclear whether gains come from the relaxation or from the extra parameters / decoupled feature space. That is a real uncertainty, but not fatal: the method works empirically, and the relaxation is a standard technique.\n\nVerdict: worth a serious referee. I would send it to review with requests for matched-FLOPs tables, error bars, a PatchMerger comparison, and at least one experiment that anneals tau or compares against straight-through training. Who should read it: anyone working on token reduction or efficient ViTs. Recommendation: engage, but push for those additions.","headline":"A modest but consistent improvement over ToMe via learned decoupled merging embeddings; the main open question is whether the soft training surrogate really transfers to hard inference as claimed.","tokens_in":19782,"tokens_out":6165,"would_cite":true,"duration_ms":54395,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"The paper claims that token merging in Vision Transformers should be driven by a separate learned embedding, and that training this embedding through a continuous relaxation of the merging step improves the speed-accuracy trade-off of…","keywords":["Vision Transformers","token merging","token reduction","decoupled embedding","continuous relaxation","efficient inference","image classification","semantic segmentation"],"falsifier":"Train DTEM modularly on ImageNet at $\\tau=0.1$, then measure agreement between the token pairs selected by hard bipartite matching at inference and the pairs receiving the largest soft weights during training; if the agreement is low, or if lowering the training temperature to $0.01$ does not raise both agreement and final accuracy, the surrogate is not faithful to the inference-time merging policy.","tokens_in":18780,"feed_emoji":"⚡","tokens_out":8959,"duration_ms":78724,"temperature":0.7,"pith_summary":"Token merging makes Vision Transformers faster by combining similar tokens, but existing methods decide what to merge using the same intermediate features the ViT uses for encoding. This paper claims those features are the wrong tool: the merging policy needs its own representation, learned specifically for merging. DTEM adds a per-block linear projection that maps detached ViT features into a small similarity space, and trains that projection by replacing the hard matching step with a soft, differentiable version. If the claim is right, a frozen pre-trained ViT can be sped up by training only this small module, with consistent accuracy gains over ToMe across classification, captioning, and segmentation. The paper also notes that it has so far been validated only on vision tasks and does not reduce training-time compute.","feed_headline":"Decoupled embeddings lift ViT token merging at matched FLOPs","feed_subtitle":"DTEM trains a tiny separate embedding so frozen ViTs cut up to 37% FLOPs while beating ToMe on accuracy.","key_machinery":"The central object is the decoupled embedding module, $Z = f(X; \\phi_l)$, one affine layer per block that maps the block input to a low-dimensional similarity space ($d'=64$ for ViT-S/B, $128$ for ViT-L) and defines token similarity by cosine distance. The training mechanism is a continuous relaxation of Bipartite Soft Matching: a soft top-$r$ selection built from temperature-softmax steps with suppression of already-selected source nodes (Eqs. 6-7), normalized to a soft adjacency matrix (Eq. 8), followed by soft merging updates that reweight features and effective token sizes (Eqs. 9-10). This relaxation is what allows gradients from the task loss to reach the merging policy, and as $\\tau \\to 0$ it converges to the discrete matching used at inference. The paper also uses detached ViT features as input to the embedding module, so the module alone is trained by the relaxed merging signal.","core_discovery":"The central claim is that a decoupled, per-block token embedding—computed by a single linear layer from detached ViT features and used only to compute cosine similarities—learns a better merging policy than the ViT's own keys or token embeddings. To make this learnable, DTEM replaces ToMe's hard bipartite soft matching with a temperature-softmax relaxation of sequential top-r selection, normalized into a soft adjacency matrix, and merges tokens with asynchronous weighted updates. At inference the soft operators are swapped for hard bipartite soft matching, so the method keeps ToMe's speed while changing which tokens get merged. In modular training on frozen DeiT-S/B and MAE-B/L models, DTEM reports +0.15 to +0.47 accuracy over ToMe at 35% FLOP reduction and +0.47 to +1.64 at 50% reduction; in end-to-end fine-tuning on DeiT-S it reaches 79.85% top-1 at 2.9 GFLOPs, a 37.2% FLOP cut. The same recipe improves CIDEr scores by +2.3 to +6.0 in captioning and mIoU by +0.32 to +1.3 in segmentation at matched compute.","pith_inferences":["The soft-to-hard transfer gap is the main extension risk: because training runs at $\\tau=0.1$ while inference is hard matching, the reported gains could shrink if the learned embedding overfits the relaxed operator.","The same decoupling idea could be dropped into other token-reduction mechanisms—clustering-based merging or importance-based pruning—by replacing only the similarity source with a trained detached projection.","The visualization evidence suggests the learned embedding encodes a foreground/background preference; that signal could be reused as a saliency map or combined with a semantic loss to make merging more content-aware.","Jointly learning which block drops how many tokens (a reduction profile) together with the decoupled embedding is a natural next step that the paper identifies as orthogonal."],"forward_implications":["A frozen pre-trained ViT can be made faster without fine-tuning its weights: only the per-block embedding modules are trained, and the paper reports gains starting from 0.31% of ImageNet and after one epoch.","A single DTEM model trained at reduction rate $r$ can be evaluated at any lower rate, so one training run covers multiple inference FLOPs budgets.","Because the embedding module adds less than 1% of parameters and FLOPs, the accuracy gain is nearly pure policy improvement rather than extra capacity.","The decoupled-policy recipe transfers across tasks: classification, image captioning, and semantic segmentation all show better accuracy at matched compute than merging with backbone features.","In end-to-end fine-tuning, DTEM improves on both ToMe and EViT at equal FLOPs and also matches longer-trained pruning methods at 100 epochs with only the task loss."],"supporting_citations":[{"why":"Defines Bipartite Soft Matching and ToMe, the target grouping operator, baseline, and inference procedure that DTEM inherits.","marker":"[2]"},{"why":"EViT is the main comparison baseline for modular and end-to-end token reduction, and its fine-tuning setup is reused.","marker":"[18]"},{"why":"Supplies the differentiable top-k operator that DTEM's soft grouping revises.","marker":"[36]"},{"why":"Provides the DeiT pre-trained models and training recipe used for classification experiments.","marker":"[29]"},{"why":"Provides the MAE pre-trained models used to test DTEM on frozen backbones.","marker":"[11]"},{"why":"Supplies the merge-then-unmerge approach for applying token merging to semantic segmentation components.","marker":"[3]"},{"why":"Segmenter is the segmentation model and dataset setup on which DTEM's mIoU gains are measured.","marker":"[28]"},{"why":"GIT is the captioning model whose frozen encoder is modularly adapted by DTEM.","marker":"[32]"}],"fun_headline_variants":["DTEM: Decoupled tokens beat ViT keys for merging","Frozen ViTs cut 37% FLOPs with decoupled merge embeddings","Soft relaxed merging learns better token policies for ViTs","Modular DTEM lifts ViT accuracy at 35-50% FLOP cuts"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the softened merging decisions used during training are a faithful stand-in for the hard merging decisions made at inference; the paper shows this only as the smoothing temperature goes to zero, while training runs at a nonzero temperature ($\\tau = 0.1$).","fun_headline_variants_meta":{"raw":{"variants":["DTEM: Decoupled tokens beat ViT keys for merging","Frozen ViTs cut 37% FLOPs with decoupled merge embeddings","Soft relaxed merging learns better token policies for ViTs","Modular DTEM lifts ViT accuracy at 35-50% FLOP cuts"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000133,"raw_usage":{"total_tokens":1193,"prompt_tokens":1062,"completion_tokens":131,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":678,"completion_tokens_details":{"reasoning_tokens":65}},"tokens_in":678,"tokens_out":131,"duration_ms":2102,"temperature":1.0,"reasoning_tokens":65,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T15:50:47.616614+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train DTEM modularly on ImageNet at $\\tau=0.1$, then measure agreement between the token pairs selected by hard bipartite matching at inference and the pairs receiving the largest soft weights during training; if the agreement is low, or if lowering the training temperature to $0.01$ does not raise both agreement and final accuracy, the surrogate is not faithful to the inference-time merging policy.","supporting_citations":[{"cited_title":"Bolya, C.-Y","cited_arxiv_id":null,"evidence_quote":"Defines Bipartite Soft Matching and ToMe, the target grouping operator, baseline, and inference procedure that DTEM inherits."},{"cited_title":"Liang, C","cited_arxiv_id":null,"evidence_quote":"EViT is the main comparison baseline for modular and end-to-end token reduction, and its fine-tuning setup is reused."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the differentiable top-k operator that DTEM's soft grouping revises."},{"cited_title":"Touvron, M","cited_arxiv_id":null,"evidence_quote":"Provides the DeiT pre-trained models and training recipe used for classification experiments."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Provides the MAE pre-trained models used to test DTEM on frozen backbones."},{"cited_title":"Bolya and J","cited_arxiv_id":null,"evidence_quote":"Supplies the merge-then-unmerge approach for applying token merging to semantic segmentation components."},{"cited_title":"Strudel, R","cited_arxiv_id":null,"evidence_quote":"Segmenter is the segmentation model and dataset setup on which DTEM's mIoU gains are measured."}],"review_version":1}