Implementation:Zai org CogVideo GeneralLPIPSWithDiscriminator

Knowledge Sources	Zai_org_CogVideo
Domains	Video_Generation, Autoencoding, Adversarial_Training
Last Updated	2026-02-10 00:00 GMT

Overview

GeneralLPIPSWithDiscriminator is a combined perceptual and adversarial loss module that trains autoencoders by jointly optimizing pixel-level reconstruction, LPIPS perceptual similarity, and GAN-based adversarial sharpness with adaptive gradient-norm weighting.

Description

This module orchestrates three complementary loss signals for autoencoder training:

1. Pixel-level reconstruction loss: Computes the L1 distance between input and reconstruction, then normalizes it by a learned log-variance parameter (logvar). The negative log-likelihood formulation rec_loss / exp(logvar) + logvar allows the model to learn an uncertainty estimate that automatically balances reconstruction fidelity against other loss terms.

1. LPIPS perceptual loss: Evaluates perceptual similarity using a pre-trained LPIPS network. For video inputs (dims > 2), frames are rearranged from (b, c, t, h, w) to (b*t, c, h, w) so that LPIPS is applied independently to each frame. A random frame is sampled via pick_video_frame for efficiency.

1. Adversarial loss: A PatchGAN-style discriminator (NLayerDiscriminator by default, configurable via discriminator_config) provides adversarial gradients. The discriminator is activated only after disc_start training steps. An adaptive weighting mechanism computes the ratio of gradient norms from the reconstruction loss and the adversarial loss with respect to the last decoder layer, clamped to [0, 1e4], to dynamically balance these objectives.

The forward method supports two optimizer indices: optimizer_idx=0 for the generator (autoencoder) update and optimizer_idx=1 for the discriminator update. The discriminator update computes hinge or vanilla discriminator loss on real and fake logits. Additional regularization terms from the quantizer can be folded into the total loss via regularization_weights.

Usage

Use this loss module when training image or video autoencoders that require high-fidelity perceptual reconstructions with adversarial sharpening. It is the primary loss function for the image autoencoding pipeline in the CogVideo SAT framework, and can be extended to video via the dims parameter.

Code Reference

Source Location

• Repository: Zai_org_CogVideo
• File: sat/sgm/modules/autoencoding/losses/discriminator_loss.py
• Lines: 17-309

Signature

class GeneralLPIPSWithDiscriminator(nn.Module):
    def __init__(
        self,
        disc_start: int,
        logvar_init: float = 0.0,
        disc_num_layers: int = 3,
        disc_in_channels: int = 3,
        disc_factor: float = 1.0,
        disc_weight: float = 1.0,
        perceptual_weight: float = 1.0,
        disc_loss: str = "hinge",
        scale_input_to_tgt_size: bool = False,
        dims: int = 2,
        learn_logvar: bool = False,
        regularization_weights: Union[None, Dict[str, float]] = None,
        additional_log_keys: Optional[List[str]] = None,
        discriminator_config: Optional[Dict] = None,
    ):

Import

from sat.sgm.modules.autoencoding.losses.discriminator_loss import GeneralLPIPSWithDiscriminator

I/O Contract

Inputs

Name	Type	Required	Description
inputs	torch.Tensor	Yes	Original input images or video frames, shape (B, C, H, W) or (B, C, T, H, W)
reconstructions	torch.Tensor	Yes	Autoencoder reconstructions, same shape as inputs
optimizer_idx	int	Yes	0 for generator update, 1 for discriminator update
global_step	int	Yes	Current training step; discriminator activates after disc_start
last_layer	torch.Tensor	Yes	Weight tensor of the last decoder layer, used for adaptive weighting
regularization_log	Dict[str, torch.Tensor]	Yes	Dictionary of auxiliary regularization losses (e.g., from quantizer)
split	str	No	Logging prefix, defaults to "train"
weights	Union[None, float, torch.Tensor]	No	Optional per-sample weights for the NLL loss

Outputs

Name	Type	Description
loss	torch.Tensor	Scalar total loss (generator loss when optimizer_idx=0, discriminator loss when optimizer_idx=1)
log	dict	Dictionary of scalar logging values including total loss, NLL loss, reconstruction loss, perceptual loss, generator/discriminator loss, logvar, and adaptive weight

Usage Examples

# Initialize the loss module
loss_fn = GeneralLPIPSWithDiscriminator(
    disc_start=10000,
    disc_num_layers=3,
    disc_in_channels=3,
    disc_factor=1.0,
    disc_weight=1.0,
    perceptual_weight=1.0,
    disc_loss="hinge",
    dims=2,
)

# Generator step (optimizer_idx=0)
gen_loss, gen_log = loss_fn(
    inputs=original_images,
    reconstructions=reconstructed_images,
    optimizer_idx=0,
    global_step=current_step,
    last_layer=decoder.conv_out.weight,
    regularization_log={"kl_loss": kl_term},
)

# Discriminator step (optimizer_idx=1)
disc_loss, disc_log = loss_fn(
    inputs=original_images,
    reconstructions=reconstructed_images,
    optimizer_idx=1,
    global_step=current_step,
    last_layer=decoder.conv_out.weight,
    regularization_log={},
)

Related Pages

• Principle:Zai_org_CogVideo_Perceptual_Adversarial_Loss