Comfy-Org / Comfy-Org/ComfyUI

TrainLoraNode: x0-space MSE on FLOW models introduces an implicit sigma^2 loss weighting

Open
#15,847 0 comments 0 reactions 0 assignees View on GitHub
Potential Bug
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

Selecting `loss_function = MSE` should apply an unweighted mean squared error
on the training objective, or — if a timestep weighting is applied — that
weighting should be visible and configurable rather than implied by the choice
of prediction space.

### Actual Behavior

For FLOW models, `TrainLoraNode` computes the loss in x0 space:

```python
x0_pred = model_wrap(xt, sigma, ...)
loss = loss_fn(x0_pred.float(), x0.float())
```

while sigma is sampled by drawing a uniform random `percent` and mapping it
through `model_sampling.percent_to_sigma(...)`.

For ComfyUI's `CONST` flow parameterization, x0-space MSE is **not** equivalent
to unweighted velocity MSE. It is exactly velocity MSE multiplied by `sigma^2`.
The node therefore applies an implicit sigma-dependent weighting that is not
visible or configurable in the UI, and there is currently no
`weighting_scheme` / timestep-sampling input on the node.

**Derivation.** For the `CONST` parameterization in
`comfy/model_sampling.py`:

```
x_t = sigma * eps + (1 - sigma) * x0
x0_pred = x_t - sigma * v_pred
```

with velocity target `v_target = eps - x0`. Then:

```
x0_pred - x0 = sigma*eps + (1-sigma)*x0 - sigma*v_pred - x0
= sigma * (eps - x0 - v_pred)
= -sigma * (v_pred - v_target)
```

and therefore:

```
||x0_pred - x0||^2 = sigma^2 * ||v_pred - v_target||^2
```

so `L_x0 = sigma^2 * L_velocity` for the same sample and prediction.

**Scale.** For two examples with comparable velocity-space squared error, the
relative weight supplied by the x0-space loss is proportional to `sigma^2`.
Comparing sigma 0.90 and 0.05 gives `(0.90/0.05)^2 = 324`.

This should **not** be read as "the effective training gradient is 324x
larger". Actual gradient contributions also depend on the model's errors, its
Jacobians, the `percent -> sigma` mapping and the sampled data. The claim is
narrower: for comparable velocity error at two fixed sigma values, the current
loss parameterization itself supplies a ~3.2e2 weighting ratio.

I am not arguing that any particular weighting must become the default. Several
weighting schemes are legitimate — Diffusers' flow trainers expose `none`,
`sigma_sqrt`, `logit_normal`, `mode` and `cosmap` explicitly. The issue is that
here the weighting is imposed implicitly by the choice of prediction space and
is not surfaced to the user.

### Steps to Reproduce

**Numerical verification of the identity** (no training required). Using
ComfyUI's own `CONST` class with an arbitrary velocity and fixed sigmas, the
ratio between `x0-MSE / sigma^2` and directly computed velocity-MSE is
`1.000000` to test precision:

```
sigma=0.05 L_x0 = 0.0072 sigma^2 * L_v = 0.0072 ratio 1.000000
sigma=0.50 L_x0 = 0.8203 sigma^2 * L_v = 0.8203 ratio 1.000000
sigma=0.90 L_x0 = 2.3532 sigma^2 * L_v = 2.3532 ratio 1.000000
```

**Behavioural A/B.** A paired experiment with the same 4-image dataset, model,
seed, rank, alpha, optimizer, learning rate, 160 steps and timestep sampling.
The only intended variable:

- **A.** stock x0-space MSE
- **B.** velocity-equivalent formulation (see below), tested in a local build

Note that the loss values of the two runs are not directly comparable, because
the change alters the scale of the loss itself. The comparison was made on the
resulting adapters and on paired generations at fixed seeds.

### Debug Logs

```powershell
Adapter magnitudes after the paired 160-step runs:

median |down| for attention.out:
velocity-equivalent (B) 0.4303
stock x0-space (A) 0.3248 (-24%)

Removing the compensation produced **smaller** adapter movement, yet the paired
generations at fixed seeds were visibly worse: one severely malformed output,
stronger texture degradation, and systematic apparent aging of the subject
relative to run B. The compensated run was structurally more stable.

I do not take this as proof that unweighted velocity MSE is optimal for any
particular task. It does show that the implicit `sigma^2` weighting materially
changes training behaviour and is not merely an algebraic curiosity.
```

### Other

**Environment**

- ComfyUI commit `4da9e2dbead52fc1e68beae33fe3d7ad63b63241`, base tag `v0.33.3`
- PyTorch 2.10.0, Python 3.13.12
- macOS 26.6.2, Apple Silicon M5 Pro, device MPS
- Model: Z-Image Base (FLOW / `CONST` sampling), bf16, batch size 1

The behaviour described is present in stock ComfyUI. The velocity-equivalent
formulation below was tested in a local build only.

**Velocity-equivalent formulation** (what I tested — deliberately not called
"the fix", since other weighting schemes may be intentional):

```python
per_elem = (x0_pred.float() - x0.float()).pow(2)
s2 = batch_sigmas.detach().float().pow(2).clamp_min(1e-4)
loss = (per_elem / s2.view(-1, *([1] * (per_elem.ndim - 1)))).mean()
```

For batch size 1 this is algebraically equivalent to velocity MSE under the
`CONST` equations above. A general batched implementation would need per-example
loss reduction before applying the corresponding per-example sigma weighting.
Computing the loss on the velocity directly, before the x0 conversion, would be
equivalent and cleaner.

**Request**

Would it make sense to make the FLOW loss parameterization / timestep weighting
explicit? Options could include:

- direct velocity-space loss for FLOW models;
- x0 loss with an explicitly selected weighting;
- a user-configurable `weighting_scheme` input, as the Diffusers flow trainers
have;
- at minimum, documenting that MSE in x0 space implies a `sigma^2` weighting
under `CONST` sampling.

**Possibly related**

- #13089 — "Lora Training node, loss value consistently exploding", where the
reporter expects loss to stay below 1.0. I do not know whether that report has
the same cause, but the `sigma^2` factor described here does mean that raw
loss magnitude on FLOW models is dominated by which sigma was drawn, which may
be worth considering when interpreting loss values from this node.

Contributor guide

Open the contributing guide

Research direction

Start at TrainLoraNode's FLOW loss path and read the CONST parameterization in comfy/model_sampling.py, then verify the x0-MSE/sigma² relationship with the numerical example in the issue. Done should make the FLOW loss parameterization or timestep weighting explicit and configurable, or document the existing sigma² weighting, without assuming one default.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.