{"id":"f32769fb-cc75-4e50-bdb3-d10eb5272c74","arxiv_id":"2608.01298","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"UDT, a U-Net diffusion transformer with data-adaptive token merging for down/upsampling, improves ImageNet generation FID while cutting training epochs from 1400 to 40 compared with SiT-XL.","lead":"A new U-Net-shaped diffusion transformer, UDT, uses token merging for downsampling and upsampling instead of fixed spatial pooling. On ImageNet 256x256 it reaches better FID than the SiT baseline with up to 40x fewer training epochs.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Headline 40x convergence claim rests on single-run FID point estimates with no seeds or checkpoints; a 0.3 FID margin can be within run-to-run noise, so the central claim is not yet pinned down.","rationale":"The reader's weakest_assumption focuses on whether token merging preserves fine-grained information; the paper's Table 1(b) and ToMe's prior dense-prediction results make that mechanism plausible, and I do not see a clear internal contradiction there. However, the reader's rationale also flags single-run metrics, missing error bars, and unpinned code, and that is the concern I find most load-bearing for the headline claim. The claimed 40x convergence speedup is a point estimate with a 0.3 FID margin over its reference, so seed variance alone could determine whether the headline holds. Because the paper provides no repeated-seed evidence, no released checkpoints, and only unpinned code, the conditional verdict remains appropriate. A three-seed reproduction of the single most important 40-epoch result, plus a re-evaluation of the SiT reference under the same harness, would settle this concern and could justify moving to accept.","tokens_in":22845,"tokens_out":16804,"duration_ms":149306,"concrete_test":"Re-run the exact UDT-XL/2+REPA recipe from Appendix A for 40 epochs with three independent training seeds (varying only the data-shuffle seed), evaluate FID using the paper's ADM protocol (50K samples, eval seed 0), and recompute the SiT-XL/2 1400-epoch reference with the same evaluation harness. If the mean 40-epoch FID plus one standard deviation is below 7.9, the speedup claim is supported; if the interval overlaps or exceeds 7.9, the central claim must be restated as comparable instead of 40x faster.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central quantitative claim compares one UDT-XL/2+REPA run at 40 epochs (FID 7.6, no CFG) against a literature SiT-XL/2 result at 1400 epochs (FID 7.9). Section 4.2 and Tables 2/4 report only single runs; Appendix A.2 gives the evaluation protocol but no training-seed variation, error bars, or confidence intervals. The claimed advantage is a 0.3 FID margin, which is the same order as typical seed-to-seed FID variation for class-conditional ImageNet 256 training at this scale. If a second seed moves the 40-epoch FID above 7.9, the headline '40x faster convergence' is no longer supported and the result becomes 'comparable performance in far fewer epochs,' which is a materially weaker claim. The same single-run structure underlies the 80-epoch table entries and the CFG comparisons in Table 3. This is not an objection to the architecture: Table 1(b) already shows data-adaptive merging substantially outperforms a learned projection in the same U-shape, so the mechanism is plausible. The unsecured step is specifically the statistical reliability of the headline speedup number, which the paper does not address.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"This paper introduces UDT, a U-Net-shaped diffusion transformer that replaces fixed learnable spatial downsampling with data-adaptive token merging (ToMe) and exact unmerging, while preserving the token hidden dimension. The authors report that UDT accelerates convergence substantially, e.g., UDT-XL/2+REPA reaches FID 7.6 at 40 epochs without CFG versus SiT-XL/2's 7.9 FID at 1400 epochs (~40x), and achieves strong CFG results (FID 1.38 at 320 epochs with SD-VAE and 1.35 at 500 epochs with VA-VAE). The paper includes ablations of NMerge, r schedules, merging components, advanced techniques, and a range of drop-in replacements (DiT, JiT, LightningDiT, MMDiT), plus reduced-data and 512x512 experiments.","tokens_in":23054,"tokens_out":8269,"duration_ms":71952,"significance":"The core architectural idea is simple and plausible: merging semantically similar tokens is a transformer-native way to create a U-Net hierarchy, and the paper provides credible evidence that this outperforms learned projection-based downsampling in the same U-shape (Table 1b), as well as strong results across model sizes. If the convergence claims survive statistical scrutiny, the contribution is significant because it offers a drop-in, parameter-free (in the learned-parameter sense) U-Net DiT backbone compatible with REPA and T2I models. The paper also ships code and extensive experiments, which is a strength. The main caveat is that the headline quantitative claims rest on single-run FID point estimates without repeated-seed variation, so the magnitudes of the speedups are not yet pinned down.","major_comments":[{"comment":"The central speedup claim (7.6 FID at 40 epochs vs 7.9 FID at 1400 epochs) is supported only by single-run FID estimates; Appendix A.2 specifies seed=0 for evaluation but no training-seed variation, standard deviations, or confidence intervals are reported anywhere. A 0.3 FID margin is of the same order as typical run-to-run variation for class-conditional ImageNet training at this scale, and the 40-epoch point appears only in a figure curve rather than in a numerical table. Please add at least three training seeds for the headline configuration (and ideally for Tables 2 and 4), report mean +/- std or confidence intervals, and provide a checkpoint-level table for the 40/80/320/500 epoch numbers. Without this, the '40x faster convergence' claim is not statistically distinguishable from 'comparable performance in far fewer epochs.'","section":"Section 4.2, Tables 2-4, Fig. 1(b), Appendix A.2"},{"comment":"The 7.9 FID baseline is ambiguous. The abstract attributes 7.9 to 'SiT ... at 1400 epochs (w/o CFG)', while Table 4 reports SiT-XL/2+REPA at 800 epochs with FID 7.9, and Section 4.2 says both '20-40x faster' and 'outperforming SiT-XL/2 + REPA'. Please state explicitly which checkpoint and configuration (SiT vs SiT+REPA, 800 vs 1400 epochs, CFG/no-CFG) is used for each speedup factor, and compute the factors consistently. If the relevant baseline is SiT-XL/2 at 1400 epochs, the source of that number should be cited rather than inferred from Table 4.","section":"Abstract and Section 4.2"},{"comment":"The r-schedule notation contains apparent typos that make the architecture specification incomplete; for example, UDT-L/2 reads '15 (Enc 2-5), 12 (Enc 2-13)' and UDT-XL/2 reads '12 (Enc 2), 14 (Enc 6-11)', leaving the intervening encoder blocks unspecified. Please list the exact per-block r values for every model size, since the schedule is a central design choice and the paper elsewhere emphasizes that the improvement is purely architectural.","section":"Appendix A.1, Table 10"}],"minor_comments":[{"comment":"Please distinguish curves by markers or line styles in addition to color, since color-only distinction is hard to read and the figures may be printed in grayscale.","section":"Fig. 1(b) and Fig. 4"},{"comment":"State the training seed(s) used for each run, not only the evaluation seed, and report the number of training runs that each reported FID is based on.","section":"Appendix A.2"},{"comment":"Indicate clearly that the reported 'epochs' are training epochs and that FID is computed on 50K samples without class-balanced sampling unless stated; Appendix G should be referenced from the caption.","section":"Table 3 caption"},{"comment":"The term 'auto-guidance' appears in the figure caption but is never defined in the text; either define it or remove the reference.","section":"Appendix G, Fig. 7 caption"},{"comment":"The 'data-adaptive' terminology is used for a fixed key-similarity heuristic adopted from ToMe; please clarify that no per-sample learned adaptation is involved, and consider adding a convolutional projection baseline in Table 1(b) to strengthen the comparison against learned downsampling.","section":"Section 3 and Table 1(b)"}],"recommendation":"major_revision","confidential_remarks":"The single-run issue is the key obstacle. I would be willing to accept after repeated-seed evidence is added for the headline claims; the paper's broad experiments are valuable, and the mechanism is credible. No concerns about novelty or fit with the journal's scope."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Key point: this is a real architectural idea, not a repackaging. Using ToMe token merging as the down/upsampling operator in a U-Net-shaped diffusion transformer, while preserving token hidden dimension, is new and well-motivated. The authors show data-adaptive merging beats a learned projection in the same U-shape (Table 1b) and beats fixed spatial downsampling in representation quality (Figure 2), and the effect carries across model sizes, patch sizes, VAE choices, and even MMDiT. That breadth is the paper's strength.\n\nThe load-bearing claim, though, is the 40x convergence speedup: UDT-XL/2+REPA reaches 7.6 FID at 40 epochs without CFG versus SiT-XL/2's 7.9 at 1400 epochs. Every FID number in the paper is a single run with no seed variation, and a 0.3 FID gap is within typical run-to-run noise at this scale. So the exact factor is not established. The stress-test note is right that a single seed could move the number. But it's worth keeping perspective: the paper also reports 7.7 FID at 80 epochs and 7.0 at 60 epochs without REPA, both still beating the 7.9 baseline by more than 0.3. Even if the 40-epoch run bounces, the convergence advantage is real; only the precise factor is shaky.\n\nThe weaker step is hyperparameter sensitivity. NMerge and the r schedule are tuned per model, and the paper doesn't report how sensitive FID is to those choices. Since the architecture's whole point is that merging is data-adaptive, I'd like to see an ablation where r is varied at fixed model size, and whether the baselines got comparable tuning. Also, the appendix gives a protocol but code is not pinned to a commit and no checkpoints are released, which makes independent reproduction harder.\n\nOn the core mechanism: yes, merging by key similarity and unmerging by copy-back can in principle drop detail-critical tokens. But the design keeps full resolution in the first and last blocks, uses skip connections, and the ablations support the design. I don't see this as a load-bearing flaw, just an assumption that should be stress-tested with high-resolution or fine-detail datasets.\n\nWho this is for: anyone working on efficient DiT training or U-Net transformer hybrids. It deserves a serious referee, and I'd accept it with the expectation of revisions asking for multi-seed results, sensitivity analysis, and checkpoints. The central architectural claim is likely to survive; the headline number needs pinning down.","headline":"Real architectural contribution with an over-sold headline: single-run FIDs make the exact 40x speedup unverified, but the convergence story holds up.","tokens_in":23650,"tokens_out":2473,"would_cite":true,"duration_ms":22254,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"UDT's token-merging U-Net reaches SiT's 1400-epoch FID in 40 epochs","keywords":["diffusion transformers","token merging","U-Net architecture","image generation","representation alignment","flow matching","training convergence","ImageNet generation"],"falsifier":"Measure the copy-back discrepancy at the bottleneck: unmerge tokens with the recorded indices and compare them with the original full-resolution features, on images dominated by fine texture such as fur, fabric, or foliage. If the reconstruction error grows sharply with merge rate or texture density, and FID degrades correspondingly, the information-preservation assumption fails; if unmerged features are near-identical, the assumption holds.","tokens_in":22601,"feed_emoji":"🧩","tokens_out":5828,"duration_ms":48139,"temperature":0.7,"pith_summary":"The paper argues that diffusion transformers can gain the convergence benefits of U-Nets without learnable spatial downsampling: replace fixed grid pooling with data-adaptive token merging in a U-shaped encoder-decoder, while keeping the token dimension constant. It claims that this single architectural change makes a flow-based transformer on ImageNet 256x256 reach, without classifier-free guidance, FID 7.6 after 40 epochs, where the isotropic baseline needs 1400 epochs for FID 7.9. With classifier-free guidance and representation alignment, it reports FID 1.38 after 320 epochs. The broader point is that representation quality and denoising capacity can both improve from an architecture that compresses redundant tokens instead of blurring neighboring ones.","feed_headline":"Token-merging U-Net hits 1400-epoch FID in 40","feed_subtitle":"Data-adaptive token merging preserves detail while cutting compute, reaching 7.6 FID at 40 epochs with no guidance.","key_machinery":"The central object is Token Merging (ToMe), here repurposed from an inference-speed technique into the down/upsampling operator of a U-shape transformer. It partitions tokens by bipartite soft matching, merges the most similar key-pairs with size-weighted averaging, applies proportional attention to correct for merged-token sizes, and unmerges by copying back along recorded indices. In UDT it progressively reduces tokens from 256 to a bottleneck of 112 ($N_{\\text{Merge}} = 112$) with the hidden dimension kept at the DiT value $D$, then restores tokens symmetrically; the first encoder and last decoder block stay at full resolution and are skip-connected.","core_discovery":"The central claim is that a U-Net-shaped diffusion transformer whose downsampling and upsampling are implemented by token merging and unmerging preserves the DiT's isotropic token dimension and self-attention dynamics while giving it a true encoder-decoder hierarchy. Merging is driven by key similarity: redundant tokens such as backgrounds and flat regions are fused by weighted averaging, their merge indices are recorded, and the decoder copies merged tokens back to their original positions with skip connections carrying full-resolution information. The authors show that this beats both isotropic DiTs and earlier U-Net DiTs that use fixed 2x2 neighborhood downsamples with learnable projections, and that it aligns naturally with representation alignment because bottleneck features can be unmerged back to full token resolution for patch-wise matching.","pith_inferences":["Beyond the paper: a testable reading is that data adaptivity, not mere resolution reduction, drives the gain; ablating token merging with random pooling at the same rates would settle whether similarity-based fusion or just hierarchy matters.","Beyond the paper: the authors explicitly leave video generation and 2K resolution untested, but because token merging adapts to redundancy, one would expect the largest gains on high-resolution images where backgrounds occupy most of the frame; that expectation is not supported by the paper's evidence.","Beyond the paper: the copy-back unmerge means every original token still receives a prediction, so the merge indices naturally give a pooling/unpooling pair that could be reused by dense prediction heads or segmentation objectives.","Beyond the paper: the 10%-data result hints that merging acts as a structural prior for scarce-data regimes; a controlled experiment varying dataset size would test whether the convergence advantage grows as data shrinks."],"forward_implications":["UDT-XL/2+ with REPA reaches FID 7.6 without classifier-free guidance at 40 epochs, roughly 40x faster convergence than the SiT-XL/2 baseline's 1400 epochs, and 7.7 FID out-of-the-box at 80 epochs.","With classifier-free guidance the model reports FID 1.38 after 320 epochs using the standard latent VAE and 1.35 after 500 epochs with an improved VAE, competitive with models trained two to three times longer.","Because the down/upsampling is parameter-free and keeps the token dimension, the same blocks can replace DiT, pixel-space JiT, and the visual branch of MMDiT, improving FID in each case.","Longer token sequences such as patch size 1 and 512x512 images are handled by aggressive early merging, with roughly 1.3-2.2x cost increase instead of 4.5x, reaching FID 1.71 at 512x512 from scratch.","On 10% of ImageNet, UDT-L/2 reaches FID 10.6 at 500 epochs, beating SiT-L/2 trained on the full dataset at 80 epochs while using fewer training images."],"supporting_citations":[{"why":"Supplies the token-merging algorithm (bipartite soft matching, weighted averaging, proportional attention, index-based unmerging) that UDT uses for down/upsampling.","marker":"[4]"},{"why":"Defines the SiT baseline architecture and training protocol that UDT modifies and compares against.","marker":"[35]"},{"why":"Provides the REPA representation-alignment method and the setup UDT is shown to be compatible with.","marker":"[59]"},{"why":"Introduces the U-shaped SiT-down variant with fixed spatial downsampling and UREPA, the prior art UDT compares against.","marker":"[47]"},{"why":"Introduces U-DiT, the earlier U-Net DiT with learnable fixed-neighborhood downsampling that UDT surpasses.","marker":"[48]"},{"why":"Shows token merging can be applied to diffusion models for faster inference, the precedent UDT extends to training.","marker":"[5]"},{"why":"Supplies the improved VAE (VA-VAE) used in UDT's strongest FID results.","marker":"[57]"}],"fun_headline_variants":["UDT: token-merging DiT cuts training to 40 epochs","Data-adaptive token merging gives DiTs U-Net speed","U-Net DiT via token merging: 40-epoch FID match","Token-adaptive UDT: DiT speed with U-Net convergence","Merging tokens, not resolutions: fast DiT training"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that merging tokens by key similarity and later copying them back preserves the fine-grained information a diffusion model must reconstruct; if merging irreversibly discards detail-critical tokens, the encoder-decoder would lose exactly what later layers need.","fun_headline_variants_meta":{"raw":{"variants":["UDT: token-merging DiT cuts training to 40 epochs","Data-adaptive token merging gives DiTs U-Net speed","U-Net DiT via token merging: 40-epoch FID match","Token-adaptive UDT: DiT speed with U-Net convergence","Merging tokens, not resolutions: fast DiT training"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000242,"raw_usage":{"total_tokens":1570,"prompt_tokens":1035,"completion_tokens":535,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":651,"completion_tokens_details":{"reasoning_tokens":441}},"tokens_in":651,"tokens_out":535,"duration_ms":5319,"temperature":1.0,"reasoning_tokens":441,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T15:08:10.677373+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Measure the copy-back discrepancy at the bottleneck: unmerge tokens with the recorded indices and compare them with the original full-resolution features, on images dominated by fine texture such as fur, fabric, or foliage. If the reconstruction error grows sharply with merge rate or texture density, and FID degrades correspondingly, the information-preservation assumption fails; if unmerged features are near-identical, the assumption holds.","supporting_citations":[{"cited_title":"Bolya, C.-Y","cited_arxiv_id":null,"evidence_quote":"Supplies the token-merging algorithm (bipartite soft matching, weighted averaging, proportional attention, index-based unmerging) that UDT uses for down/upsampling."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines the SiT baseline architecture and training protocol that UDT modifies and compares against."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Provides the REPA representation-alignment method and the setup UDT is shown to be compatible with."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Introduces the U-shaped SiT-down variant with fixed spatial downsampling and UREPA, the prior art UDT compares against."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Introduces U-DiT, the earlier U-Net DiT with learnable fixed-neighborhood downsampling that UDT surpasses."},{"cited_title":"Bolya and J","cited_arxiv_id":null,"evidence_quote":"Shows token merging can be applied to diffusion models for faster inference, the precedent UDT extends to training."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the improved VAE (VA-VAE) used in UDT's strongest FID results."}],"review_version":2}