kohya-ss / kohya-ss/sd-scripts

Feature Request: Add L0 Loss as an Option for `--loss_type`

Open
#1,841 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

---

### **Description**
I would like to request the inclusion of **L0 loss** as an option for `--loss_type` in `kohya-ss/sd-scripts`. The **L0 loss** is particularly useful for mode-seeking tasks, such as denoising images with impulse noise, where it outperforms L1 and L2 losses.
L0 loss anneals from gamma = 2 to gamma = 0 during training, slowly transitioning from L2, to L1 and finally L0 loss. Epsilon is added for numerical stability.

---

### **Motivation**
The L0 loss is known for its ability to focus on the **mode** of the distribution, unlike L1 and L2 losses which target the median and mean, respectively. This makes L0 loss ideal for tasks involving sparse and extreme noise, like random-valued impulse noise or other multimodal distributions. In the context of diffusion models, L0 loss can complement existing loss functions by improving robustness and precision during fine-tuning, particularly in scenarios where sparse or extreme variations in data quality occur. It provides an additional tool for optimizing model performance when training on noisy or incomplete data.

The theoretical properties of L0 loss have been explored in the following paper:

- **Noise2Noise: Learning Image Restoration without Clean Data**
([https://arxiv.org/abs/1803.04189](https://arxiv.org/abs/1803.04189))

The paper demonstrates that L0 loss performs particularly well when a significant portion of the pixels are corrupted, as seen in the experiments on image restoration. L0 loss was specifically applied for **mode-seeking** in challenging noise conditions.

---

### **Community Implementation**
A community implementation of L0 loss exists in the following repository:
[PistonY/torch-toolbox](https://github.com/PistonY/torch-toolbox/blob/master/torchtoolbox/nn/loss.py#L34)

It provides a ready-to-use version of the L0 loss function for PyTorch.

---

### **Expected Behavior**
1. Add **L0 loss** as an option for the `--loss_type` argument.
2. Allow users to specify L0 loss parameters such as
1. number of steps until gamma fully annealed to 0 - if not set, should default to the total number of training-steps
2. minimum gamma value - defaults to 0

---

### **Potential Benefits**
- Improves performance for noise restoration tasks involving sparse corruption.
- Complements existing loss types (L1 and L2) for users working on denoising and image enhancement.

---

### **References**
- **Paper**: *Noise2Noise: Learning Image Restoration without Clean Data* ([https://arxiv.org/abs/1803.04189](https://arxiv.org/abs/1803.04189))
- **Community Code**: ([PistonY/torch-toolbox](https://github.com/PistonY/torch-toolbox/blob/master/torchtoolbox/nn/loss.py#L34))

---

Thank you for considering this feature request! Let me know if more details are needed or if I can contribute further.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by locating the --loss_type argument and the training-step handling in sd-scripts, then compare the linked torch-toolbox loss implementation. Done means L0 is selectable, its annealing and minimum-gamma parameters have defined defaults, and the training path uses them correctly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.