Sampling Procedure

[messimm]

Hi! Great work — I see the results look solid both in the paper and the repo.

One question about train–inference mismatch: during training the student sees DTS inputs with per-token heterogeneous timesteps (e.g., a fixed fraction of tokens are “cleaner” / less noised), so the attention layers consistently have access to some less-corrupted tokens as context. At inference, however, we revert to the standard homogeneous timestep setting, i.e. τ=t1 for all tokens.

What’s the key reason this generalizes without issues in practice? Do you have an intuition or an ablation (e.g., occasionally training with fully homogeneous timesteps / mask ratio = 0, or testing DTS-style inference) that supports robustness to this shift?

[hila-chefer]

Hi @messimm, thanks for your interest!
We actually discuss this in the motivation section of the paper (see attached figure). To formulate an effective self-supervised task, one must balance between introducing asymmetry and maintaining the generative behavior. DTS (unlike "vanilla" masking or diffusion forcing) offers a "continuous masking" mechanism. This means that during training, two timesteps are drawn from $p(t)$. They could be very distinct (making the task semantically challenging) or very close (maintaining the uniform noising behavior). This allows us to model all continuous options for dual timesteps, naturally bridging the gap between highly distinct steps and nearly homogeneous ones.

To answer your question about ablations: in early versions, we actually explored mixing between DTS and uniform timesteps, and found the performance to decrease. This aligns with the attached figure showing that DTS alone slightly improves generation quality, intuitively because the heterogeneous setting encourages the formation of stronger global relations.

[messimm]

Hi, thanks a lot for the detailed clarification — that makes sense, and I buy the empirical ablations.

I still find one aspect intriguing, especially from a video-generation perspective. In many video diffusion/FM setups we intentionally bias the timestep distribution toward high-noise regions (e.g., via a “scheduler scale” / reweighting of p(t)) so that highly noised samples appear more often during training. This typically helps the model learn the low-frequency / global structure early, which is crucial for semantics and physical plausibility.

With DTS, I think I understand your “continuous masking” intuition: since you draw two timesteps t,s ~ p(t), the model spans a continuum from highly distinct to nearly homogeneous noise patterns. That said, it seems this could reduce the frequency of homogeneous extreme-noise exposures. If a “very high-noise” timestep occurs with probability p under p(t), then seeing an all-tokens-at-very-high-noise case requires both draws landing in that region, which scales like p^2. Since p^2 << p for p < 1 (e.g., 0.9 -> 0.9^2 = 0.81), those extreme-noise homogeneous cases may become noticeably rarer.

Two questions:

1. How should we think about the effective noise exposure under DTS?
   Do you analyze/track the marginal distribution of per-token timesteps (or statistics like min(tau), max(tau)) compared to the base p(t)? In other words: does DTS preserve the amount of training signal at very high noise, or does it effectively reduce it in a way that is still beneficial?

2. Implementation detail (Wan 2.2): how do you inject two timesteps / per-token timesteps into an architecture that expects a single global timestep embedding?
   Did you modify the timestep conditioning to be token-wise (compute e(t) per token and add it to token embeddings / AdaLN parameters), or use some other trick (e.g., combine embeddings, separate heads, etc.)?

This phenomenon is really interesting and seems slightly counterintuitive given common timestep reweighting practices in video models — would love any extra intuition or pointers.

[hila-chefer]

Importantly, the marginal distribution per token remains the same since $t,s$ are both iid from $p(t)$. So everything kind of boils down to your selection of $p(t)$. If one wants more high noises (or more high noises after masking), the simple solution is modifying $p(t)$ to emphasize high noises even more, so the training dynamic depends on your noise selection in DTS same as with the vanilla framework.
In the paper, we tune the noise scheduling using the vanilla flow matching framework and apply the same scheduler to all baselines to show robustness and maintain a fair comparison.

[betterze]

Thank you for the helpful discussion.

I have the same confusion regarding Question 2: how is the time condition injected into the model? Since the input x_t​ involves two time steps, do you use token-wise noise levels?

Based on your response and the code, it seems that the model receives only a single time-step condition, which is injected through AdaIN and used to modulate all tokens. If that is the case, which time step is actually provided to the model—the larger t or the smaller t?

Thank you very much for the clarification.

[hila-chefer]

Hi @betterze, thanks for your interest!
This is actually highlighted in the code, the paper and the README:
implementation implementation p2
README highlights several times that there’s per token conditioning eg here
In the paper, we provide implementation details on the architecture (Supp A.3), that’s the first thing mentioned there on the architecture.

During inference you do have a single timestep, which is also injected in the code per token.

[betterze]

found it. I really appreciate your help.

[YaohouF]

Hi @hila-chefer, thank you for the inspiring work.

I have a follow-up question regarding the dual-timestep sampling. Before Eq. (4), when sampling s,t∼p(t), do you sort them before constructing the student input, so that a fixed fraction of tokens always receives the noisier timestep (similar to MAE-style masking)? By reading the section 3.3 I assumed we should not sort them for maintaining marginal timestep distribution?

[hila-chefer]

Hi @YaohouF, thanks for your interest!
Your understanding is correct- we do not sort them to maintain the marginal distribution per token.
This differs from traditional “masking” and creates two states for the self supervised objective— “local to global” where most tokens are heavily noised and the model needs to infer a lot of missing information, and “global to local” where most information is clean, and the model is required to make up local details that are missing.