CARE: Condition-Aware Representation Regularization for Diffusion Models

Fengjia Guo; Zhuoyi Yang; Jie Tang

Source-linked AI summary

CARE: Condition-Aware Representation Regularization for Diffusion Models

Fengjia Guo, Zhuoyi Yang, Jie Tang

arXiv:2609.28561v1Sep 23, 2026cs.LGcs.CV

TL;DR

Existing representation regularizers often overlook built-in conditions such as labels and text, despite their role in defining generation targets. CARE modulates representation regularization using condition similarity. It improves generation quality across class-conditional and text-to-image tasks, and the paper limits its study largely to categorical and textual conditions and image-generation diffusion models.

Problem
Common representation regularizers overlook built-in conditions such as labels and captions, although these signals determine the generation target.
Source: §1 · p. 1
Method
CARE is a lightweight, plug-and-play regularizer that uses built-in condition signals to guide the representation space without external foundation models.
Source: §1 · p. 2
Results
CARE improves generation quality across class-conditional and text-to-image tasks, with a further 16.61% FID reduction on text-to-image benchmarks.
Source: §1 · p. 2, §4.2 · p. 5
Takeaways & Limitations
CARE complements existing approaches and requires no external models, offering a distinct source of improvement.
Source: §5 · p. 8
The study focuses largely on categorical and textual conditions and on diffusion models for image generation.
Source: §5 · p. 8

Abstract

from arXiv · show

Recent advances in diffusion models highlight the importance of representation regularization for improving sample quality and training efficiency. However, commonly used regularization methods often overlook the built-in conditions (such as labels or texts) which directly determine the generation target. In this work, we demonstrate how conditioning signals affect the feature distribution and introduce the CARE (Condition-Aware REpresentation regularization). CARE is a lightweight plug-and-play regularization framework that dynamically modulates feature distribution based on condition similarity. CARE leverages built-in conditioning signals to judiciously guide the representation space, promoting tighter feature clusters for similar conditions without relying on explicit alignment losses or external supervision. Empirically, CARE consistently improves both visual fidelity and convergence stability across both class-to-image and text-to-image tasks. On ImageNet, CARE achieves a 19.08\% reduction in FID in 400k training steps, leading to a 3.5$\times$ speed-up. When applied to text-to-image generation, CARE lowers FID by 16.61\% in 200k iterations and improves semantic alignment between generated samples and text prompts. Moreover, CARE can be seamlessly integrated with existing regularization methods, yielding additional performance gains.

1. Introduction

The introduction argues that condition-agnostic representation regularizers overlook labels or text that determine generation targets. CARE uses those built-in conditions to guide representation structure, and the paper reports gains in generation quality and condition alignment.

1. Introduction

Existing representation regularizers overlook labels and captions, so similar conditioning inputs may yield scattered features that weaken semantic consistency and controllability.Source: §1 · pp. 1–2
CARE is a plug-and-play loss that uses semantic conditions to shape representations without external foundation models and can combine with existing objectives or regularizers.Source: §1 · p. 2
The authors motivate CARE with an information-theoretic interpretation of conditional representation regularization, consistent with the Platonic Representation Hypothesis.Source: §1 · p. 2
CARE improves generation quality across class-conditional and text-to-image tasks, including a further 16.61% FID reduction on text-to-image benchmarks.Figure 2 reports consistent ImageNet FID improvements throughout training.Source: §1 · p. 2, Figure 2 · p. 2
In text-to-image settings, CARE improves representation–condition alignment while reducing FID, both with and without REPA.The alignment is measured by a linear probe on intermediate representations.Source: Figure 3 · p. 2, §1 · p. 2

2. Related Work

Related work describes representation regularization for diffusion and approaches that incorporate conditioning into generative objectives. CARE instead regularizes the representation space with a lightweight term, without changing the diffusion process or adding parameters.

2.1. Diffusion Models with Regularization

REPA aligns diffusion representations with pretrained encoders, while SARA, REG, and dispersive loss extend alignment or encourage feature dispersion.REPA is associated with faster convergence and improved semantic coherence of generated samples.Source: §2.1 · p. 2

2.2. Diffusion with Condition-Integrated Objectives

Condition-integrated approaches adjust generative distributions or priors; CAR-Flow uses conditioning variables, while Issachar et al. (2025) design condition-specific priors.These approaches may require added parameters or multi-stage training.Source: §2.2 · p. 3
CARE incorporates conditional information in the representation space through lightweight regularization, without altering the diffusion process or adding parameterization.Source: §2.2 · p. 3

3. Method

CARE adds condition-aware regularization to diffusion representations, using condition similarity to modulate repulsion rather than imposing hard constraints. Its instantiations adapt similarity measurement to discrete labels or continuous condition embeddings.

3.2. Preliminaries

Diffusion learns a probability flow from noise toward data, while flow matching predicts a conditioned velocity field that guides generation toward specific data modes.The conditioning variable may be a class label or text embedding.Source: §3.2 · p. 3

3.3. CARE: Condition-Aware REpresentation Regularization

During text-to-image training, intermediate representations become more aligned with text conditions as generation quality improves.This observation motivates organizing representations to preserve conditional structure across semantically similar conditions.Source: §3.3 · p. 3
CARE modulates repulsive regularization by condition similarity, organizing intermediate representations without hard constraints or external encoders.The objective uses condition similarity and a decreasing function to control its influence on representation-space regularization.Source: §3.3 · p. 3, §3.1 · p. 3

3.4. Instantiations

CARE augments the original diffusion objective with an auxiliary regularization loss over representations and conditions.The representation-space kernel is normalized by feature dimension to stabilize its scale in high-dimensional spaces.Source: §3.3 · p. 3, §3.4 · p. 4

3.4.1. CLASS-CONDITIONAL DIFFUSION

For class-conditional generation, similarity is 1 for matching labels and 0 otherwise; ψ assigns α to matching labels and 1 to nonmatches.The class-condition rule is the discrete-label instantiation of CARE’s similarity weighting.Source: §3.4.1 · p. 4

3.4.2. FEATURE-CONDITIONAL DIFFUSION

For continuous condition embeddings, cosine rescaling gives absolute similarity while batch-softmax gives relative similarity, with temperature controlling distribution sharpness.Both formulations use ψ(s)=s·(α−1)+1, with 0 < α < 1.Source: §3.4.2 · p. 4

4.1. Experiments Settings

Experiments evaluate CARE on class-to-image and text-to-image generation, using distinct model setups and evaluation protocols. The section also specifies settings for combining CARE with REPA and for measuring representation alignment.

4.1. Experiments Settings

Class-to-image experiments use 256×256 ImageNet with SiT models trained on SD-VAE latents, and sample with 250-step SDE Euler–Maruyama without classifier guidance.FID and inception score are reported; default settings use batch size 256, λCARE=0.25, the last SiT layer, and α=0.5.Source: §4.1 · p. 4
Text-to-image experiments use a 24-layer MMDiT trained on 595k 256×256 CC3M-subset images with GPT-relabeled captions and CLIP-L condition inputs.CARE is applied to layer 8; evaluation samples 50,000 dataset prompts and reports FID and textual CLIP score.Source: §4.1 · p. 4
For REPA combinations, CARE is applied at REPA’s layer using DINOv2-B, with λREPA set to 0.5 for class-to-image and 0.25 for text-to-image.λCARE is fixed at 0.25 across tasks unless otherwise specified.Source: §4.1 · p. 4
Text-condition alignment is measured by training a single-layer linear probe on representations extracted at t=0 with the empty token, using a 9:1 train-validation split.The probe predicts text-condition embeddings by minimizing negative cosine similarity; model alignment is scored by the highest average validation cosine similarity.Source: §4.1 · pp. 4–5
Table 1 reports CARE comparisons across class-to-image and text-to-image benchmarks.Source: Table 1 · p. 5

4.2. Main Results

CARE improves generation quality across class-conditional and text-to-image tasks and can complement existing representation regularizers. It also strengthens the conditional structure of intermediate representations.

At 400k steps, CARE reduces SiT-XL/2 ImageNet FID from 17.19 to 13.91, a 19.08% relative improvement.With CFG, FID decreases from 5.36 to 4.09, a 23.69% relative gain.Source: §4.2 · pp. 5–6
CARE-trained SiT-XL/2 at 2.4M steps achieves lower FID than baseline SiT-XL/2 at 7M steps.The comparison indicates improved training efficiency in the reported convergence regime.Source: §4.2 · p. 6
At 200k text-to-image steps, CARE reduces FID from 11.32 to 9.44, a 16.61% relative improvement; it also raises textual CLIP score throughout training.At 300k steps, FID decreases from 8.77 to 7.02, a 19.95% improvement.Source: §4.2 · p. 6
Across class-to-image and text-to-image settings, CARE and REPA show complementary and synergistic effects.On SiT-B/2, adding CARE to REPA lowers FID from 24.31 to 21.64, a 10.98% relative improvement.Source: §4.2 · p. 6
CARE consistently improves condition–representation alignment during training, with additional gains when combined with REPA.The results indicate stronger conditional structure and clearer semantic organization of learned features.Source: §4.2 · p. 6

4.3. Ablation Studies

The ablations examine how CARE's condition-agnostic weight, similarity measure, loss coefficient, injection depth, and local-batch label composition affect its performance.

4.3.1. EFFECT OF CONDITION-AGNOSTIC WEIGHT α

Across class-to-image and text-to-image, α = 0.5 generally gives the best trade-off; with REPA on ImageNet, α from 0.01 to 0.001 performs best, while α = 1.0 is condition-agnostic and tends to perform worst.With REPA, the preferred α shifts toward smaller values.Source: §4.3.1 · p. 6, Table 2 · p. 6, Table 3 · p. 6

4.3.2. EFFECT OF SIMILARITY MEASUREMENT

Both linear and softmax similarity measures improve FID and CLIP-T over baseline; linear slightly favors sample quality, while softmax favors condition alignment.The two formulations represent absolute and relative similarity measurements, respectively.Source: §4.3.2 · p. 6, Table 4 · p. 7

4.3.3. EFFECT OF LOSS COEFFICIENT λCARE

Without REPA, λCARE = 0.25 lowers FID from 34.84 to 30.91 and raises IS from 41.53 to 47.22; with REPA, λCARE = 0.5 gives best FID (21.57), while 0.25 gives highest IS (68.50).Without REPA, increasing the coefficient to 0.5 slightly degrades performance; relatively small coefficients suffice for gains.Source: §4.3.3 · p. 7, Table 5 · p. 7

4.3.4. EFFECT OF INJECTION DEPTH

Injecting CARE at deeper layers consistently improves FID and IS; layer 12, the final SiT-B/2 transformer block, achieves the best results.The authors attribute this to deeper-layer regularization acting more directly on representations propagated to the output, while earlier effects may be attenuated by later transformations.Source: §4.3.4 · pp. 7–8, Table 6 · p. 8

4.3.5. ENFORCING DISTINCT LABELS WITHIN LOCAL BATCH

On ImageNet, enforcing distinct labels within local batches lowers dispersive-loss FID from 32.79 to 32.07, while CARE achieves 30.94.The authors distinguish CARE as learning relative representation structure across conditions, not merely avoiding false negatives.Source: §4.3.5 · p. 8, Table 7 · p. 8

4.4. Generalizing to other supervision signals

CARE can use supervision signals beyond text conditions. With REPA on text-image generation, replacing text conditions with DINOv2 image class tokens further improves FID.

With REPA on text-image generation, replacing text conditions with DINOv2 image class tokens improves FID from 7.44 to 7.27.The authors suggest pretrained vision-model signals may provide more semantically aligned guidance than textual conditions.Source: §4.4 · p. 8

5. Discussions and Conclusion

CARE uses built-in conditions to improve the semantic structure of diffusion representations and is connected to dispersive loss. The study focuses on categorical and textual conditions in image generation.

5. Discussions and Conclusion

Dispersive loss is a special case of CARE in which a constant mapping removes conditional supervision, yielding a condition-agnostic unsupervised contrastive loss.The cited work is Wang & He (2025).Source: §5 · p. 8
The study focuses largely on categorical and textual conditions and is limited to diffusion models for image generation.Richer multimodal supervision and video, audio, or 3D generation are identified as future directions.Source: §5 · p. 8
CARE is theoretically connected to mutual-information maximization and empirically enhances sample quality for class-to-image and text-to-image tasks.It requires no external models and complements existing approaches.Source: §5 · p. 8

A. Heuristic Interpretation of CARE via Mutual Information

The appendix offers a heuristic mutual-information interpretation of CARE, then derives a tractable, kernel-based objective and a simplified class-conditional form. The analysis is presented as intuition, not as necessary for defining or optimizing CARE.

A. Heuristic Interpretation of CARE via Mutual Information

CARE’s heuristic interpretation links structured conditional representations to dependence between representations Zθ and conditions C, with mutual information equivalent to minimizing H(C | Zθ).The appendix says this interpretation is not required to define or optimize CARE.Source: A. Heuristic Interpretation of CARE… · p. 11
Because conditional likelihood p(C | Zθ) is intractable, CARE uses a kernel-based Nadaraya–Watson estimate to encourage geometric coherence for representations with similar conditions.The derivation focuses on discrete conditions and defines X = Zθ and Y = 1{C = c}.Source: A. Heuristic Interpretation of CARE… · p. 11
A batch negative log-likelihood built from the kernel estimate is approximated by swapping summation and logarithm, yielding condition-aware supervision plus unsupervised uniformity regularization.The passage notes that this approximation is not exact.Source: A. Heuristic Interpretation of CARE… · p. 11
The surrogate introduces γ to balance its two effects; under discrete conditions and label smoothing, the objective is rewritten using X, Y, β, and γ.The smoothed class-conditional form is Lβ,γ = −log(X + βY) + (1 + γ)log(X + Y).Source: A. Heuristic Interpretation of CARE… · pp. 11–12
After several hundred training steps, r = Y/X falls to the order of 10−4, motivating a first-order approximation of the smoothed objective.The resulting concise form is L = log(αX + Y), where 0 < α < 1; the paper identifies this as CARE under class-conditioned settings.Source: Figure 5 · p. 12, A. Heuristic Interpretation of CARE… · p. 12

B. Additional Visual Results on ImageNet 256 × 256

The figure presents uncurated SiT-XL/2 samples generated after CARE training and specifies the sampler settings.

B. Additional Visual Results on ImageNet 256 × 256

Figure 6 shows uncurated samples from SiT-XL/2 trained with CARE for 2.4M iterations, using 50-step sampling.The caption asks readers to zoom in for the best view.Source: Figure 6 · p. 13
The sampling setup uses an ODE sampler with classifier-free guidance scale 4.0.Source: B. Additional Visual Results on… · p. 13