refactor: move timestep_sampling_offset to custom_attributes dict

Address reviewer feedback (kohya-ss):
1. Remove timestep_sampling_offset from subset public schema (BaseSubset,
   BaseSubsetParams, ConfigSanitizer). Use custom_attributes instead,
   following the diff_output_preservation pattern.
2. Use nested dict format: timestep_sampling = { offset = -0.5 }
   to allow future expansion (scale, mode, etc.).
3. Skip offset during validation — only apply during training
   for comparable loss metrics.
4. Add regression tests for offset=None/0.0 baseline match and
   positive offset shifting timestep mean upward.
5. Update documentation with new TOML syntax and scope notes.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

anima_train_network.py

@@ -286,11 +286,17 @@ class AnimaNetworkTrainer(train_network.NetworkTrainer):
             latents = latents.squeeze(2)  # [B, C, 1, H, W] -> [B, C, H, W]
         noise = torch.randn_like(latents)
 
-        # Per-sample timestep sampling offset from dataset subset config (timestep_sampling_offset in TOML)
-        tso = batch.get("timestep_sampling_offsets")
+        # Per-sample timestep sampling offset from custom_attributes: timestep_sampling = { offset = ... }
+        # Only applied during training; validation uses unbiased sampling for comparable loss.
+        tso = None
+        if is_train and "custom_attributes" in batch:
+            offsets = [ca.get("timestep_sampling", {}).get("offset", 0.0) for ca in batch["custom_attributes"]]
+            t = torch.tensor(offsets, dtype=torch.float32)
+            if t.abs().sum() > 0:
+                tso = t
         noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timesteps(
             args, noise_scheduler, latents, noise, accelerator.device, weight_dtype,
-            timestep_sampling_offset=tso if tso is not None and tso.abs().sum() > 0 else None,
+            timestep_sampling_offset=tso,
         )
         timesteps = timesteps / 1000.0  # scale to [0, 1] range. timesteps is float32
 

docs/timestep_sampling_offset.md

@@ -1,10 +1,10 @@
-# Per-subset timestep sampling offset (`timestep_sampling_offset`)
+# Per-subset timestep sampling offset (`custom_attributes.timestep_sampling.offset`)
 
 ## Overview
 
-For flow-matching models, the noise level (timestep) at which each sample is trained materially affects the final result. `timestep_sampling_offset` shifts the flow-matching timestep sampling distribution per dataset subset toward higher- or lower-noise regions.
+For flow-matching models, the noise level (timestep) at which each sample is trained materially affects the final result. The `timestep_sampling.offset` custom attribute shifts the flow-matching timestep sampling distribution per dataset subset toward higher- or lower-noise regions.
 
-It is applied as an offset to the pre-sigmoid normal sample of the logit-normal (sigmoid / shift / flux_shift) schedule. Default `0.0` leaves sampling unchanged.
+It is applied as an offset to the pre-sigmoid normal sample of the logit-normal (sigmoid / shift / flux_shift) schedule. Default `0.0` (or omitted) leaves sampling unchanged.
 
 - **Negative** value → biases toward lower-noise timesteps (detail-focused steps)
 - **Positive** value → biases toward higher-noise timesteps (structure-focused steps)
@@ -28,20 +28,22 @@ SGA analyzes flow-matching fine-tuning as a quadratic form governed by a Neural
 
 ## Usage
 
-In your dataset TOML config, set `timestep_sampling_offset` per subset:
+This feature uses `custom_attributes` to avoid expanding the subset public schema. Set it per subset in your dataset TOML config:
 
 ```toml
 [[datasets.subsets]]
 image_dir = "closeup_shots"
-timestep_sampling_offset = -0.5   # lower noise → detail refinement
+[datasets.subsets.custom_attributes]
+timestep_sampling = { offset = -0.5 }
 
 [[datasets.subsets]]
 image_dir = "full_body_shots"
-timestep_sampling_offset = 0.5    # higher noise → structure learning
+[datasets.subsets.custom_attributes]
+timestep_sampling = { offset = 0.5 }
 
 [[datasets.subsets]]
 image_dir = "other"
-# default 0.0 = unchanged
+# no custom_attributes needed — default is no offset
 ```
 
 Useful magnitudes are typically in the range of `-1.0` to `1.0`. The offset is applied before the sigmoid transform, so the effect on the final sigma distribution is nonlinear and saturates at large values.
@@ -50,4 +52,5 @@ Useful magnitudes are typically in the range of `-1.0` to `1.0`. The offset is a
 
 - Applies to `sigmoid`, `shift`, and `flux_shift` timestep sampling modes.
 - `uniform` and `sigma` (density-based) modes are not affected.
-- Currently wired in `anima_train_network.py`; other trainers can opt in with a one-line change.
+- Currently consumed by `anima_train_network.py`. Other trainers (FLUX, SD3, etc.) do not read this attribute; setting it for those trainers is a no-op.
+- The offset is applied only during **training**. Validation uses unbiased sampling for comparable loss metrics.

library/config_util.py

@@ -74,7 +74,6 @@ class BaseSubsetParams:
     token_warmup_min: int = 1
     token_warmup_step: float = 0
     custom_attributes: Optional[Dict[str, Any]] = None
-    timestep_sampling_offset: float = 0.0
     validation_seed: int = 0
     validation_split: float = 0.0
     resize_interpolation: Optional[str] = None
@@ -201,7 +200,6 @@ class ConfigSanitizer:
         "caption_prefix": str,
         "caption_suffix": str,
         "custom_attributes": dict,
-        "timestep_sampling_offset": Any(float, int),
         "resize_interpolation": str,
     }
     # DO means DropOut

library/dataset.py

@@ -1003,7 +1003,6 @@ class BaseDataset(torch.utils.data.Dataset):
         flippeds = []  # 変数名が微妙
         text_encoder_outputs_list = []
         custom_attributes = []
-        timestep_sampling_offsets = []
         masks = []
         masked_images = []
 
@@ -1012,7 +1011,6 @@ class BaseDataset(torch.utils.data.Dataset):
             subset = self.image_to_subset[image_key]
 
             custom_attributes.append(subset.custom_attributes)
-            timestep_sampling_offsets.append(subset.timestep_sampling_offset)
 
             # in case of fine tuning, is_reg is always False
             loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
@@ -1238,7 +1236,6 @@ class BaseDataset(torch.utils.data.Dataset):
         example["flippeds"] = flippeds
 
         example["network_multipliers"] = torch.FloatTensor([self.network_multiplier] * len(captions))
-        example["timestep_sampling_offsets"] = torch.FloatTensor(timestep_sampling_offsets)
 
         if self.debug_dataset:
             example["image_keys"] = bucket[image_index : image_index + self.batch_size]

library/subset.py

@@ -33,7 +33,6 @@ class BaseSubset:
         token_warmup_min: int,
         token_warmup_step: Union[float, int],
         custom_attributes: Optional[Dict[str, Any]] = None,
-        timestep_sampling_offset: float = 0.0,
         validation_seed: Optional[int] = None,
         validation_split: Optional[float] = 0.0,
         resize_interpolation: Optional[str] = None,
@@ -61,7 +60,6 @@ class BaseSubset:
         self.token_warmup_step = token_warmup_step  # N（N<1ならN*max_train_steps）ステップ目でタグの数が最大になる
 
         self.custom_attributes = custom_attributes if custom_attributes is not None else {}
-        self.timestep_sampling_offset = timestep_sampling_offset
 
         self.img_count = 0
 
@@ -99,7 +97,6 @@ class DreamBoothSubset(BaseSubset):
         token_warmup_min,
         token_warmup_step,
         custom_attributes: Optional[Dict[str, Any]] = None,
-        timestep_sampling_offset: float = 0.0,
         validation_seed: Optional[int] = None,
         validation_split: Optional[float] = 0.0,
         resize_interpolation: Optional[str] = None,
@@ -128,7 +125,6 @@ class DreamBoothSubset(BaseSubset):
             token_warmup_min,
             token_warmup_step,
             custom_attributes=custom_attributes,
-            timestep_sampling_offset=timestep_sampling_offset,
             validation_seed=validation_seed,
             validation_split=validation_split,
             resize_interpolation=resize_interpolation,
@@ -172,7 +168,6 @@ class FineTuningSubset(BaseSubset):
         token_warmup_min,
         token_warmup_step,
         custom_attributes: Optional[Dict[str, Any]] = None,
-        timestep_sampling_offset: float = 0.0,
         validation_seed: Optional[int] = None,
         validation_split: Optional[float] = 0.0,
         resize_interpolation: Optional[str] = None,
@@ -201,7 +196,6 @@ class FineTuningSubset(BaseSubset):
             token_warmup_min,
             token_warmup_step,
             custom_attributes=custom_attributes,
-            timestep_sampling_offset=timestep_sampling_offset,
             validation_seed=validation_seed,
             validation_split=validation_split,
             resize_interpolation=resize_interpolation,
@@ -241,7 +235,6 @@ class ControlNetSubset(BaseSubset):
         token_warmup_min,
         token_warmup_step,
         custom_attributes: Optional[Dict[str, Any]] = None,
-        timestep_sampling_offset: float = 0.0,
         validation_seed: Optional[int] = None,
         validation_split: Optional[float] = 0.0,
         resize_interpolation: Optional[str] = None,
@@ -270,7 +263,6 @@ class ControlNetSubset(BaseSubset):
             token_warmup_min,
             token_warmup_step,
             custom_attributes=custom_attributes,
-            timestep_sampling_offset=timestep_sampling_offset,
             validation_seed=validation_seed,
             validation_split=validation_split,
             resize_interpolation=resize_interpolation,

tests/library/test_flux_train_utils.py

@@ -218,3 +218,115 @@ def test_different_timestep_count(args, device):
     assert timesteps.shape == (2,)
     # Check that timesteps are within the proper range
     assert torch.all(timesteps < 500)
+
+
+# Tests for timestep_sampling_offset
+class TestTimestepSamplingOffset:
+    """Regression tests for per-subset timestep sampling offset."""
+
+    @pytest.mark.parametrize("mode", ["sigmoid", "shift", "flux_shift"])
+    def test_none_offset_matches_baseline(self, args, noise_scheduler, device, mode):
+        """offset=None should produce identical results to no offset."""
+        args.timestep_sampling = mode
+        args.sigmoid_scale = 1.0
+        bsz = 64
+        latents = torch.randn(bsz, 4, 8, 8)
+        noise = torch.randn(bsz, 4, 8, 8)
+        dtype = torch.float32
+
+        torch.manual_seed(42)
+        _, ts_baseline, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype
+        )
+        torch.manual_seed(42)
+        _, ts_none, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype, timestep_sampling_offset=None
+        )
+        assert torch.allclose(ts_baseline, ts_none)
+
+    @pytest.mark.parametrize("mode", ["sigmoid", "shift", "flux_shift"])
+    def test_zero_offset_matches_baseline(self, args, noise_scheduler, device, mode):
+        """offset=0.0 should produce identical results to no offset."""
+        args.timestep_sampling = mode
+        args.sigmoid_scale = 1.0
+        bsz = 64
+        latents = torch.randn(bsz, 4, 8, 8)
+        noise = torch.randn(bsz, 4, 8, 8)
+        dtype = torch.float32
+
+        torch.manual_seed(42)
+        _, ts_baseline, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype
+        )
+        torch.manual_seed(42)
+        _, ts_zero, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype,
+            timestep_sampling_offset=torch.zeros(bsz)
+        )
+        assert torch.allclose(ts_baseline, ts_zero)
+
+    @pytest.mark.parametrize("mode", ["sigmoid", "shift", "flux_shift"])
+    def test_positive_offset_increases_mean_timestep(self, args, noise_scheduler, device, mode):
+        """Positive offset should shift timestep distribution upward (higher noise)."""
+        args.timestep_sampling = mode
+        args.sigmoid_scale = 1.0
+        bsz = 256
+        latents = torch.randn(bsz, 4, 8, 8)
+        noise = torch.randn(bsz, 4, 8, 8)
+        dtype = torch.float32
+
+        torch.manual_seed(42)
+        _, ts_baseline, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype
+        )
+        torch.manual_seed(42)
+        _, ts_offset, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype,
+            timestep_sampling_offset=torch.full((bsz,), 1.0)
+        )
+        assert ts_offset.mean() > ts_baseline.mean()
+
+    def test_offset_does_not_affect_uniform(self, args, noise_scheduler, device):
+        """Uniform sampling ignores offset (offset code path is not reached)."""
+        args.timestep_sampling = "uniform"
+        bsz = 64
+        latents = torch.randn(bsz, 4, 8, 8)
+        noise = torch.randn(bsz, 4, 8, 8)
+        dtype = torch.float32
+
+        torch.manual_seed(42)
+        _, ts_baseline, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype
+        )
+        torch.manual_seed(42)
+        _, ts_offset, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype,
+            timestep_sampling_offset=torch.full((bsz,), 1.0)
+        )
+        assert torch.allclose(ts_baseline, ts_offset)
+
+    def test_per_sample_offset_broadcasting(self, args, noise_scheduler, device):
+        """Different offsets per sample should produce different shifts."""
+        args.timestep_sampling = "sigmoid"
+        args.sigmoid_scale = 1.0
+        bsz = 4
+        latents = torch.randn(bsz, 4, 8, 8)
+        noise = torch.randn(bsz, 4, 8, 8)
+        dtype = torch.float32
+        offset = torch.tensor([-1.0, 0.0, 0.0, 1.0])
+
+        torch.manual_seed(42)
+        _, ts_baseline, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype
+        )
+        torch.manual_seed(42)
+        _, ts_offset, _ = get_noisy_model_input_and_timesteps(
+            args, noise_scheduler, latents, noise, device, dtype,
+            timestep_sampling_offset=offset
+        )
+        # Samples with offset=0 should match baseline
+        assert torch.allclose(ts_offset[1], ts_baseline[1])
+        assert torch.allclose(ts_offset[2], ts_baseline[2])
+        # Negative offset → lower timestep, positive → higher
+        assert ts_offset[0] < ts_baseline[0]
+        assert ts_offset[3] > ts_baseline[3]