PaperlayerView exact revision on arXiv

Source-linked AI summary

ReState: Noise-Level-Aware State Correction for Training-Free Object Removal in Few-Step Diffusion Models

Donghwan Lee, Sanghyun Jo, Sihyun Kim, Ziseok Lee, Kyungsu Kim

arXiv:2610.21287v1cs.CV

TL;DR

Training-free object removal often leaves object-shaped residuals in few-step distilled diffusion, where a structured masked prediction forms early and attention rerouting acts later. ReState corrects the masked clean prediction toward retrieved background evidence with a noise-level-aware weight while retaining object-excluding attention. It achieves the best masked metrics among training-free methods across the evaluated paired benchmarks and few-step backbones, without additional denoiser evaluations.

  • Problem

    On few-step distilled diffusion models, training-free attention forcing can improve removal yet leave object contours and textures in the mask, raising how to change content already being committed by a short trajectory.

    Source: §1 · p. 1

  • Method

    ReState retrieves background evidence from unmasked source regions and corrects the masked clean prediction toward it with a noise-level-aware weight, while retaining object-excluding attention and the native denoiser-evaluation count.

    Source: §1 · pp. 1–2

  • Results

    ReState achieves the best masked PSNR, SSIM, and LPIPS among training-free methods across two paired benchmarks and two main few-step backbones, improving masked PSNR by 3.17 dB over the strongest baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    In the evaluated settings, few-step removal requires direct state correction while generative freedom remains, rather than attention rerouting alone.

    Source: §6 · p. 10

    The retrieved background evidence is a proposal built from clean source latents outside the mask, not a ground-truth reconstruction.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

Training-free object removal commonly redirects self-attention away from the target region, but on few-step distilled diffusion models this often leaves object-shaped residuals. We identify a temporal mismatch: a structured prediction forms inside the mask at the first forward, and the short trajectory offers few later chances to revise it, while attention rerouting acts mainly in later steps. Across four distilled backbones, magnitude-matched single-forward interventions show that background-directed state correction affects the final result more than attention rerouting. We therefore propose ReState, which corrects the masked clean prediction toward background evidence retrieved from the input using a noise-level-aware weight that vanishes at the final forward. ReState retains object-excluding attention as a complementary constraint and adds no denoiser evaluations. Across two paired benchmarks and two few-step backbones, ReState achieves the best masked PSNR, SSIM, and LPIPS among training-free methods, improving PSNR by 3.17 dB on average over the strongest baseline. On OBER-Test, it also outperforms the 50-step attention-forcing reference on all three metrics at 6.7x lower latency. The same configuration transfers to two additional distilled backbones without retuning. Project page: https://tdrlee.github.io/ReState

1 Introduction

Few-step distilled models can commit structured content inside the mask before attention forcing can remove it. ReState directly corrects the masked prediction toward retrieved background evidence while retaining object-excluding attention, and improves benchmark results without extra denoiser evaluations.

Diagnosis

  • Across four distilled backbones, matched single-forward interventions show that direct background-directed correction has a larger terminal removal effect than attention forcing.On the Turbo backbones, the first-forward effect of attention forcing is not significantly different from zero, while matched direct correction is significant.Source: §1 · pp. 1–2

Method

  • ReState moves the masked clean prediction toward background evidence retrieved from unmasked source regions, with correction strength keyed to noise level and vanishing at the final forward.Object-excluding attention remains active as a complementary constraint, and retrieval reuses the first generation forward without adding denoiser evaluations.Source: §1 · pp. 1–2

Motivation

  • Few-step models can fix a structured image inside the mask at the first forward, leaving the short trajectory unable to absorb it and attention forcing unable to remove the resulting residual.The figure contrasts this with multi-step SDXL, where the background forms gradually along the trajectory.Source: §1 · p. 1, Figure 1 · p. 2

Results

  • ReState achieves the best masked PSNR, SSIM, and LPIPS among training-free methods across two paired benchmarks and two main few-step backbones, improving masked PSNR by 3.17 dB over the strongest baseline.On OBER-Test, it surpasses every compared 50-step attention-forcing configuration in SSIM and LPIPS, runs 6.7× faster than the 50-step reference, and transfers without per-backbone tuning to SDXL-Lightning and Hyper-SDXL.Source: §1 · p. 2

2 Related Work

Related work spans removal-specific training or optimization, training-free attention routing, and direct manipulation of diffusion states. ReState’s distinction is its early correction toward retrieved background evidence for the short trajectories of few-step models.

Training-based and optimization-based removal

  • Training-based and optimization-based removal methods use removal-specific data, adaptation, or dedicated training, whereas ReState operates on a frozen pretrained model without removal-specific training.A separate line trains fast removal models that generate results in one to four forwards.Source: §2 · p. 3

Training-free attention-based removal

  • Training-free attention methods suppress target-region keys to draw information from surrounding context, but none of the described methods supplies an explicit background target for the masked prediction.The surveyed approaches vary in suppression rules, schedules, and guidance mechanisms.Source: §2 · p. 3
  • ReState uses hard object-excluding attention as a complementary constraint, not as its contribution, and directly corrects the masked clean prediction when the trajectory is most sensitive.The paper distinguishes this from relying on attention rerouting alone.Source: §2 · p. 3

Direct state manipulation in diffusion sampling

  • ReState resembles prior diffusion state manipulation but uses retrieved pseudo-target background evidence, computes no measurement-model gradient, and runs on a distilled backbone’s native four-step schedule.The paper identifies direct state manipulation itself as not novel; its contribution is measuring the effective early intervention and its noise-dependent schedule.Source: §2 · p. 3

Few-step distilled diffusion

  • Prior work on few-step distilled diffusion has studied generation quality, speed, editing, and inpainting; this work tests whether a single-forward removal intervention survives to the terminal image.That terminal-effect analysis motivates combining attention rerouting with direct prediction correction.Source: §2 · pp. 3–4

3 Method

ReState combines object-excluding attention with a masked-state correction toward background evidence retrieved from the source. The correction is strongest while noise remains and vanishes at the final forward, without adding denoiser evaluations.

3.2 Existing Object-Excluding Attention

  • Attention forcing suppresses attention to target-region keys, but does not specify which background content the masked prediction should approach.ReState uses the hard form at every forward.Source: §3.2 · p. 5

3.3 Mask-Excluded Background Retrieval

  • ReState corrects only the masked clean prediction toward retrieved background evidence, while retaining object-excluding attention as a route constraint.The retrieval uses first-forward features to aggregate clean source latents from unmasked positions; the same forward supplies the denoiser output, so retrieval adds no denoiser evaluation.Source: §3.3 · p. 5, §3.2 · p. 5, §3.4 · p. 6
  • The retrieved proposal aggregates clean source latents at unmasked positions whose first-forward features match each masked position among the top fraction r.The proposal is built from unmasked source content, rather than a ground-truth reconstruction.Source: §3.3 · p. 5

3.4 Noise-Level-Aware Direct Correction

  • The correction minimizes a trade-off between staying near the model prediction and moving masked positions toward the retrieved evidence, weighted by noise remaining after the upcoming transition.The quadratic objective separates over latent positions and has a closed-form minimizer, requiring no gradient or inner iterations.Source: §3.4 · p. 5
  • Correction leaves the mask exterior unchanged, contracts masked distance to the evidence by 1 − ωt, and vanishes at the final transition.Using σt−1 encodes the terminal boundary: when σt−1 = 0, ωt = 0, leaving final synthesis to proceed under object-excluding attention.Source: §3.1 · p. 4, §3.4 · pp. 5–6

4 Experiments

Across the reported few-step comparisons, ReState leads training-free methods on masked metrics, while matched interventions and component ablations show that direct correction acts early and attention forcing adds most near final synthesis. Its unchanged configuration also transfers to additional backbones, and it beats the 50-step reference on OBER-Test at lower latency.

Main removal results

  • 3.17 dB: ReState improves masked PSNR over DesignEdit on average and leads masked PSNR, SSIM, and LPIPS in all four backbone–benchmark settings.Table 2 compares methods on OBER-Test and RORD under the same frozen backbone and nominal four-step schedule.Source: §4 · p. 7

Transfer to additional backbones

  • The same ReState configuration improves masked PSNR by 1.46 to 2.39 dB over DesignEdit on SDXL-Lightning and Hyper-SDXL without per-backbone tuning.The reported gains extend to every masked metric on both benchmarks.Source: §4 · p. 7

Intervention asymmetry and timing

  • At the first forward, magnitude-matched direct correction has a larger terminal removal effect than attention forcing on all four tested distilled backbones.On the Turbo family, the first-forward attention-forcing effect is not significantly different from zero; attention forcing grows more effective as noise decreases.Source: §4 · pp. 7–8

Component and timing ablation

  • 17.62 dB: direct correction alone exceeds the 15.50 dB from full-schedule attention forcing alone on SDXL-Turbo.Adding attention forcing only at the final forward raises the result to 18.65 dB, and applying it throughout reaches 18.75 dB; last-forward attention forcing without correction reaches 12.15 dB.Source: §4 · pp. 8–9, Table 4 · p. 9

Efficiency and multi-step comparison

  • On OBER-Test, ReState improves all three masked metrics over the 50-step attention-forcing reference at 6.7× lower latency.ReState retains the native four denoiser evaluations; its additional retrieval and masked update add 51% latency over the no-intervention SDXL-Turbo baseline.Source: §4 · p. 9, Table 5 · p. 9

Robustness

  • Across seeds, ReState ranks first on all three masked metrics for every reported backbone and benchmark, with lower masked-PSNR variation than DesignEdit.The comparison uses seeds 0, 1, and 2.Source: §4 · p. 9

5 Discussion

The discussion attributes ReState’s component roles to a timing mismatch: direct correction supplies background direction early, while attention forcing contributes mainly during late synthesis. It also bounds the findings to the tested settings and identifies background-evidence quality and textural detail as remaining constraints.

5 Discussion

  • A four-step trajectory preserves first-forward perturbations more strongly than a 50-step reference, while attention forcing contributes mainly at the final forward.The discussion reports a first-forward survival slope of 1.38 on the four-step trajectory versus 0.30 on the 50-step reference, and 1.03 of attention forcing’s 1.13 dB added benefit at the final forward.Source: §5 · pp. 9–10

Scope and limitations

  • ReState transferred unchanged to 50 steps yields 17.57 dB, while a diagnostic noise window recovers 18.83 dB.The discussion attributes this difference to correction remaining strong across many low-noise forwards in the 50-step schedule.Source: §5 · p. 10
  • The removal findings cover two paired benchmarks and two main backbones, with further full-removal evaluation on SDXL-Lightning and Hyper-SDXL.An oracle background raises masked PSNR from 18.75 to 22.32 dB, and the fill has 0.51× the high-band energy of the true background.Source: §5 · p. 10

6 Conclusion

ReState addresses residual objects in few-step distilled diffusion by correcting the masked prediction toward retrieved background evidence while retaining object-excluding attention. Across the evaluated settings, it achieves the best masked metrics among training-free methods without adding denoiser evaluations.

6 Conclusion

  • Across two paired benchmarks and two few-step backbones, ReState achieves the best masked metrics among training-free methods without additional denoiser evaluations.The conclusion attributes the remaining few-step removal challenge to the need for direct state correction while generative freedom remains.Source: §6 · p. 10
  • Attention-based removal on few-step distilled diffusion often leaves object-shaped residuals.The conclusion states that, in the evaluated settings, attention rerouting alone is insufficient while generative freedom remains.Source: §6 · p. 10
  • ReState combines noise-level-aware correction toward retrieved background evidence with object-excluding attention as a complementary constraint.The correction weight vanishes at the final forward.Source: §6 · p. 10

A Method Details

The method operates on clean predictions across diffusion parameterizations, constrains self-attention keys in object regions, and retrieves background evidence from the input. Its closed-form correction leaves unmasked positions unchanged and fades to zero at the final transition.

A.1 Exact Diffusion Parameterization

  • The correction is applied in clean-prediction space and mapped back to the model’s native output type, leaving the sampler transition unmodified.The method is stated to be agnostic to output parameterization and to deterministic or stochastic samplers.Source: §A.1 · p. 15

A.2 Object-Excluding Attention Details

  • Object-excluding attention biases object-region key columns for all query rows and heads, while leaving cross-attention layers unchanged.The intervention covers self-attention layers on the image-token grid; masks are max-pooled to each token grid.Source: §A.2 · p. 15

A.3 Background Retrieval Details

  • Background retrieval matches masked-position features against unmasked features and forms a weighted average of clean source latents at retrieved positions.It retains the top 1% of background positions, with softmax weights at temperature τ = 0.1; features come from the first generation forward, adding no auxiliary pass.Source: §A.3 · p. 16

A.4 Derivation of the Closed-Form Correction

  • The closed-form correction changes only masked clean predictions, with its weight increasing with remaining noise and reaching zero at σt−1 = 0.For SDXL-Turbo with λ = 8, the four-forward weights are 0.993, 0.954, 0.794, and 0; keying the weight to σt instead leaves final-forward correction and collapses perceptual quality.Source: §A.4 · pp. 16–17

B Experimental Details

The experiments use paired removal benchmarks, masked image-quality metrics, and image-level paired statistical tests. Baseline comparisons share a forward-only pipeline, while the method reuses the first forward for retrieval and stays at the native denoiser count.

B Experimental Details

  • OBER-Test uses all 163 released paired images, while RORD contributes one manually reviewed single-target frame from each of 343 sequences, paired with its object-free ground truth.RORD samples and masks are converted to 512 × 512 mask-centered crops, preserving native resolution for small objects and minimally downsampling larger ones.Source: B Experimental Details · p. 17

B.1 Metrics and Statistics

  • Masked PSNR and SSIM are computed over masked regions, while masked LPIPS compares a composite that matches ground truth outside the mask.PSNR and SSIM exclude a 4-pixel border; LPIPS uses a padded mask bounding box, and results are averaged per image.Source: §B.1 · p. 18

B.2 Auxiliary-Pass Accounting of Baselines

  • Metric differences are paired by image and tested with sign-flip permutation tests; ReState’s retrieval probe reuses the first forward and adds no denoiser pass.ReState stays at the native forward count, unlike baselines that duplicate batches or make reference calls.Source: §B.1 · p. 18, §B.2 · p. 19, Table B2 · p. 19

C.1 Complete First-Forward Analysis

Matched single-forward tests find that direct background correction has a stronger first-forward terminal effect than attention forcing on distilled backbones. The controls and retention analysis show that this asymmetry depends on the coarse distilled setting and that early perturbations persist more on the few-step trajectory.

C.1 Complete First-Forward Analysis

  • With perturbation norm fixed, only directions toward retrieved or oracle background reduce the source-to-ground-truth residual; random and orthogonal directions do not.Across first-forward runs, the direct-versus-attention effect correlates with masked PSNR and LPIPS changes, and most images with a positive effect improve on both metrics.Source: §C.1 · pp. 19–20
  • At the first forward, direct correction and its gap over attention forcing are significant on the Turbo family, while attention forcing is not significantly different from zero.Across all four distilled backbones, the gap remains significant after Holm correction; 71–82% of images lie above the diagonal, compared with 49% for the original-weight control.Source: Table C1 · p. 20, §C.1 · p. 20, Figure C1 · p. 20

C.2 Multi-Step Reference Intervention

  • On the 50-step SDXL reference trajectory, both single-forward interventions have effects an order of magnitude smaller, with no comparable first-forward asymmetry.A single early intervention of either kind is largely absorbed by the remaining trajectory.Source: §C.2 · p. 20, Table C2 · p. 21, Figure C2 · p. 21

C.3 Diagnostic Schedule Controls

  • The intervention asymmetry appears only when distilled weights meet the coarse schedule, but the two control configurations are diagnostic rather than a causal factorization.Original SDXL weights on the four-step grid show moderate, symmetric effects.Source: §C.3 · pp. 20–21
  • Densifying SDXL-Turbo to 8 or 16 evaluations shrinks the first direct correction’s terminal effect without unlocking an early attention-forcing effect.At a fixed budget on the 50-step reference, late eight-forward attention forcing beats the high-noise window by 0.9 dB; only the full schedule reaches the reference quality.Source: §C.3 · pp. 21–22, Table C3 · p. 21, Table C4 · p. 22

C.4 Retention Analysis

  • First-forward perturbation retention has a median of 1.38 on four-step SDXL-Turbo and 0.30 on the 50-step reference.The few-step trajectory preserves, and can amplify, a first-forward perturbation, whereas the dense trajectory attenuates it.Source: §C.4 · p. 22

C.5 Residual-Frequency Analysis

On OBER-Test, direct correction removes object-shaped high-frequency structure but produces a smoother-than-real fill; across a 32× strength sweep, removal quality remains nearly flat within the tested correction structure.

C.5 Residual-Frequency Analysis

  • On OBER-Test, masked high-band energy falls from 8.3× the true-background level without intervention to 1.92× with attention forcing, 0.30× with direct correction, and 0.51× with ReState.Attention forcing leaves object-shaped excess; ReState eliminates it, while the remaining gap to parity reflects a smoother-than-real fill.Source: Table C5 · p. 22, §C.5 · p. 22

D Method Ablations

  • The ablations vary one component of the default SDXL-Turbo configuration on OBER-Test, with reported significance tested by two-sided paired permutation against that configuration.Tables report point estimates.Source: D Method Ablations · pp. 22–23

D.1 Correction Strength

  • Across a 32× λ sweep, masked PSNR varies by at most 0.14 dB, and the default λ = 8 is within 0.01 dB of the optimum.Realized weights shift substantially, so this insensitivity holds within the closed-form profile, which is strong while noise remains and zero at the final forward.Source: §D.1 · p. 23

D.2 Correction Timing and Schedule Structure

Correction windows ending later in the four-forward schedule achieve higher removal quality, while the terminal forward should remain uncorrected. The ablations also show that evidence quality, per-forward application, and similarity weighting matter to the tested outcomes.

D.2 Correction Timing and Schedule Structure

  • Keying ω to σt leaves a final-forward correction of about 0.79 and degrades LPIPS from 0.194 to 0.284; restoring the terminal skip is not significantly different from default.A terminal weight sweep likewise degrades LPIPS as ω4 rises from 0 to 0.794.Source: §D.2 · pp. 23–24, Table D2 · p. 24
  • With ω4 = 0 retained, the closed-form weight profile is not where performance comes from; adequate overall strength and terminal release are the operative factors in the four-step schedule.At T=1, every forward is terminal and the correction vanishes identically.Source: §D.2 · p. 24
  • On OBER-Test, correction windows ending at the third forward reach 18.47–18.75 dB, versus 18.18–18.35 dB for windows ending at the second and 16.85 dB for first-forward-only correction.Object-excluding attention remains active in every row; uncorrected forwards can re-predict masked content.Source: §D.2 · p. 23, Table D2 · p. 24

D.3 Background Evidence Source

  • Replacing retrieved background evidence with a global mean costs 0.59 dB, while using the ground-truth background as an oracle target gains 3.57 dB.The ablation attributes the difference to evidence quality rather than the correction mechanism.Source: §D.3 · p. 24

D.4 Per-Forward Correction versus One-Shot Latent Initialization

  • At the same target and four-denoiser-evaluation budget, per-forward prediction correction beats writing the retrieved proposal into the latent once on masked LPIPS on both benchmarks.The comparison also reports masked PSNR and SSIM win rates of 69–90%.Source: Table D3 · p. 25, §D.4 · p. 25

D.5 Retrieval Ratio

  • Across retained fractions r ∈ [0.005, 1], masked PSNR stays within 0.15 dB of default, while LPIPS is never significantly better at larger fractions.The authors retain r = 0.01 because the shallow PSNR peak is not statistically distinguishable from default.Source: §D.5 · p. 25

D.6 Probe Feature Source

ReState transfers without tuning to two additional distilled backbones and compares favorably with native 50-step attention-forcing configurations on OBER-Test. Probe-source tests, baseline retuning, and schedule-transfer experiments further specify where the method works and what its correction schedule requires.

D.6 Probe Feature Source

  • The shared probe is unchanged by attention suppression (18.72 versus 18.75 dB; p ≥0.46 on every metric), and retaining suppression requires no separate pass.The default harvest uses the same pass as the first generation forward.Source: §D.6 · pp. 25–26
  • Near-clean probe features degrade to 17.44 dB and LPIPS 0.214, while second-forward features remain at par; feature quality also tracks denoiser layer rather than native resolution.Mid-block features are not significantly different from the default, whereas encoder-side and decoder features perform worse; the authors read the noise pattern as consistent with content suppression.Source: §D.6 · p. 26
  • Recomputing correspondence at each generation forward slightly raises quality (19.02 versus 18.75 dB; LPIPS 0.190 versus 0.194), although retrieved positions change substantially.The bank values remain fixed to clean source latents, so correspondence is the quantity that changes.Source: §D.6 · p. 26

E.1 Full Quantitative Tables

  • Table E1 collects the full quantitative comparison; OSOR is excluded from the matched comparison because its OBER-Test evaluation is not guaranteed to be out-of-distribution.OSOR’s SDXL-Inpainting variant reaches 23.19 dB masked PSNR on OBER-Test and 18.77 dB on RORD under the stated protocol.Source: §E.1 · p. 27

E.2 Full Removal on Additional Distilled Backbones

  • On both additional backbones and benchmarks, ReState improves every masked metric over DesignEdit by 1.46–2.39 dB masked PSNR with the same frozen configuration and no per-backbone tuning.Differences are significant at p = 1.0 × 10−4; margins are smaller than on SDXL-Turbo, where attention forcing has a positive first-forward effect.Source: §E.2 · p. 27, Table E2 · p. 28

E.3 Multi-Step Baseline Comparison

  • Transferring ReState unchanged to 50 steps underperforms at 17.57 dB and LPIPS 0.231; a correction window recovers performance to the level of the strongest 50-step attention configuration.The authors identify the native few-step schedule as the method’s operating point because dense schedules need an explicit window to preserve uncorrected synthesis freedom.Source: §E.3 · pp. 28–29
  • At four steps, ReState (18.75 dB, SSIM 0.498, LPIPS 0.194) exceeds every 50-step attention configuration on SSIM and LPIPS, within 0.04 dB of the best attention PSNR, at 6.7× lower latency.Attention-forcing baselines that degrade on the four-step port recover at their native 50-step count.Source: §E.3 · p. 28, Table E3 · p. 29

E.4 Limited Re-tuning of Baselines at Four Steps

  • Cross-benchmark retuning leaves the selected four-step baselines 3.03–3.41 dB below ReState in masked PSNR, with every masked-metric gap significant.Attentive Eraser recovers from 7.09 to 15.34 dB on OBER-Test, while DesignEdit’s official full-schedule masking remains its best configuration.Source: §E.4 · p. 29, Table E4 · p. 30

E.5 Seed Robustness

Across three seeds, ReState ranks first on all masked metrics for every seed, backbone, and benchmark, with lower masked-PSNR variability than DesignEdit. The section also includes qualitative results on OBER-Test and RORD.

E.5 Seed Robustness

  • ReState ranks first on all three masked metrics in all 36 seed-backbone-benchmark cells, and its paired advantage over DesignEdit is significant in all 36 tests.Its masked-PSNR standard deviation is 0.02–0.21 dB, versus 0.31–0.81 dB for DesignEdit.Source: Table E5 · p. 30, §E.5 · p. 30

E.6 Qualitative Results on OBER

E.7 Qualitative Results on RORD

Related papers

Loading 2610.21287v1…