TRACE: TRajectory-Aware Concept Erasure for Flow Matching Models

Concept unlearning is essential for suppressing harmful and copyrighted content in text-to-image (T2I) models while preserving their generation utility. While flow matching (FM) models are increasingly widespread, current unlearning techniques are predominantly diffusion-oriented, redirecting model predictions toward safe anchor-conditioned targets or unconditional generations at isolated states. However, directly transferring these diffusion-unlearning objectives to FM models can overlook their inherent continuous flow structure. We theoretically analyze how such adaptation can damage the underlying flow geometry. Flow-aware unlearning is essential for FM models because mismatching the flow paths can introduce conflicting velocity supervision and impose a positive regression floor, which causes a risk of utility degradation during unlearning. To overcome these structural disruptions, we propose TRACE, a flow-aware unlearning framework that erases target concepts via safe flow trajectory replacements. TRACE first jointly selects anchors for different generation seeds using a flow conflict proxy. It then fine-tunes the model along the selected teacher trajectories to replace harmful or undesirable generations while preserving a safe flow structure. Across the eight main-table settings, experiments on Stable Diffusion 3.5 Medium and FLUX.2 4B demonstrate that TRACE outperforms state-of-the-art baselines, reducing target-concept CLIP similarity by up to 7.10%, increasing retained-prompt CLIP alignment by up to 7.78%, and reducing FID by up to 16.55%. Furthermore, under sequential unlearning of 15 concepts on SD 3.5 Medium, TRACE maintains a 74.93% anchor-win rate and 15.29 FID, retaining a clear advantage over the compared baselines.