data
RectifiedFlowConfig
extends
GenerativeModelConfigRectifiedFlowConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', base_channels: int = 128, channel_mult: tuple[int, ...] = (lambda: (1, 2, 2, 2))(), num_res_blocks: int = 4, attention_resolutions: tuple[int, ...] = (lambda: (16,))(), dropout: float = 0.15, resnet_groups: int = 32, fir: bool = False, fir_kernel: tuple[int, ...] = (lambda: (1, 3, 3, 1))(), progressive: Progressive = 'none', progressive_input: ProgressiveInput = 'none', embedding_type: TimeEmbedding = 'positional', fourier_scale: float = 16.0, data_centered: bool = True, skip_rescale: bool = True, init_scale: float = 0.0, scale_by_sigma: bool = False, t_schedule: TimeSchedule | int = 'uniform', time_eps: float = 0.001, solver: str = 'dopri5', rtol: float = 1e-05, atol: float = 1e-05, trace_method: TraceMethod | None = None, trace_noise: TraceNoise = 'rademacher', exact_trace_max_dim: int = 32)Frozen configuration for every Rectified Flow variant.
Parameters
sample_sizeint or tuple of int= 32Spatial resolution of a sample.
in_channelsint= 3Channels per sample.
out_channelsint= 3Kept equal to
in_channels: the field is a velocity in the
sample's own space, so it cannot change the dimension.act_fn(silu, swish, relu, gelu)= "silu"Activation inside the velocity field. The reference
implementation calls this
swish, which is the same function.base_channelsint= 128Width of the first stage — the reference configs'
nf.channel_multtuple of int= (1, 2, 2, 2)Per-stage width multipliers; the length sets the ladder depth.
num_res_blocksint= 4Residual blocks per resolution. Four, not two, on CIFAR-10.
attention_resolutionstuple of int= (16,)Spatial sizes at which self-attention is applied — one resolution
only, unlike the multi-resolution attention of the U-Net Flow
Matching borrows.
dropoutfloat= 0.15Dropout inside residual blocks.
resnet_groupsint= 32Upper bound on GroupNorm groups; the actual count is
min(channels // 4, resnet_groups), as in the reference.firbool= FalseWhether resampling goes through a finite-impulse-response filter
rather than nearest-neighbour / average pooling. Off at
32 x 32, on at 256 x 256.fir_kerneltuple of int= (1, 3, 3, 1)Separable filter taps, normalised to unit sum at build time.
progressive(none, output_skip)= "none"Whether the output is accumulated through a resolution pyramid.
progressive_input(none, input_skip)= "none"Whether a downsampled copy of the input is injected at each stage.
embedding_type(positional, fourier)= "positional"Where the time embedding's frequencies come from.
fourier_scalefloat= 16.0Standard deviation of the random frequencies when
embedding_type="fourier". Ignored otherwise.skip_rescalebool= TrueDivide residual and attention sums by , keeping
the activation variance from growing with depth.
init_scalefloat= 0.0Gain applied to the last convolution of each block. Zero makes
every residual branch start as the identity, which is what the
reference configs set.
scale_by_sigmabool= FalseDivide the network output by the conditioning value
t * 999.
All four released 256x256 configs set this; the CIFAR-10 one does
not, hence the default. It is part of the trained function, so a
checkpoint and its flag have to agree.t_schedule(uniform, t0, t1)= "uniform"Which times the objective draws.
"uniform" is the
rectified-flow / reflow objective; the others pin t and turn
the same loss into k-step distillation.time_epsfloat= 1e-3Lower end of the time interval. Times are drawn from
[time_eps, 1] and solves run over the same range, matching the
reference implementation.solverstr= "dopri5"Method handed to
lucid.diffeq.odeint. dopri5 is the
embedded Runge–Kutta 4(5) pair the reference calls rk45.rtolfloat= 1e-5, 1e-5Solver tolerances — the reference's
ode_tol.atolfloat= 1e-5, 1e-5Solver tolerances — the reference's
ode_tol.trace_method(exact, hutchinson)= "exact"How the divergence is obtained when scoring likelihood.
None
picks by dimension. Irrelevant to training.trace_noise(rademacher, gaussian)= "rademacher"Probe distribution for the Hutchinson estimate.
exact_trace_max_dimint= 32Dimension below which
trace_method=None resolves to
"exact".Notes
in_channels and out_channels must agree: the field lives in
the sample's own space.
Examples
>>> from lucid.models.generative.rectified_flow import RectifiedFlowConfig
>>> RectifiedFlowConfig().t_schedule
'uniform'
>>> RectifiedFlowConfig(sample_size=32, in_channels=3).data_dim
3072
>>> RectifiedFlowConfig(t_schedule=4).is_distillation
TrueUsed by 3
Constructors
1dunder
__init__
→None__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', base_channels: int = 128, channel_mult: tuple[int, ...] = (lambda: (1, 2, 2, 2))(), num_res_blocks: int = 4, attention_resolutions: tuple[int, ...] = (lambda: (16,))(), dropout: float = 0.15, resnet_groups: int = 32, fir: bool = False, fir_kernel: tuple[int, ...] = (lambda: (1, 3, 3, 1))(), progressive: Progressive = 'none', progressive_input: ProgressiveInput = 'none', embedding_type: TimeEmbedding = 'positional', fourier_scale: float = 16.0, data_centered: bool = True, skip_rescale: bool = True, init_scale: float = 0.0, scale_by_sigma: bool = False, t_schedule: TimeSchedule | int = 'uniform', time_eps: float = 0.001, solver: str = 'dopri5', rtol: float = 1e-05, atol: float = 1e-05, trace_method: TraceMethod | None = None, trace_noise: TraceNoise = 'rademacher', exact_trace_max_dim: int = 32)Properties
3int: Flattened width of one sample.
bool: Whether t_schedule pins the time rather than drawing it.
True for "t0", "t1" and any integer k. A distilling
model is trained to be stepped a fixed number of times, so its
loss is only meaningful on reflow pairs.
TraceMethod: trace_method, or the one implied by the dimension.