Conditional Embedding Perturbation (CEP) by Koratahiu · Pull Request #1235 · Nerogar/OneTrainer

Koratahiu · 2025-12-30T14:17:40Z

This draft implements the Conditional Embedding Perturbation (CEP) strategy proposed in the paper:
Slight Corruption in Pre-training Data Makes Better Diffusion Models (NeurIPS 2024 spotlight)

This method aims to improve the generation quality and diversity of diffusion models by mitigating the impact of "perfect" overfitting to training pairs. The paper demonstrates theoretically that standard training can cause the generated distribution to collapse to the empirical distribution of the training data.

CEP addresses this by introducing slight, dimension-scaled noise to the conditional embeddings (e.g., text encoder outputs) during training. By optimizing the objective, the model is forced to learn a smoother conditional manifold, reducing the distance to the true data distribution and preventing memorization.

Implementation Details

Adds a perturbation term $\delta$ to the text embeddings before they are passed to the model.
The noise is sampled from a Uniform distribution and scaled by the embedding dimension, ensuring the corruption remains "slight" regardless of the model architecture (SD 1.5 vs SDXL vs Flux).
All models are supported with UI

Usage

Enable Conditional Embedding Perturbation (CEP) (below timestep shifting)
Set CEP Gamma to 1

TODO

To be tested

Koratahiu · 2026-02-07T12:08:57Z

This has been tested with SDXL, Chroma, and Zib, and it works very well.
It is especially beneficial for Zib, which relies on semantic patterns that CEP mitigates through its perturbation noise.

dxqb

interesting. code comments added, but I want to try it myself too. I remember that some people on Discord have tested it. Would be great if they could post their conclusions here also.

Flux2 was added in the meantime

dxqb · 2026-02-12T10:04:08Z

modules/ui/TrainingTab.py


-
+        # Conditional Embedding Perturbation (CEP)
+        cep_label = components.label(frame, 10, 0, "Conditional Embedding Perturbation (CEP)",


is there a gamma that is a no-op?
in that case we wouldn't need an enabled switch. This is how most other parameters in OneTrainer work, that there is a 0.0 which doesn't do anything for example

Yeah, 0 is a no-op.
1 is the paper's default value (slight noise based on the dimension of the TEs), 2 is double that, and so on.

dxqb · 2026-02-12T10:04:47Z

modules/modelSetup/BaseChromaSetup.py

                text_encoder_dropout_probability=config.text_encoder.dropout_probability,
            )

+            if config.cep_enabled:


I'd prefer this call in Model.encode_text
this is where similar functionality is implemented (such as caption dropout)

Would model.encode_text do it on-the-fly without caching?
One benefit of this method is that it doesn't need re-caching

model.encode_text takes the cached output and returns it, but it can (and does) still modify the cached output before returning it. doesn't mean you have to cache the perturbation, it can applied to the cached value.

dxqb · 2026-02-12T10:08:44Z

modules/ui/TrainingTab.py

            components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting")

-
+        # Conditional Embedding Perturbation (CEP)


I think this option fits better near "Caption Dropout Probability"

I think it fits both: injected 'noise' applied to the TE conditioning. However, wouldn't TE settings require per-model setting application? I'm trying to avoid that

dxqb · 2026-02-15T18:40:56Z

modules/modelSetup/mixin/ModelSetupNoiseMixin.py


        return noise

+    def _apply_conditional_embedding_perturbation(


doesn't use self
@staticmethod and no self

dxqb · 2026-02-15T19:16:01Z

modules/modelSetup/mixin/ModelSetupNoiseMixin.py

+
+            # gamma controls perturbation magnitude (Paper uses gamma=1.0 as default baseline)
+            # Calculate scaling factor: sqrt(gamma / d)
+            scale = math.sqrt(gamma / d)


I think this should be
scale = gamma / math.sqrt(d)

Yeah, you're right; I had (1/√d) in my mind when I wrote this

dxqb · 2026-02-15T19:47:26Z

modules/modelSetup/BaseChromaSetup.py

            )

+            if config.cep_enabled:
+                text_encoder_output = self._apply_conditional_embedding_perturbation(


should CEP also be applied during validation? it currently is - validation uses the same predict().
theoretically I guess not, because you want validation to be deterministic and comparable across time. but the effect might be minor.

Isn't this also the case with caption dropout (which is in model.encode_text)?

good point, but that's definitely not good. I've added it here: #957 (comment)

dxqb · 2026-02-15T20:01:17Z

* The noise is sampled from a Uniform distribution and scaled by the embedding dimension, ensuring the corruption remains "slight" regardless of the model architecture (SD 1.5 vs SDXL vs Flux).

I think it might still need tuning per model, because the magnitude of embeddings are different by text encoder
the paper/PR only corrects for dimension

Koratahiu added 5 commits December 30, 2025 17:04

initial CEP

8b013b0

gamma 1

e4c2521

adjust position

7cd25b8

fix transformer (CEP)

39af0d3

Merge branch 'master' of https://github.com/Nerogar/OneTrainer into cep

df574e0

Koratahiu marked this pull request as ready for review February 7, 2026 12:08

dxqb requested changes Feb 12, 2026

View reviewed changes

dxqb reviewed Feb 12, 2026

View reviewed changes

Merge branch 'upstream' into pr-1235

85a5d38

dxqb reviewed Feb 15, 2026

View reviewed changes

add flux2

a860bb5



		# Conditional Embedding Perturbation (CEP)
		cep_label = components.label(frame, 10, 0, "Conditional Embedding Perturbation (CEP)",

		components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting")


		# Conditional Embedding Perturbation (CEP)

Uh oh!

Conversation

Koratahiu commented Dec 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Implementation Details

Usage

TODO

Uh oh!

Koratahiu commented Feb 7, 2026

Uh oh!

dxqb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dxqb Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dxqb Feb 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dxqb Feb 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dxqb commented Feb 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Koratahiu commented Dec 30, 2025 •

edited

Loading

dxqb left a comment •

edited

Loading

dxqb Feb 12, 2026 •

edited

Loading

dxqb Feb 15, 2026 •

edited

Loading

dxqb Feb 15, 2026 •

edited

Loading

dxqb commented Feb 15, 2026 •

edited

Loading