Data shape determines how features are structured, how patterns are separated, and how distributions cover the underlying domain. Poor data shape can make models learn noise rather than generalizable structure. This paper studies robust feature-centric data reshaping: generating feature transformations that remain useful, stable, and reproducible under noisy evaluation and imperfect data conditions. We view reshaping operation sequence search as reward-guided diffusion generation, and robust reshaping as searching for regions in the latent reward landscape rather than isolated high-reward transformations. The key challenge is dual instability: noisy evaluators distort local reward guidance, while stochastic generative trajectories can converge to inconsistent solutions. We propose FCDiff, a flat-consensus diffusion framework that addresses both failures through a micro-macro decomposition. The micro layer replaces point-estimate reward guidance with Gaussian-smoothed, Monte Carlo averaged gradients, steering generation toward locally flat reward regions. The macro layer aggregates independently guided trajectories with a weighted Frechet-mean barycenter, selecting consensus-supported basins and filtering stochastic outliers. Across an 8-dataset headline cohort under heavy-tailed evaluator noise, FCDiff attains the best aggregate rank on lower-tail reliability and robustness against both search-based AutoFE and robustness-oriented generative baselines, with statistically significant accuracy gains over every generative baseline. Our results show that robust data reshaping requires searching for flat, consensus-supported regions rather than sharp single-trajectory optima.
Figures & tables
Fig. 1 : Robust data reshaping at a glance: two coupled instabilities. (a) Evaluator noise. Search runs over a latent reward landscape scored by a noisy evaluator (gray band); under a small latent perturbation ±δ , a sharp peak loses its reward while a flat region keeps it. (b) Trajectory stochasticity. Independent guided trajectories from the same noise zT end at different latents; we aggregate them into a weighted consensus (star) that down-weights stochastic outliers. FCDiff targets regions that are both flat and consensus-supported.
Fig. 2 : FCDiff as a symmetric micro–macro framework for robust data reshaping. A VAE–reward–diffusion backbone builds a continuous latent reward geometry over discrete reshaping sequences. The micro layer handles local reward instability by smoothing point-estimate guidance into flat-region guidance. The macro layer handles trajectory-level instability by aggregating multiple guided endpoints into a consensus latent. Together, FCDiff searches for flat, consensus-supported regions rather than sharp single-trajectory optima.
Micro (Local)
Macro (Global)
Failure Mode
Noisy reward gradients
Stochastic trajectory divergence
Operator
Gaussian smoothing
Fréchet mean aggregation
Objective Relaxation
R(z)→Eδ[R(z+δ)]
max→ consensus center
Geometry
Flat regions
Broad basins
Effect
Reduces local sensitivity
Reduces cross-run variance
Output
Stable trajectory endpoint
Consensus latent solution
TABLE I : Structural symmetry of the micro and macro components in FCDiff. The two layers address complementary sources of instability and jointly enable region-level robust reshaping.
ml_586
ml_589
ml_607
ml_620
Group
Method
Perf.
CVaR
Worst
Perf.
CVaR
Worst
Perf.
CVaR
Worst
Perf.
CVaR
Worst
Search
RL
0.682
0.662
0.662
0.678
0.662
0.662
0.666
0.651
0.651
0.660
0.647
0.647
Evo
0.698
0.658
0.658
0.691
0.652
0.652
0.688
0.661
0.661
0.657
0.638
0.638
LLM
0.688
0.654
0.635
0.685
0.664
0.639
0.666
0.659
0.658
0.656
0.640
0.638
Robust-gen
Std
0.687
0.684
0.683
0.674
0.671
0.671
0.663
0.659
0.658
0.652
0.652
0.652
Adv
0.687
0.683
0.683
0.674
0.671
0.671
0.663
0.659
0.659
0.652
0.650
0.650
TABLE II : Main results (8-dataset headline cohort, 3 metrics: Perf., CVaR@5, Worst). All methods optimise under Student- t(ν=3,σ=1) noise on every reward query; downstream performance reported on the clean 5-fold CV evaluator ( 1− RAE for the four ml_* regression datasets, macro-F1 for spectf / svmgd3 / spam / ap_omn ). Search baselines: n=5 pipelines; robust-gen + Ours : up to n=100 /dataset. Per-cell global winners bolded (ties bolded jointly). The per-cell deltas look small, but they are systematically in Ours’ favour: 13/24 per-cell wins (incl. ties), and statistically separable from every robust-gen baseline on Perf. (paired Wilcoxon, p<0.05 , n=8 ); full statistics in § IV-B .
Mean rank (1 = best)
p (vs FCDiff)
Method
Perf.
CVaR@5
Worst
Perf.
CVaR@5
Worst
FCDiff (ours)
2.25
2.00
3.12
—
—
—
Std-Diff
4.88
3.25
3.25
.008
.203
.438
Adv-PGD
5.88
4.75
4.25
.008
.016
.125
Grad-Clip
5.25
4.88
4.88
.008
.016
.031
Medoid-Cons
7.00
5.25
5.00
.008
.008
.094
TABLE III: Aggregate reliability summary across the 8-dataset headline cohort. Mean rank is averaged over the 8 datasets within each metric (1 = best, 8 = worst). p (vs FCDiff) is the two-sided paired Wilcoxon signed-rank test on the per-dataset deltas. Bold = best mean rank in column; p bolded when p<0.05 . FCDiff attains the lowest mean rank on every metric, with statistically significant gains across the four robust-gen baselines on Perf. and CVaR@5; the Worst-metric tie with Std-Diff is the regime in which both methods occupy the same single-trajectory peak, while CVaR@5 separates them.
Fig. 3 : Diagnostics for the micro and macro layers (cohort: 3-dataset SDEdit-init side cohort in (a)(b); 8-dataset headline cohort in (c)).
Fig. 4 : Geometry (RQ5) and efficiency (RQ6) diagnostics.
Worst-10% Utility ( ↑ )
Utility Var. ( ↓ , ×10−4 )
Variant
586
spectf
svm3
586
spectf
svm3
DIFFT (Std-Diff, K=M=1 )
0.729
0.500
1.103
2.6
2.7
140
+ smoothing ( K>1 , M=1 )
0.731
0.806
2.171
0.98
0.55
53
+ consensus ( K=1 , M>1 )
0.743
0.545
1.833
0.75
0.89
20
FCDiff (both layers)
0.748
0.845
2.361
0.52
0.66
15
TABLE IV : DIFFT-backbone head-to-head on the SDEdit-anchored side cohort (3 datasets where the variants admit clean discrimination). DIFFT (Std-Diff) is the vanilla DIFFT backbone with K=M=1 from Table II ; the other rows add one or both of our ensemble layers. Numbers are mean over 20 seeds × 50 samples per cell. Worst-10% Utility : surrogate-space lower tail (higher = better). Utility Var. : per-trajectory utility dispersion (lower = better), reported in units of 10−4 . The (syntactic) invalid-decode rate is 0.0 for all four variants on all three datasets of this SDEdit-anchored cohort and is therefore omitted from the table. Per-cell winners are bolded .
Method
ml_586
spectf
svmgd3
spam
Decode valid
best-of- M
.727
.771
.859
.957
100%
consensus
.745
.798
.875
.961
100%
Δ (cons. − best-of- M )
+.018
+.027
+.016
+.004
—
TABLE V : Consensus latent vs. best-of- M ( M=20 , 50 seeds, Student- t(ν=3,σ=1) ). Cells report Worst-10% Utility (higher is better). Syntactic decode validity is identical (and perfect) for both on every dataset; consensus dominates on the tail metric, with the gap shrinking toward zero on easy datasets where reward noise has minimal effect on best-of- M .
Variant
586
589
607
620
spectf
svm3
spam
ap_om
Worst-10% Utility (higher = better)
full_FCDiff
1469.1
360.5
1962.2
1002.6
1.14
494.2
1522.6
7.0
standard_M
40.0
5.2
20.4
13.1
0.93
11.0
14.0
2.4
single_FCDiff
1351.8
322.0
1852.4
930.8
1.07
433.8
1417.1
4.4
Wall-clock (s, lower = faster)
full_FCDiff
58.8
59.8
60.0
49.1
23.0
48.7
59.5
57.6
TABLE VI : Compute-matched comparison at M=16 on the 8-dataset headline cohort (this slice is from the RQ6 M -sweep {4,8,16,32} ; the headline ablation uses M=20 ). Best-of- M ( standard_M ) and single guided trajectory ( single_FCDiff ) are matched-compute baselines for the consensus output ( full_FCDiff ). Per-cell winners on the Worst-10% Utility row are bolded .
Training data attribution for diffusion models aims to identify the training samples that influence a generated instance, but existing methods either require costly per-sample gradient computation or query-specific model optimization. Moreover, most methods attribute changes in a proxy loss rather than changes in the actual model's generative behavior. We address these limitations by formulating attribution directly with a local score discrepancy measure, which applies to any diffusion variant (including DDPM, EDM, and flow matching), and by showing that such measure can be estimated without retraining, as a preconditioned gradient similarity. We instantiate this estimator as Training-data Influence via score Discrepancy (TID), which uses Kronecker-factored curvature to avoid random projections and per-sample gradient storage. We then distill TID into TIDE, a forward-only student trained online to reproduce the teacher's rankings from the diffusion model's internal activations. Under counterfactual evaluation on CIFAR-10, ArtBench-10, and MS-COCO, TID matches or outperforms state-of-the-art approaches, while TIDE retains most of TID's accuracy at four to five orders of magnitude lower per-query cost, attributing generated samples in milliseconds and faster than the generation itself.
Shixuan Liu, Joan Serrà, Kin Wai Cheuk +6
University of Illinois at Urbana-Champaign · Sony AI · University of Texas at Austin +1
Training data attribution (TDA) should enable generative model interpretability and foster a variety of related downstream tasks. Nonetheless, current TDA approaches lack reliability and robustness, preventing their adoption in real-world setups. In this paper, we take a decisive step towards more reliable and robust TDA for diffusion models. We propose to perform TDA with mirrored unlearning and noise-consistent skew (MUCS). The idea is to fine-tune a second model with bounded mirrored gradient ascent, and to measure the normalized skew of this model with respect to the original one using consistent noise samples. We show that, while being conceptually simple and generic, MUCS systematically outperforms existing methods on three different datasets by a large margin. We additionally study the effect that core design choices have on final performance, and analyze novel aspects regarding the overlap of influential instances across generated items and the potential of ensembling TDA approaches. We believe that our findings may have broader implications for more general unlearning setups, as well as for tasks requiring the comparison of diffusion losses.
Model collapse arises when generative models are trained on synthetic data produced by earlier models. The phenomenon has attracted considerable attention because of its societal and technical implications. However, previous studies have reached seemingly contradictory conclusions: replacing real data with synthetic data causes collapse (Shumailov et al.), yet accumulating real data alongside synthetic data can prevent it. For diffusion models, we study an intermediate regime typical of finite-budget pipelines: all past datasets and the real data are kept, but each new model is trained on a fixed-size sample from this growing pool, so the real fraction vanishes without any data being removed. Experiments on a 2D spiral dataset as well as the image benchmarks (MNIST, Fashion-MNIST, and CIFAR-10) show that replacement protocol degrades dataset rapidly as in the literature, whereas the fixed budget degrades only partially, sparing some features. A linear-response model of the multi-generational parameter dynamics, analyzed by stochastic recursion, confirms that the two protocols differ: some features will be fragile and lost within a few generations for both protocols, while some will be robust and preserved over practically unbounded horizons under the fixed budget protocol.
Hanna Malet, Gabriel Turinici
CEREMADE CNRS Université Paris Dauphine - PSL Paris 75116 FRANCE · Cornell University Ithaca, NY, USA