EDM2 loss weighting?

[Seedmanc]

So what's that and why isn't it listed in improvements over original? The settings look so complicated, there's no way to use it without a guide.

[67372a]

Hey @Seedmanc , that is true.

I'm just not good about updating documentation, the fork is mainly used by me and a smaller group that I tend to communicate with and keep up to speed on the details.

It's based on https://arxiv.org/abs/2312.02696.

Here are some reasonable defaults, with a quick explainer:

edm2_loss_weighting = "True"
If it's enabled

edm2_loss_weighting_optimizer = "LoraEasyCustomOptimizer.fmarscrop.FMARSCropV2ExMachina"
The fully qualified optimizer class name for the optimizer to use to optimize the EDM2 loss weighting model.

edm2_loss_weighting_optimizer_lr = "2e-2"
The LR for the loss weighting optimizer, typically 2e-2 for new weights, 5e-3 for established weights.

edm2_loss_weighting_optimizer_args = "{'update_strategy':'cautious', 'gamma':0.0, 'betas':(0.99,0.9999,0.999), 'adaptive_clip':0}"
A JSON object of the optimizer args.

edm2_loss_weighting_lr_scheduler = "True"
If a inverse sqrt LR scheduler is applied to the LR.

edm2_loss_weighting_lr_scheduler_warmup_percent = "0.1"
How many steps to of the total training to warmup

edm2_loss_weighting_lr_scheduler_constant_percent = "0.9"
How many steps to of the total training to keep the LR at the defined amount, after which it will decay using inverse sqrt.

edm2_loss_weighting_max_grad_norm = "0"
Max grad norm for the edm2 loss weighting grads

edm2_loss_weighting_generate_graph_output_dir = "D:/ai/training/loss_weighting"
Where to save image graphs showing the loss weighting amounts across timesteps.

edm2_loss_weighting_generate_graph_every_x_steps = "10"
How often to save graphs.

edm2_loss_weighting_generate_graph = "True"
If graphs should be generated

edm2_loss_weighting_num_channels = "448"
The number of channels for the edm2 loss weighting model, more channels means it has more degrees of freedom and can fit more precisely, 448 is the highest Nvidia went in their paper.

edm2_loss_weighting_generate_graph_y_limit = "5"
The y-limit for the graph, if unset, is dynamic, keeps the scale from jumping around.

edm2_loss_weighting_initial_weights = ""
The full path to an existing edm2 weights file, as it makes sense to reuse weights if already determined for a given dataset and primary model.

[Seedmanc]

So is it like an alternative to validation loss? Which is better? That one requires only a few lines in extra params to use.
I'm coming from vanilla ETS and I wonder about the loss improvements in this fork.

[67372a]

@Seedmanc validation loss and EDM2 loss weighting are different things.

Validation loss simply calculates loss for reporting and analysis purposes, based on a percentage of images that are withheld from training and only used for validation. The intent being to see how well it can generalize past images it is explicitly trained on, which can help determine if it is learning effectively or overfitting.

EDM2 loss weighting basically brings up the loss for timesteps where it is lower, so that it can learn more effectively.

[Seedmanc]

Is it compatible with vpred models? Those usually require zero-terminal.

[67372a]

@Seedmanc yes, edm2 is compatible with vpred and ztsnr