data

MeanFlowOutput

extendsModelOutput

MeanFlowOutput(loss: Tensor, prediction: Tensor, target: Tensor)

source edit

What a MeanFlow training step reports.

Attributes

lossTensor

The adaptively weighted regression error, a scalar.

predictionTensor

(B, C, H, W) — the model's average velocity

u_\theta(z_t, r, t)

at the sampled times.

targetTensor

The stop-gradiented target the prediction is regressed onto. Reported because the gap between the two is the objective, and because a target that has drifted from the prediction's scale is the first symptom of a mis-specified tangent.

Examples

>>> import lucid
>>> from lucid.models.generative.mean_flow._model import MeanFlowOutput
>>> out = MeanFlowOutput(
...     loss=lucid.zeros(()),
...     prediction=lucid.zeros(1, 3, 32, 32),
...     target=lucid.zeros(1, 3, 32, 32),
... )
>>> out.prediction.shape, out.target.shape
((1, 3, 32, 32), (1, 3, 32, 32))
The target is returned alongside because it is computed, not given —
it depends on the model's own output at another time, so a caller
cannot reconstruct it afterwards.

Used by 1

lucid.models.generative.mean_flow

Constructors

dunder

init

→None

__init__(loss: Tensor, prediction: Tensor, target: Tensor)

source edit

Initialise the head. See the class docstring for parameters.

>>> import lucid >>> from lucid.models.generative.mean_flow._model import MeanFlowOutput >>> out = MeanFlowOutput( ... loss=lucid.zeros(()), ... prediction=lucid.zeros(1, 3, 32, 32), ... target=lucid.zeros(1, 3, 32, 32), ... ) >>> out.prediction.shape, out.target.shape ((1, 3, 32, 32), (1, 3, 32, 32)) The target is returned alongside because it is computed, not given — it depends on the model's own output at another time, so a caller cannot reconstruct it afterwards.