data
MeanFlowConfig
extends
GenerativeModelConfigMeanFlowConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 4, out_channels: int = 4, act_fn: GenerativeActivation = 'silu', patch_size: int = 2, hidden_size: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, frequency_embedding_size: int = 256, num_classes: int = 1000, class_dropout: float = 0.1, time_conditioning: TimeConditioning = 't_interval', time_sampler: TimeSampler = 'lognorm', lognorm_mean: float = -0.4, lognorm_std: float = 1.0, ratio_r_not_t: float = 0.25, adaptive_weight_power: float = 1.0, adaptive_weight_eps: float = 0.001, guidance_scale: float = 1.0, guidance_mix: float = 0.0, guidance_interval: tuple[float, float] = (0.0, 1.0))Configuration for the MeanFlow family.
Parameters
sample_sizeint= 32Spatial extent of the field the model operates on. For the
paper's ImageNet models this is the VAE latent,
32, not the
256-pixel image.in_channelsint= 4, 4Latent channels. Four is the standard VAE tokenizer's width;
the CIFAR-10 experiment works in pixel space at three.
out_channelsint= 4, 4Latent channels. Four is the standard VAE tokenizer's width;
the CIFAR-10 experiment works in pixel space at three.
patch_sizeint= 2Side of the square patch each token covers. The suffix in the
paper's variant names —
B/2 is the Base backbone at patch 2.hidden_sizeint= 768Transformer width.
depthint= 12Number of transformer blocks.
num_headsint= 12Attention heads per block.
mlp_ratiofloat= 4.0Feed-forward expansion inside each block.
num_classesint= 1000Label vocabulary for class conditioning. The extra index
num_classes is the unconditional token the guidance dropout
substitutes.class_dropoutfloat= 0.1Probability of replacing the label with the unconditional token
during training, which is what leaves the model able to produce
the unconditional field guidance needs.
time_conditioning(t_interval, t_r, t_r_interval, interval)= "t_interval"Which time variables the network is conditioned on. The
Jacobian-vector product is always taken with respect to
regardless of this choice — the
encoding changes the network's inputs, not the identity.
time_sampler(lognorm, uniform)= "lognorm"Distribution the pair
(r, t) is drawn from.lognorm_meanfloat= -0.4, 1.0Parameters of the normal that
"lognorm" squashes through the
logistic. The paper's ImageNet setting; its CIFAR-10 setting is
(-2.0, 2.0).lognorm_stdfloat= -0.4, 1.0Parameters of the normal that
"lognorm" squashes through the
logistic. The paper's ImageNet setting; its CIFAR-10 setting is
(-2.0, 2.0).ratio_r_not_tfloat= 0.25Fraction of the batch that gets . The rest
trains the instantaneous velocity, which is what anchors the
field. Zero reduces the method to Flow Matching and does not
produce usable one-step samples.
adaptive_weight_powerfloat= 1.0 in the loss weight .
0 is the plain squared error; 0.5 is close to the
pseudo-Huber loss of prior one-step work.adaptive_weight_epsfloat= 1e-3 above — keeps the weight finite where the error is
near zero.
guidance_scalefloat= 1.0 in the guided target. One disables guidance and
recovers the plain objective.
guidance_mixfloat= 0.0, which mixes the model's own class-conditional
average velocity into the target alongside the class-unconditional
one. The effective scale a sampler sees is
.
guidance_intervaltuple of float= (0.0, 1.0)The range of over which guidance is applied at all.
The larger models restrict it;
XL/2+ uses (0.3, 0.8).Notes
Reference: Geng, Deng, Bai, Kolter, and He, "Mean Flows for One-step Generative Modeling", arXiv:2505.13447, 2025. Architecture and training settings are Table 4; the ablations that pin the defaults are Table 1.
The backbone is DiT's, unchanged — the paper is explicit that it keeps "the DiT architecture blocks untouched" and that architectural improvements are orthogonal. What differs from a diffusion DiT is only the conditioning: two time variables rather than one, each embedded and passed through a two-layer MLP before being summed.
Examples
>>> from lucid.models.generative.mean_flow import MeanFlowConfig
>>> config = MeanFlowConfig()
>>> config.hidden_size, config.depth, config.patch_size
(768, 12, 2)
The variant names encode the backbone size and the patch side, so
L/2 is the Large backbone tokenised at two:
>>> large = MeanFlowConfig(hidden_size=1024, depth=24, num_heads=16)
>>> large.num_patches
256Used by 3
Constructors
1dunder
__init__
→None__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 4, out_channels: int = 4, act_fn: GenerativeActivation = 'silu', patch_size: int = 2, hidden_size: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, frequency_embedding_size: int = 256, num_classes: int = 1000, class_dropout: float = 0.1, time_conditioning: TimeConditioning = 't_interval', time_sampler: TimeSampler = 'lognorm', lognorm_mean: float = -0.4, lognorm_std: float = 1.0, ratio_r_not_t: float = 0.25, adaptive_weight_power: float = 1.0, adaptive_weight_eps: float = 0.001, guidance_scale: float = 1.0, guidance_mix: float = 0.0, guidance_interval: tuple[float, float] = (0.0, 1.0))Initialise the head. See the class docstring for parameters.