Merge pull request #2138 from bghira/feature/reflexflow-by-default

ReflexFlow: enable by default when flow-matching scheduled sampling is enabled

Author: bghira

documentation/OPTIONS.md

@@ -691,7 +691,7 @@ See the [DATALOADER.md](DATALOADER.md#automatic-dataset-oversubscription) guide
 
 - **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.
+- **Default**: Auto-enable for flow-matching models when `--scheduled_sampling_max_step_offset` > 0; override with `--scheduled_sampling_reflexflow=false`.
 
 ### `--scheduled_sampling_reflexflow_alpha`
 

documentation/experimental/SCHEDULED_SAMPLING.md

@@ -65,7 +65,7 @@ The solver used for the rollout generation steps.
 
 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`: Enable ReflexFlow enhancements during rollout (auto-enabled for flow-matching models when scheduled sampling is active; pass `--scheduled_sampling_reflexflow=false` to opt out).
 *   `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).

simpletuner/helpers/models/common.py

@@ -258,6 +258,7 @@ def __init__(self, config: dict, accelerator):
         self.setup_training_noise_schedule()
         self.diff2flow_bridge = None
         self.setup_diff2flow_bridge()
+        self._maybe_enable_reflexflow_default()
         self.assistant_adapter_name = "assistant"
         self.assistant_lora_loaded = False
 
@@ -2428,6 +2429,29 @@ def sample_flow_sigmas(self, batch: dict, state: dict) -> tuple[torch.Tensor, to
         timesteps = sigmas * 1000.0
         return sigmas, timesteps
 
+    def _maybe_enable_reflexflow_default(self) -> bool:
+        """
+        Enable ReflexFlow automatically when scheduled sampling is active on flow-matching models
+        and the user did not explicitly set the flag.
+        """
+        try:
+            offset_value = getattr(self.config, "scheduled_sampling_max_step_offset", 0)
+            max_offset = float(offset_value or 0)
+        except Exception:
+            return False
+
+        if max_offset <= 0:
+            return False
+
+        if getattr(self.config, "scheduled_sampling_reflexflow", None) is not None:
+            return False
+
+        if self.PREDICTION_TYPE is not PredictionTypes.FLOW_MATCHING:
+            return False
+
+        setattr(self.config, "scheduled_sampling_reflexflow", True)
+        return True
+
     def prepare_batch(self, batch: dict, state: dict) -> dict:
         """
         Moves the batch to the proper device/dtype,
@@ -2548,6 +2572,7 @@ def prepare_batch(self, batch: dict, state: dict) -> dict:
                 batch["timesteps"],
             ).to(device=self.accelerator.device, dtype=self.config.weight_dtype)
 
+        self._maybe_enable_reflexflow_default()
         if getattr(self.config, "scheduled_sampling_max_step_offset", 0) > 0:
             effective_prob = float(getattr(self.config, "scheduled_sampling_probability", 0.0) or 0.0)
             prob_start = float(getattr(self.config, "scheduled_sampling_prob_start", effective_prob) or effective_prob)

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

@@ -319,7 +319,7 @@ def register_advanced_fields(registry: "FieldRegistry") -> None:
             tab="training",
             section="loss_functions",
             subsection="advanced",
-            default_value=False,
+            default_value=None,
             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,

tests/test_scheduled_sampling_rollout.py

@@ -202,6 +202,56 @@ def test_flow_matching_rollout_updates_sigmas_and_caches(self):
         assert "_reflexflow_biased_pred" in updated
 
 
+class ReflexFlowDefaultToggleTests(unittest.TestCase):
+    def _make_model(self, prediction_type, config):
+        class _StubModel(ModelFoundation):
+            PREDICTION_TYPE = prediction_type
+
+            def __init__(self, cfg):
+                self.config = cfg
+
+            def model_predict(self, prepared_batch, custom_timesteps: list | None = None):
+                return prepared_batch
+
+            def _encode_prompts(self, *args, **kwargs):
+                return None
+
+            def convert_text_embed_for_pipeline(self, text_encoder_output, pooling_encode_output=None):
+                return text_encoder_output
+
+            def convert_negative_text_embed_for_pipeline(self, text_encoder_output, pooling_encode_output=None):
+                return text_encoder_output
+
+        return _StubModel(config)
+
+    def test_auto_enables_when_unset_for_flow_matching(self):
+        config = SimpleNamespace(scheduled_sampling_max_step_offset=3, scheduled_sampling_reflexflow=None)
+        model = self._make_model(PredictionTypes.FLOW_MATCHING, config)
+
+        changed = model._maybe_enable_reflexflow_default()
+
+        self.assertTrue(changed)
+        self.assertTrue(config.scheduled_sampling_reflexflow)
+
+    def test_respects_user_opt_out(self):
+        config = SimpleNamespace(scheduled_sampling_max_step_offset=3, scheduled_sampling_reflexflow=False)
+        model = self._make_model(PredictionTypes.FLOW_MATCHING, config)
+
+        changed = model._maybe_enable_reflexflow_default()
+
+        self.assertFalse(changed)
+        self.assertFalse(config.scheduled_sampling_reflexflow)
+
+    def test_skips_non_flow_matching_models(self):
+        config = SimpleNamespace(scheduled_sampling_max_step_offset=3, scheduled_sampling_reflexflow=None)
+        model = self._make_model(PredictionTypes.EPSILON, config)
+
+        changed = model._maybe_enable_reflexflow_default()
+
+        self.assertFalse(changed)
+        self.assertIsNone(config.scheduled_sampling_reflexflow)
+
+
 class _FlowLossModel(ModelFoundation):
     PREDICTION_TYPE = PredictionTypes.FLOW_MATCHING
     NAME = "dummy"