Trajectory Divergence Map (TDM) for Shape-Aware Editing
- TDM is a latent-space divergence map computed from token-wise velocity differences, highlighting regions for large structural edits.
- It guides selective KV feature injection by comparing source-image and target-prompt trajectories, thus enabling precise shape-aware modifications.
- The method is training-free and mask-free, preserving non-target areas while supporting extensive geometric transformations.
Searching arXiv for the specified paper and closely related uses of similar terminology. Trajectory Divergence Map (TDM) is the central region-localization mechanism in "Follow-Your-Shape: Shape-Aware Image Editing via Trajectory-Guided Region Control" for prompt-driven, shape-aware image editing in a flow-based generative model. In that framework, TDM is computed by comparing token-wise velocity differences between a source-image inversion trajectory and a target-prompt editing trajectory, then converting those differences into a spatial edit map that guides selective feature injection. The method is described as training-free and mask-free, and its stated purpose is to enable large structural edits while strictly preserving non-target content (Long et al., 11 Aug 2025).
1. Motivation and problem setting
TDM is introduced in the context of rectified flow / flow-matching image editing, where generation and inversion are treated as trajectories in latent space driven by a learned velocity field . The source reconstruction trajectory is framed as a stable denoising path conditioned on the source prompt , whereas the edited trajectory is driven by the target prompt . Because the target prompt alters the velocity field, the edited denoising path deviates from the reconstruction path. The paper states that this deviation is spatially informative: regions intended for modification exhibit significant divergence, while background and unaffected regions follow nearly identical trajectories (Long et al., 11 Aug 2025).
This localization signal is presented as especially relevant for large-scale shape transformation rather than small attribute transfer. The target use case is not merely texture modification but edits that change geometry, contour, and semantic identity while keeping scene layout and non-target content intact. The paper contrasts TDM with three prior classes of control strategy: binary segmentation masks, which are too rigid for shape change; cross-attention masks, which are said to be noisy and inconsistent under large structural changes; and unconditional or global feature injection, which preserves source structure but also suppresses intended edits. TDM is proposed as a content-aware, model-internal, spatially localized alternative derived from the generative process itself.
2. Formal definition
The paper defines the source inversion latent sequence as and the editing latent sequence as . For spatial token at timestep , the raw trajectory divergence is
This is a token-wise distance between source-conditioned and target-conditioned velocity predictions, evaluated at their respective trajectory latents rather than at a shared latent point. The paper is explicit that TDM is therefore a latent-space or token-space dynamical measure, not an image-space quantity, not a cross-attention map, and not a gradient-derived saliency map (Long et al., 11 Aug 2025).
Each timestep map is then min-max normalized across spatial tokens:
To obtain a temporally consistent edit map, the normalized scores are fused over an editing window 0 using token-wise softmax weighting:
1
The aggregated map is smoothed with Gaussian convolution,
2
and the paper then states that the smoothed attention map is binarized with threshold 3, yielding the final mask through 4. In the notation used by the paper, 5, 6, and 7 are continuous quantities, whereas the thresholded 8 is binary. This dual use of soft and binary forms is a core feature of the pipeline rather than an inconsistency.
3. Computation pipeline and scheduling
The practical computation begins by inverting the source image under the source prompt to obtain the initial noisy latent code 9. The inversion stage yields both the latent trajectory 0 and stored source attention features 1. Editing then proceeds along a denoising trajectory 2 conditioned on 3. At selected timesteps, TDM compares the source-prompt velocity on the inversion path with the target-prompt velocity on the editing path, normalizes the resulting token divergences, aggregates them over time, smooths them spatially, and thresholds them into the edit mask 4 (Long et al., 11 Aug 2025).
A significant implementation detail is that the comparison is not performed by evaluating both prompts at the same latent. Instead, the two velocity fields are evaluated at 5 and 6, respectively. The paper presents this as important because the divergence signal is meant to reflect the actual separation between reconstruction and editing trajectories rather than a prompt-conditioned perturbation of a single state.
The scheduling of TDM is also explicit. The method states that per-step TDM is unstable in early, high-noise denoising stages, so TDM-guided injection is not activated immediately. This is why the overall process includes early stabilization before trajectory-guided region control becomes active. The paper further states that TDM requires no external segmentation masks, no manual annotation, no extra supervision, and no training of a localization model.
Implementation details explicitly mentioned in experiments include the base model FLUX.1-[dev], the PyTorch framework, NVIDIA A100 GPU with 40GB memory, timestep 28, guidance scale 2.0, 7, ControlNet conditioning interval 8, and Depth and Canny strengths of 9 and 0. The paper does not provide a separate runtime or FLOP analysis specifically for TDM.
4. Function within Follow-Your-Shape
TDM’s immediate role is to determine where editing should be allowed. High-divergence regions are treated as editable; low-divergence regions are treated as preservation zones. This role is realized through a three-stage editing process centered on Scheduled KV Injection (Long et al., 11 Aug 2025).
In Stage 1, described as initial trajectory stabilization, the method performs unconditional or global source KV injection for the first 1 timesteps. The paper states that this operation enforces a global reconstruction objective and is equivalent to setting 2, so all regions are temporarily preserved from the inversion path.
In Stage 2, once the latent becomes more structured, TDM is computed and used to build the edit mask 3. That mask controls blending between source inversion KV features and current target KV features. For values, the paper gives
4
and states that an analogous operation is performed for keys:
5
The interpretation is direct. Where 6, target features are used and editing is allowed. Where 7, source inversion features are injected and original content is preserved.
In Stage 3, the method combines ControlNet for explicit geometric guidance with TDM-guided feature injection for selective semantic preservation. The ControlNet residual is written as
8
and the final attention output is
9
TDM therefore determines the spatial mixing pattern for 0 and 1 inside self-attention. The stated consequence is selective preservation: background and non-target regions receive source inversion features, while target edit regions receive target features, avoiding the edit suppression associated with unconditional global injection.
5. Interpretation, distinctions, and limitations
The paper’s basic intuition is that if a region does not need to change, source reconstruction and target-prompt denoising should induce similar local motion in latent space. If a region must change substantially in contour, part arrangement, or object identity, then the model must push that token along a different trajectory. TDM is thus a map of where the model’s denoising dynamics disagree most strongly between reconstruction and edit (Long et al., 11 Aug 2025).
This interpretation distinguishes TDM from several nearby localization signals. A saliency map generally highlights visually important regions, whereas TDM is edit-conditioned and trajectory-conditioned. Cross-attention maps indicate where text tokens attend spatially, but the paper argues that they are often noisy or semantically diffuse under large structural edits. Simple feature-difference maps are representational; TDM uses predicted denoising velocities along trajectories and is therefore dynamical. Gradient-based localization depends on backpropagated sensitivities to an objective; TDM is described as a forward-process comparison between source and target trajectory dynamics.
Several misconceptions are addressed by the method itself. TDM is not an externally provided segmentation mask. It is not computed in image space. It is not identical to a cross-attention mask. It is also not purely binary: the pipeline uses continuous raw, normalized, and temporally fused divergence values before producing a binary mask for KV injection.
The authors also note limitations relevant to TDM. The method is prompt-sensitive, because prompt-guided inversion and denoising trajectories can change with wording. It is hyperparameter-sensitive, including injection timing and conditioning strength. TDM is unstable in early high-noise stages, motivating the scheduled design. In the video-extension discussion, the authors state that TDM becomes less stable across frames, with editing regions fluctuating temporally. The paper therefore does not present TDM as a universally robust localization mechanism.
6. Experimental evidence and benchmark context
The empirical evaluation centers on ReShapeBench, a benchmark introduced for shape-aware editing and described as comprising 120 new images and enriched prompt pairs specifically curated for shape-aware editing (Long et al., 11 Aug 2025). The benchmark is presented as targeting the regime where TDM is meant to be most useful: large structural transition, cross-contour changes, cross-semantic changes, and subject continuity.
Qualitative evidence in the paper attributes stronger large-scale shape transformation and better preservation of non-target regions to the proposed region-controlled editing strategy centered on TDM and scheduled injection. The figures described in the text emphasize fewer background distortions, fewer ghosting artifacts, and better completion of difficult shape transformations. Figure 1 is said to visualize TDM evolving across timesteps as a dynamically localized editing region.
The best reported numbers on ReShapeBench are as follows:
The paper explicitly connects these gains to the proposed region-controlled editing strategy and to the edit mask 3, especially for background preservation, text alignment, and overall image quality.
The ablation evidence is narrower than the method’s centrality might suggest. The paper does not provide a clean ablation isolating “with TDM versus without TDM,” nor a direct replacement of TDM with cross-attention masks or segmentation masks. What it does provide is an ablation over 4, the number of initial stabilization steps before TDM-guided editing becomes active: 5 gives Aesthetic 6.51, PSNR 32.79, LPIPS 6, CLIP Sim 31.05; 7 gives 6.55, 34.38, 9.88, 32.56; 8 gives 6.57, 35.79, 8.23, 33.71; 9 gives 6.52, 31.25, 10.52, 29.41; and 0 gives 6.48, 30.41, 12.37, 27.66. The paper interprets this as evidence that too little stabilization hurts trajectory quality, too much stabilization suppresses editability, and a moderate value of 1 is best.
7. Terminological distinctions and related uses
The label “TDM” is not unique across arXiv literature. In "Learning Few-Step Diffusion Models by Trajectory Distribution Matching" and "TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward," TDM denotes Trajectory Distribution Matching rather than Trajectory Divergence Map (Luo et al., 9 Mar 2025, Luo et al., 8 Mar 2026). Those works concern few-step diffusion distillation and reinforcement learning for deterministic few-step diffusion trajectories, not spatial edit localization in flow-based image editing.
A second nearby but distinct usage appears in "Diverse Trajectory Forecasting with Determinantal Point Processes," which does not define a Trajectory Divergence Map explicitly. The closest reusable object there is the DPP kernel matrix 2, built from pairwise trajectory similarities and per-trajectory quality weights; this is a set-level diversity structure over candidate trajectories rather than a token-space edit-localization map (Yuan et al., 2019).
A more conceptually analogous but mathematically different construction appears in "Trajectory-free approximation of phase space structures using the trajectory divergence rate," where the relevant object is a scalar field over phase space, 3, describing instantaneous normal attraction or repulsion of adjacent trajectories (Jr. et al., 2017). That dynamical-systems quantity can be visualized as a trajectory divergence map in phase space, but it is defined over vector fields and invariant manifolds rather than over latent tokens in a flow-based image editor.
Within the image-editing context of Follow-Your-Shape, Trajectory Divergence Map therefore has a specific meaning: a training-free, mask-free latent/token-space map derived from inversion-versus-editing velocity disagreement and used to spatially gate Scheduled KV Injection for shape-aware editing. A plausible implication is that the term “trajectory divergence” has become useful across several subfields for naming structure extracted from differences between trajectories, but the mathematical object, domain of definition, and downstream use vary substantially across those literatures.