DOCSEARCHbyAlgolia
Submit Docs

rectified-flow-pytorch

repository·main·Indexed 19 days ago

https://github.com/lucidrains/rectified-flow-pytorch

A PyTorch implementation of rectified flow and related research for generative modeling. The library provides the RectifiedFlow and Reflow classes for flow alignment, a Trainer class based on accelerate, and a variety of specialized flow formulations including LapFlow (Laplacian multiscale flow), EquilibriumMatching, NanoFlow, MeanFlow, LsdFlow, and FiT.

Tokens
22.3K
Snippets
57
Records
75
Agent score
66%

What's inside rectified-flow-pytorch

  3. How SoFlow loss components work

    main

    SoFlow optimizes a combined loss function consisting of two main parts:

    1. Flow Matching Loss (flow_matching_loss): Standard rectified flow objective that minimizes the difference between the predicted velocity and the target flow (noise to data).
    2. Solution Consistency Loss (solution_consistency_loss): Enforces consistency between different points along the flow trajectory. It uses an Euler-parameterized solution to ensure that predicting the velocity at time $t$ and stepping to $l$ is consistent with the target trajectory.

    The forward method of SoFlow randomly masks batches between these two losses based on lambda_flow_matching to balance training.

  4. Understand NanoFlow's multi-objective output requirement

    main

    For NanoFlow to function, the underlying model must be designed to output a tensor with a specific structure. The model is expected to return a tensor of shape (batch, 3, *data_shape).

    The second dimension (size 3) is unbinded into three distinct components:

    1. pred_flow: The predicted velocity/flow.
    2. pred_clean: The predicted clean data.
    3. pred_noise: The predicted noise.

    If your model currently only predicts one of these, you must wrap it or modify it to output all three to use the NanoFlow multi-objective training logic.

  5. How RectifiedFlow prediction objectives work

    main

    The predict parameter determines how the model's output is interpreted and how the target flow is derived during training:

    1. flow: The model directly predicts the velocity field $v_t$. The target is the straight-line interpolation between noise and data.
    2. noise: The model predicts the noise $\epsilon$. The flow is derived as $(x_t - \epsilon) / t$. This is often more stable for certain architectures.
    3. clean: The model predicts the clean data $x_1$. The flow is derived as $(x_1 - x_t) / (1 - t)$.

    If mean_variance_net=True, the model outputs both a mean and a standard deviation, allowing for probabilistic flow matching.

  6. How LsdFlow handles conditioning

    main

    Conditioning in LsdFlow is controlled by the accept_cond flag during initialization.

    • If accept_cond=True: The model is expected to accept (x, s, t, cond) as arguments. Both forward and sample methods must receive a cond tensor. If cond is missing, an assertion error will be raised.
    • If accept_cond=False: The model is expected to accept (x, s, t) as arguments. Passing a cond tensor to forward or sample will result in an assertion error.

    This ensures that the conditioning logic is explicitly managed and prevents accidental mismages between the flow wrapper and the underlying neural network.

  7. Use XMWrapper for explorative modeling

    main

    The XMWrapper is a PyTorch Module designed for explorative modeling (specifically for forward XM as described in Alexi Gladstone et al.). It wraps an existing flow_model to allow for evaluating multiple candidates for each sample in a batch and selecting the candidate with the minimum loss.

    Key features:

    • Candidate Evaluation: It repeats inputs $K$ times (where $K$ is the number of candidates) and performs forward passes to find the best candidate per sample.
    • Batch Management: It can handle large numbers of candidates by chunking the computation using max_batch_size to avoid OOM errors.
    • Loss Reduction Support: It automatically detects if the underlying flow_model supports a loss_reduction parameter in its forward method to facilitate per-sample loss calculation.
    from rectified_flow_pytorch.xm_wrapper import XMWrapper

# Assuming flow_model is your existing Rectified Flow model
wrapper = XMWrapper(
    flow_model = flow_model,
    candidates = 5,           # Number of candidates to evaluate per sample
    max_batch_size = 128       # Maximum batch size for chunked forward passes
)

# The wrapper can be used in a training loop like a standard module
loss = wrapper(*args, **kwargs)

  8. How SplitMeanFlow handles conditioning

    main

    Conditioning is handled via the accept_cond flag.

    • If accept_cond=True: The forward and sample methods expect a cond argument. This tensor is passed directly to the underlying model as the fourth positional argument.
    • If accept_cond=False: The cond argument must be None. If you attempt to pass a cond tensor when accept_cond is False, an assertion error will be raised.

    This ensures that the model's signature remains consistent with the training logic.

  9. SoFlow sampling schedules and $r$ values

    main

    The parameter $r$ controls the step size in the Solution Consistency loss. The value of $r$ decays over training steps according to the r_schedule:

    • 'constant': Returns r_end.
    • 'linear': Linearly interpolates from r_init to r_end.
    • 'cosine': Uses a cosine schedule.
    • 'exponential': (Default/Best) Uses an exponential decay from r_init to r_end. This is the recommended schedule per the paper.

  11. How SelfFlow training works

    main

    SelfFlow implements a training procedure where a student model learns to match a teacher model's representations while also learning the flow between noise and data.

    1. Timestep Sampling: Two timesteps are sampled using schedule_fn: one for the teacher and one for the student.
    2. Dual-Timestep Masking: A spatial mask is generated. For masked patches, the student uses the teacher's timestep; otherwise, it uses its own. This forces the student to handle varying levels of 'cleanliness'.
    3. Flow Loss: The student predicts the flow (velocity) between the interpolated state and the target data. The loss is calculated using loss_fn (typically MSE).
    4. Representation Alignment: The student's hidden state (at student_align_layer) is passed through a projector and compared to the teacher's hidden state (at teacher_align_layer) using repr_loss_fn (typically cosine similarity).
    5. Total Loss: total_loss = flow_loss + (repr_loss * repr_loss_weight).

  12. How ValueFlow training paths work

    main

    The ValueFlow class implements a multi-modal training objective that can switch between different learning paradigms based on the input provided:

    1. State Generation: If prob_state_generation is set, the model learns to generate states (e.g., images) conditioned on returns. This is done via standard flow matching between noise and the flow_state.
    2. PPO (Behavioral Cloning): When explicit_target_return is provided, the model learns to match the flow between noise and the target return. This is essentially a conditional flow matching on the value space.
    3. Q-Learning (Bootstrapping): When next_state and next_action are provided, the model uses two losses:
      • DCFM (Dynamic Conditional Flow Matching): Matches the flow of the velocity of the next state.
      • BCFM (Bootstrapped Conditional Flow Matching): Matches the flow between noise and the bootstrapped target return (reward + discounted next Q-value).