incorporate zero centered gradient penalty on real and fake images for stability, paper out of Cornell and Brown

Author: lucidrains

README.md

@@ -278,3 +278,12 @@ $ accelerate launch train.py
     url     ={https://api.semanticscholar.org/CorpusID:269214195}
 }
 ```
+
+```bibtex
+@inproceedings{Huang2025TheGI,
+    title   = {The GAN is dead; long live the GAN! A Modern GAN Baseline},
+    author  = {Yiwen Huang and Aaron Gokaslan and Volodymyr Kuleshov and James Tompkin},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:275405495}
+}
+```

gigagan_pytorch/gigagan_pytorch.py

@@ -122,6 +122,7 @@ def gradient_penalty(
     outputs,
     grad_output_weights = None,
     weight = 10,
+    center = 0.,
     scaler: GradScaler | None = None,
     eps = 1e-4
 ):
@@ -151,7 +152,7 @@ def gradient_penalty(
         gradients = maybe_scaled_gradients * inv_scale
 
     gradients = rearrange(gradients, 'b ... -> b (...)')
-    return weight * ((gradients.norm(2, dim = 1) - 1) ** 2).mean()
+    return weight * ((gradients.norm(2, dim = 1) - center) ** 2).mean()
 
 # hinge gan losses
 
@@ -2308,9 +2309,11 @@ def train_discriminator_step(
                 # detach output of generator, as training discriminator only
 
                 images.detach_()
+                images.requires_grad_()
 
                 for rgb in rgbs:
                     rgb.detach_()
+                    rgb.requires_grad_()
 
             # main divergence loss
 
@@ -2352,13 +2355,22 @@ def train_discriminator_step(
                 gp_loss = 0.
 
                 if apply_gradient_penalty:
-                    gp_loss = gradient_penalty(
+                    real_gp_loss = gradient_penalty(
                         real_images,
                         outputs = [real_logits, *real_multiscale_logits],
                         grad_output_weights = [1., *(self.multiscale_divergence_loss_weight,) * len(real_multiscale_logits)],
                         scaler = self.D_opt.scaler
                     )
 
+                    fake_gp_loss = gradient_penalty(
+                        images,
+                        outputs = [fake_logits, *fake_multiscale_logits],
+                        grad_output_weights = [1., *(self.multiscale_divergence_loss_weight,) * len(fake_multiscale_logits)],
+                        scaler = self.D_opt.scaler
+                    )
+
+                    gp_loss = real_gp_loss + fake_gp_loss
+
                     if not torch.isnan(gp_loss):
                         total_gp_loss += (gp_loss.item() / grad_accum_every)
 

gigagan_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.2.20'
+__version__ = '0.3.0'