diff --git a/src/refiners/training_utils/config.py b/src/refiners/training_utils/config.py
index 097e74f..6c21461 100644
--- a/src/refiners/training_utils/config.py
+++ b/src/refiners/training_utils/config.py
@@ -165,6 +165,7 @@ class OptimizerConfig(BaseModel):
                     betas=self.betas,
                     weight_decay=self.weight_decay,  # type: ignore
                     safeguard_warmup=True,
+                    use_bias_correction=True,  # recommended for diffusion models
                 )