Merge pull request #2133 from bghira/feature/reflexflow

add ReflexFlow enhancements to scheduled sampling rollout for flow-matching models (2512.04904v1)

Author: bghira

documentation/OPTIONS.md

@@ -687,6 +687,30 @@ See the [DATALOADER.md](DATALOADER.md#automatic-dataset-oversubscription) guide
 - **What**: The order of the solver used for rollout.
 - **Default**: 2.
 
+### `--scheduled_sampling_reflexflow`
+
+- **What**: Enable ReflexFlow-style enhancements (anti-drift + frequency-compensated weighting) during scheduled sampling for flow-matching models.
+- **Why**: Reduces exposure bias when rolling out flow-matching models by adding directional regularization and bias-aware loss weighting.
+- **Default**: False. Requires `--scheduled_sampling_max_step_offset` > 0.
+
+### `--scheduled_sampling_reflexflow_alpha`
+
+- **What**: Scaling factor for the frequency-compensation weight derived from exposure bias.
+- **Default**: 1.0.
+- **Why**: Higher values up-weight regions with larger exposure bias during rollout for flow-matching models.
+
+### `--scheduled_sampling_reflexflow_beta1`
+
+- **What**: Weight for the ReflexFlow anti-drift (directional) regularizer.
+- **Default**: 10.0.
+- **Why**: Controls how strongly the model is encouraged to align its predicted direction with the target clean sample when using scheduled sampling on flow-matching models.
+
+### `--scheduled_sampling_reflexflow_beta2`
+
+- **What**: Weight for the ReflexFlow frequency-compensation (loss reweighting) term.
+- **Default**: 1.0.
+- **Why**: Scales the reweighted flow-matching loss, matching the β₂ knob described in the ReflexFlow paper.
+
 ---
 
 ## 🎯 CREPA (Cross-frame Representation Alignment)

documentation/experimental/SCHEDULED_SAMPLING.md

@@ -61,6 +61,17 @@ The solver used for the rollout generation steps.
 *   **Choices:** `unipc` (recommended, fast & accurate), `euler`, `dpm`, `rk4`.
 *   `unipc` is generally the best trade-off between speed and accuracy for these short sampling bursts.
 
+### Flow Matching + ReflexFlow
+
+For flow-matching models (`--prediction_type flow_matching`), scheduled sampling now supports ReflexFlow-style exposure bias mitigation:
+
+*   `scheduled_sampling_reflexflow`: Enable ReflexFlow enhancements during rollout.
+*   `scheduled_sampling_reflexflow_alpha`: Scale the exposure-bias-based loss weight (frequency compensation).
+*   `scheduled_sampling_reflexflow_beta1`: Scale the directional anti-drift regularizer (default 10.0 to mirror the paper).
+*   `scheduled_sampling_reflexflow_beta2`: Scale the frequency-compensated loss (default 1.0).
+
+These reuse the rollout predictions/latents you already compute, avoiding an extra gradient pass, and help keep biased rollouts aligned with the clean trajectory while emphasizing missing low-frequency components early in denoising.
+
 ### Performance Impact
 
 > ⚠️ **Warning:** Enabling rollout requires running the model in inference mode *inside* the training loop.

simpletuner/helpers/models/common.py

@@ -2739,6 +2739,7 @@ def loss(self, prepared_batch: dict, model_output, apply_conditioning_mask: bool
         """
         target = self.get_prediction_target(prepared_batch)
         model_pred = model_output["model_prediction"]
+        extra_sample_loss = None
         if target is None:
             raise ValueError("Target is None. Cannot compute loss.")
 
@@ -2766,6 +2767,36 @@ def loss(self, prepared_batch: dict, model_output, apply_conditioning_mask: bool
         elif self.PREDICTION_TYPE == PredictionTypes.FLOW_MATCHING:
             # Flow matching always uses L2 loss
             loss = (model_pred.float() - target.float()) ** 2
+            if getattr(self.config, "scheduled_sampling_reflexflow", False):
+                clean_pred = prepared_batch.get("_reflexflow_clean_pred")
+                biased_pred = prepared_batch.get("_reflexflow_biased_pred")
+                beta2 = getattr(self.config, "scheduled_sampling_reflexflow_beta2", 1.0)
+                beta2 = 1.0 if beta2 is None else float(beta2)
+                if clean_pred is not None and biased_pred is not None:
+                    exposure = (biased_pred - clean_pred).detach()
+                    norm_dims = tuple(range(1, exposure.dim()))
+                    exposure_norm = exposure.abs().sum(dim=norm_dims, keepdim=True).clamp_min(1e-6)
+                    alpha = float(getattr(self.config, "scheduled_sampling_reflexflow_alpha", 1.0) or 0.0)
+                    if alpha != 0.0:
+                        weight = 1.0 + alpha * exposure / exposure_norm
+                        loss = loss * weight
+                if beta2 != 1.0:
+                    loss = loss * beta2
+
+                adr_scale = float(getattr(self.config, "scheduled_sampling_reflexflow_beta1", 10.0) or 0.0)
+                if adr_scale != 0.0:
+                    biased_latents = prepared_batch.get("noisy_latents")
+                    clean_latents = prepared_batch.get("latents")
+                    if biased_latents is not None and clean_latents is not None:
+                        target_vec = clean_latents - biased_latents
+                        flat_target = target_vec.reshape(target_vec.shape[0], -1)
+                        flat_pred = model_pred.reshape(model_pred.shape[0], -1)
+                        target_norm = torch.norm(flat_target, dim=1, keepdim=True).clamp_min(1e-6)
+                        pred_norm = torch.norm(flat_pred, dim=1, keepdim=True).clamp_min(1e-6)
+                        target_dir = flat_target / target_norm
+                        pred_dir = flat_pred / pred_norm
+                        adr = (pred_dir - target_dir).pow(2).sum(dim=1)
+                        extra_sample_loss = adr_scale * adr
         elif self.PREDICTION_TYPE in [
             PredictionTypes.EPSILON,
             PredictionTypes.V_PREDICTION,
@@ -2868,7 +2899,10 @@ def loss(self, prepared_batch: dict, model_output, apply_conditioning_mask: bool
                 mask_image = (mask_image > 0).to(dtype=loss.dtype, device=loss.device)
                 loss = loss * mask_image
 
-        loss = loss.mean(dim=list(range(1, len(loss.shape)))).mean()
+        loss = loss.mean(dim=list(range(1, len(loss.shape))))
+        if extra_sample_loss is not None:
+            loss = loss + extra_sample_loss.to(device=loss.device, dtype=loss.dtype)
+        loss = loss.mean()
         return loss
 
     def auxiliary_loss(self, model_output, prepared_batch: dict, loss: torch.Tensor):

simpletuner/helpers/scheduled_sampling/rollout.py

@@ -50,6 +50,154 @@ def _prediction_to_x(
     return model_pred
 
 
+def _flow_timestep_to_sigma(noise_schedule, timestep: int, num_train_timesteps: int) -> float:
+    """
+    Map an integer timestep index to a flow-matching sigma value.
+    Prefer the scheduler's native sigmas if available, otherwise fall back to a linear map.
+    """
+    try:
+        schedule_sigmas = getattr(noise_schedule, "sigmas", None)
+        if schedule_sigmas is not None and int(timestep) < len(schedule_sigmas):
+            return float(schedule_sigmas[int(timestep)])
+    except Exception:
+        pass
+
+    denom = max(float(num_train_timesteps - 1), 1.0)
+    return float(timestep) / denom
+
+
+def _apply_flow_matching_rollout(model, prepared_batch: dict, noise_schedule, config) -> dict:
+    """
+    Scheduled sampling rollout for flow-matching models with optional ReflexFlow caches.
+    This runs tiny self-inference bursts to replace noisy_latents/timesteps/sigmas.
+    """
+    plan = prepared_batch.get("scheduled_sampling_plan")
+    if plan is None:
+        return prepared_batch
+
+    rollout_steps = getattr(plan, "rollout_steps", None)
+    source_ts = getattr(plan, "source_timesteps", None)
+    target_ts = getattr(plan, "target_timesteps", None)
+    if rollout_steps is None or source_ts is None or target_ts is None:
+        return prepared_batch
+
+    latents = prepared_batch["latents"]
+    noise = prepared_batch.get("input_noise", prepared_batch.get("noise"))
+    if noise is None:
+        return prepared_batch
+
+    device = latents.device
+    dtype = latents.dtype
+    base_noisy = prepared_batch["noisy_latents"]
+    base_timesteps = prepared_batch["timesteps"]
+    base_sigmas = prepared_batch.get("sigmas")
+    num_train_timesteps = int(getattr(getattr(noise_schedule, "config", None), "num_train_timesteps", 1000) or 1000)
+    denom = max(float(num_train_timesteps - 1), 1.0)
+
+    reflex_enabled = bool(getattr(config, "scheduled_sampling_reflexflow", False))
+    clean_preds = torch.zeros_like(latents) if reflex_enabled else None
+    biased_preds = torch.zeros_like(latents) if reflex_enabled else None
+
+    new_noisy = base_noisy.clone()
+    new_timesteps = base_timesteps.clone()
+    new_sigmas = base_sigmas.clone() if base_sigmas is not None else torch.zeros_like(base_timesteps, dtype=dtype)
+
+    bsz = latents.shape[0]
+    for i in range(bsz):
+        offset = int(rollout_steps[i].item())
+        target_t = int(target_ts[i].item())
+        source_t = int(source_ts[i].item())
+        target_t = max(0, min(target_t, num_train_timesteps - 1))
+        source_t = max(0, min(source_t, num_train_timesteps - 1))
+        source_frac = float(source_t) / denom
+        target_frac = float(target_t) / denom
+
+        # Always record the "clean" prediction at the target timestep for ReflexFlow weighting.
+        if reflex_enabled:
+            clean_batch = _slice_batch_for_index(prepared_batch, i, device)
+            clean_batch["noisy_latents"] = base_noisy[i : i + 1]
+            clean_batch["timesteps"] = base_timesteps[i : i + 1]
+            if base_sigmas is not None:
+                clean_batch["sigmas"] = base_sigmas[i : i + 1]
+            clean_batch.pop("scheduled_sampling_plan", None)
+            if getattr(config, "controlnet", False):
+                clean_out = model.controlnet_predict(prepared_batch=clean_batch)
+            else:
+                clean_out = model.model_predict(prepared_batch=clean_batch)
+            clean_pred = clean_out.get("model_prediction", clean_out) if isinstance(clean_out, dict) else clean_out
+            clean_preds[i : i + 1] = clean_pred.to(device=device, dtype=dtype)
+
+        # No rollout requested; keep original noisy/step/prediction.
+        if offset <= 0 or source_t <= target_t:
+            if reflex_enabled:
+                biased_preds[i : i + 1] = clean_preds[i : i + 1]
+            continue
+
+        source_sigma = _flow_timestep_to_sigma(noise_schedule, source_t, num_train_timesteps)
+        target_sigma = _flow_timestep_to_sigma(noise_schedule, target_t, num_train_timesteps)
+
+        current = (1 - source_frac) * latents[i : i + 1] + source_frac * noise[i : i + 1]
+        current_frac = source_frac
+        current_sigma = source_sigma
+
+        for t in range(source_t, target_t, -1):
+            if t <= 0:
+                break
+            next_t = max(target_t, t - 1)
+            next_frac = float(next_t) / denom
+            next_sigma = _flow_timestep_to_sigma(noise_schedule, next_t, num_train_timesteps)
+
+            mini_batch = _slice_batch_for_index(prepared_batch, i, device)
+            mini_batch["noisy_latents"] = current
+            mini_batch["timesteps"] = torch.tensor([t], device=device, dtype=base_timesteps.dtype)
+            mini_batch["sigmas"] = torch.tensor([current_sigma], device=device, dtype=dtype)
+            mini_batch.pop("scheduled_sampling_plan", None)
+
+            if getattr(config, "controlnet", False):
+                model_out = model.controlnet_predict(prepared_batch=mini_batch)
+            else:
+                model_out = model.model_predict(prepared_batch=mini_batch)
+
+            model_pred = model_out.get("model_prediction", model_out) if isinstance(model_out, dict) else model_out
+            delta_t = next_frac - current_frac
+            current = current + delta_t * model_pred
+            current_frac = next_frac
+            current_sigma = next_sigma
+
+        # One final prediction at the target state for FC weighting
+        if reflex_enabled:
+            final_batch = _slice_batch_for_index(prepared_batch, i, device)
+            final_batch["noisy_latents"] = current
+            final_batch["timesteps"] = torch.tensor([target_t], device=device, dtype=base_timesteps.dtype)
+            final_batch["sigmas"] = torch.tensor([max(current_sigma, target_sigma)], device=device, dtype=dtype)
+            final_batch.pop("scheduled_sampling_plan", None)
+            if getattr(config, "controlnet", False):
+                final_out = model.controlnet_predict(prepared_batch=final_batch)
+            else:
+                final_out = model.model_predict(prepared_batch=final_batch)
+            final_pred = final_out.get("model_prediction", final_out) if isinstance(final_out, dict) else final_out
+            biased_preds[i : i + 1] = final_pred.to(device=device, dtype=dtype)
+
+        new_noisy[i : i + 1] = current.to(device=device, dtype=dtype)
+        new_timesteps[i] = torch.as_tensor(target_t, device=device, dtype=base_timesteps.dtype)
+        if new_sigmas is not None:
+            new_sigmas[i] = torch.as_tensor(target_sigma, device=device, dtype=dtype)
+
+    prepared_batch["noisy_latents"] = new_noisy
+    prepared_batch["timesteps"] = new_timesteps
+    if new_sigmas is not None:
+        prepared_batch["sigmas"] = new_sigmas
+
+    if reflex_enabled:
+        prepared_batch["_reflexflow_clean_pred"] = clean_preds.detach()
+        prepared_batch["_reflexflow_biased_pred"] = biased_preds.detach()
+        prepared_batch["_reflexflow_pre_rollout_noisy"] = base_noisy
+        prepared_batch["_reflexflow_pre_rollout_timesteps"] = base_timesteps
+        prepared_batch["_reflexflow_pre_rollout_sigmas"] = base_sigmas
+
+    return prepared_batch
+
+
 @torch.no_grad()
 def apply_scheduled_sampling_rollout(model, prepared_batch: dict, noise_schedule, config) -> dict:
     """
@@ -60,6 +208,9 @@ def apply_scheduled_sampling_rollout(model, prepared_batch: dict, noise_schedule
     if plan is None:
         return prepared_batch
 
+    if model.PREDICTION_TYPE is PredictionTypes.FLOW_MATCHING:
+        return _apply_flow_matching_rollout(model, prepared_batch, noise_schedule, config)
+
     if model.PREDICTION_TYPE not in [PredictionTypes.EPSILON, PredictionTypes.V_PREDICTION]:
         return prepared_batch
 

simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py

@@ -310,6 +310,90 @@ def register_advanced_fields(registry: "FieldRegistry") -> None:
         )
     )
 
+    registry._add_field(
+        ConfigField(
+            name="scheduled_sampling_reflexflow",
+            arg_name="--scheduled_sampling_reflexflow",
+            ui_label="Enable ReflexFlow Enhancements",
+            field_type=FieldType.CHECKBOX,
+            tab="training",
+            section="loss_functions",
+            subsection="advanced",
+            default_value=False,
+            help_text="Apply ReflexFlow anti-drift and frequency-compensation weighting during scheduled sampling for flow-matching models.",
+            tooltip="Adds ADR directional regularization and exposure-bias weighting to rollout samples.",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            order=39,
+            dependencies=[
+                FieldDependency(field="scheduled_sampling_max_step_offset", operator="greater_than", value=0, action="show")
+            ],
+            documentation="OPTIONS.md#--scheduled_sampling_reflexflow",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="scheduled_sampling_reflexflow_alpha",
+            arg_name="--scheduled_sampling_reflexflow_alpha",
+            ui_label="ReflexFlow FC Alpha",
+            field_type=FieldType.NUMBER,
+            tab="training",
+            section="loss_functions",
+            subsection="advanced",
+            default_value=1.0,
+            help_text="Scaling for exposure-bias-based loss reweighting (frequency compensation) during ReflexFlow.",
+            tooltip="Higher values up-weight regions with larger exposure bias during rollout.",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            order=40,
+            dependencies=[
+                FieldDependency(field="scheduled_sampling_reflexflow", operator="equals", value=True, action="show")
+            ],
+            documentation="OPTIONS.md#--scheduled_sampling_reflexflow_alpha",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="scheduled_sampling_reflexflow_beta1",
+            arg_name="--scheduled_sampling_reflexflow_beta1",
+            ui_label="ReflexFlow ADR Weight",
+            field_type=FieldType.NUMBER,
+            tab="training",
+            section="loss_functions",
+            subsection="advanced",
+            default_value=10.0,
+            help_text="Directional regularization strength for ReflexFlow anti-drift rectification.",
+            tooltip="Scales the unit-direction alignment term between predicted velocity and target direction.",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            order=41,
+            dependencies=[
+                FieldDependency(field="scheduled_sampling_reflexflow", operator="equals", value=True, action="show")
+            ],
+            documentation="OPTIONS.md#--scheduled_sampling_reflexflow_beta1",
+        )
+    )
+
+    registry._add_field(
+        ConfigField(
+            name="scheduled_sampling_reflexflow_beta2",
+            arg_name="--scheduled_sampling_reflexflow_beta2",
+            ui_label="ReflexFlow FC Weight",
+            field_type=FieldType.NUMBER,
+            tab="training",
+            section="loss_functions",
+            subsection="advanced",
+            default_value=1.0,
+            help_text="Weight for the ReflexFlow frequency-compensated loss term.",
+            tooltip="Scales the exposure-bias-reweighted flow-matching loss (β2 in the paper).",
+            importance=ImportanceLevel.EXPERIMENTAL,
+            order=42,
+            dependencies=[
+                FieldDependency(field="scheduled_sampling_reflexflow", operator="equals", value=True, action="show")
+            ],
+            documentation="OPTIONS.md#--scheduled_sampling_reflexflow_beta2",
+        )
+    )
+
     # Flow Matching Configuration
     registry._add_field(
         ConfigField(