Add timestep distribution visualization and loss weighting functionality (#2384)

* feat(timestep-visualization): add functionality to visualize sampled timestep distribution and loss weighting

* Potential fix for pull request finding: fix FLUX weighting handling

Co-authored-by: Copilot Autofix powered by AI <[email protected]>

* Potential fix for pull request finding: reduce batch size

Co-authored-by: Copilot Autofix powered by AI <[email protected]>

---------

Co-authored-by: Copilot Autofix powered by AI <[email protected]>

anima_train.py

@@ -53,6 +53,8 @@ def train(args):
     deepspeed_utils.prepare_deepspeed_args(args)
     setup_logging(args, reset=True)
 
+    flux_train_utils.log_timestep_sampling_info(args)
+
     # backward compatibility
     if not args.skip_cache_check:
         args.skip_cache_check = args.skip_latents_validity_check
@@ -763,4 +765,7 @@ if __name__ == "__main__":
     if args.attn_mode == "sdpa":
         args.attn_mode = "torch"  # backward compatibility
 
-    train(args)
+    if args.show_timesteps:
+        anima_train_utils.show_timesteps(args)
+    else:
+        train(args)

anima_train_control_net_lllite.py

@@ -193,6 +193,8 @@ def train(args):
     deepspeed_utils.prepare_deepspeed_args(args)
     setup_logging(args, reset=True)
 
+    flux_train_utils.log_timestep_sampling_info(args)
+
     if not args.skip_cache_check:
         args.skip_cache_check = args.skip_latents_validity_check
 
@@ -817,4 +819,7 @@ if __name__ == "__main__":
     if args.attn_mode == "sdpa":
         args.attn_mode = "torch"
 
-    train(args)
+    if args.show_timesteps:
+        anima_train_utils.show_timesteps(args)
+    else:
+        train(args)

anima_train_network.py

@@ -44,6 +44,8 @@ class AnimaNetworkTrainer(train_network.NetworkTrainer):
         train_dataset_group: Union[DatasetGroup, MinimalDataset],
         val_dataset_group: Optional[DatasetGroup],
     ):
+        flux_train_utils.log_timestep_sampling_info(args)
+
         if args.fp8_base or args.fp8_base_unet:
             logger.warning("fp8_base and fp8_base_unet are not supported. / fp8_baseとfp8_base_unetはサポートされていません。")
             args.fp8_base = False
@@ -450,5 +452,8 @@ if __name__ == "__main__":
     if args.attn_mode == "sdpa":
         args.attn_mode = "torch"  # backward compatibility
 
-    trainer = AnimaNetworkTrainer()
-    trainer.train(args)
+    if args.show_timesteps:
+        anima_train_utils.show_timesteps(args)
+    else:
+        trainer = AnimaNetworkTrainer()
+        trainer.train(args)

docs/anima_train_network.md

@@ -500,6 +500,19 @@ The `--weighting_scheme` option specifies loss weighting by timestep:
 - `none`: Same as uniform.
 - `logit_normal`, `mode`: Additional schemes from SD3 training. See the [`sd3_train_network.md` guide](sd3_train_network.md) for details.
 
+#### Visualizing the Timestep Distribution
+
+To check how the above settings actually affect the sampled timesteps, two features are available (both shared with FLUX training):
+
+* **Training log:** At the start of every run, a one-line summary of the timestep sampling configuration is logged. It explicitly states whether `--discrete_flow_shift` is applied for the chosen `--timestep_sampling` (only `sigma` and `shift` use it; with `sigmoid`, `uniform`, or `flux_shift` it is **ignored**). This makes it easy to notice when a shift value you set has no effect.
+* `--show_timesteps=<console|image>`: Visualize the *actual* sampled-timestep distribution and the loss weighting for the current settings, then exit without training. `console` prints an ASCII histogram; `image` shows a matplotlib plot (requires `matplotlib`). The distribution is shown from noisy (top / left, t=1000) to clean (bottom / right, t=0).
+* `--show_timesteps_resolution=<H | H,W>`: Image resolution (in pixels) assumed by `--show_timesteps` for resolution-dependent sampling such as `flux_shift`. A single value is used for both height and width; two comma-separated values are `H,W` (`W,H` gives the same result). Default `1024`. This is useful because the `--resolution` argument is `None` when the dataset is configured via a `.toml` file.
+
+Example:
+```
+--show_timesteps console --show_timesteps_resolution 1024
+```
+
 #### Caption Dropout
 
 Caption dropout uses the `caption_dropout_rate` setting from the dataset configuration (per-subset in TOML). When using `--cache_text_encoder_outputs`, the dropout rate is stored with each cached entry and applied during training, so caption dropout is compatible with text encoder output caching.
@@ -531,6 +544,14 @@ Note: Currently, only Anima supports combining `caption_dropout_rate` with text
 
 `--weighting_scheme`でタイムステップごとの損失の重み付けを指定します。
 
+#### タイムステップ分布の可視化
+
+上記の設定が実際のタイムステップにどう影響するかを確認するため、2つの機能があります（FLUX学習と共通）。
+
+* **学習ログ:** 学習開始時に、タイムステップサンプリング設定の概要を1行ログ出力します。選択した`--timestep_sampling`に対して`--discrete_flow_shift`が適用されるか（使用するのは`sigma`と`shift`のみ。`sigmoid`、`uniform`、`flux_shift`では**無視**される）を明示するため、設定したシフト値が効いていない場合に気づきやすくなります。
+* `--show_timesteps=<console|image>`: 現在の設定で実際にサンプリングされるタイムステップ分布とloss weightingを可視化して終了します（学習は行いません）。`console`はASCIIヒストグラム、`image`はmatplotlibで表示します（`matplotlib`が必要）。分布はノイズ側（上／左、t=1000）からクリーン側（下／右、t=0）の順で表示されます。
+* `--show_timesteps_resolution=<H | H,W>`: `flux_shift`等の解像度依存サンプリングで`--show_timesteps`が想定する画像解像度（ピクセル）。数値が1つなら縦横両方に、カンマ区切りで2つなら`H,W`に使用します（`W,H`でも結果は同じ）。デフォルト`1024`。`.toml`データセット設定を使用すると`--resolution`が`None`になるため、この指定が役立ちます。
+
 #### キャプションドロップアウト
 
 キャプションドロップアウトにはデータセット設定（TOMLでのサブセット単位）の`caption_dropout_rate`を使用します。`--cache_text_encoder_outputs`使用時は、ドロップアウト率が各キャッシュエントリとともに保存され、学習中に適用されるため、テキストエンコーダー出力キャッシュと同時に使用できます。

docs/flux_train_network.md

@@ -337,6 +337,10 @@ The `--timestep_sampling` option specifies how timesteps (0-1) are sampled:
 
 `--discrete_flow_shift` only applies when `--timestep_sampling` is set to `shift`.
 
+#### Visualizing the Timestep Distribution
+
+You can visualize the actual sampled-timestep distribution and loss weighting for your current settings — and confirm whether `--discrete_flow_shift` is actually applied — using the `--show_timesteps` and `--show_timesteps_resolution` options. A one-line summary of the configuration is also logged at the start of every training run. These options are shared between FLUX and Anima training; see the [Anima guide](anima_train_network.md#visualizing-the-timestep-distribution) for details.
+
 #### Model Prediction Processing
 
 The `--model_prediction_type` option specifies how to interpret and process model predictions:
@@ -369,6 +373,8 @@ FLUX.1の学習には多くの未知の点があり、いくつかの設定は
 - モデル予測の処理方法（`--model_prediction_type`）
 - 推奨設定の組み合わせ
 
+タイムステップ分布の可視化（`--show_timesteps` / `--show_timesteps_resolution`）や、`--discrete_flow_shift`が適用されるかの確認方法については、FLUXとAnimaで共通のため[Animaのガイド](anima_train_network.md#タイムステップ分布の可視化)を参照してください。
+
 </details>
 
 ### 6.3. Layer-specific Rank Configuration / 各層に対するランク指定

flux_train.py

@@ -66,6 +66,8 @@ def train(args):
     deepspeed_utils.prepare_deepspeed_args(args)
     setup_logging(args, reset=True)
 
+    flux_train_utils.log_timestep_sampling_info(args)
+
     # temporary: backward compatibility for deprecated options. remove in the future
     if not args.skip_cache_check:
         args.skip_cache_check = args.skip_latents_validity_check
@@ -855,4 +857,7 @@ if __name__ == "__main__":
     args_util.verify_command_line_training_args(args)
     args = args_util.read_config_from_file(args, parser)
 
-    train(args)
+    if args.show_timesteps:
+        flux_train_utils.show_timesteps(args)
+    else:
+        train(args)

flux_train_control_net.py

@@ -72,6 +72,8 @@ def train(args):
     deepspeed_utils.prepare_deepspeed_args(args)
     setup_logging(args, reset=True)
 
+    flux_train_utils.log_timestep_sampling_info(args)
+
     # temporary: backward compatibility for deprecated options. remove in the future
     if not args.skip_cache_check:
         args.skip_cache_check = args.skip_latents_validity_check
@@ -889,4 +891,7 @@ if __name__ == "__main__":
     args_util.verify_command_line_training_args(args)
     args = args_util.read_config_from_file(args, parser)
 
-    train(args)
+    if args.show_timesteps:
+        flux_train_utils.show_timesteps(args)
+    else:
+        train(args)

flux_train_network.py

@@ -49,6 +49,8 @@ class FluxNetworkTrainer(train_network.NetworkTrainer):
         super().assert_extra_args(args, train_dataset_group, val_dataset_group)
         # sdxl_train_util.verify_sdxl_training_args(args)
 
+        flux_train_utils.log_timestep_sampling_info(args)
+
         self.model_type = args.model_type  # "flux" or "chroma"
         if self.model_type != "chroma":
             self.use_clip_l = True
@@ -546,5 +548,8 @@ if __name__ == "__main__":
     args_util.verify_command_line_training_args(args)
     args = args_util.read_config_from_file(args, parser)
 
-    trainer = FluxNetworkTrainer()
-    trainer.train(args)
+    if args.show_timesteps:
+        flux_train_utils.show_timesteps(args)
+    else:
+        trainer = FluxNetworkTrainer()
+        trainer.train(args)

library/anima_train_utils.py

@@ -7,6 +7,7 @@ import os
 import time
 from typing import Optional
 
+from library import device_utils
 import numpy as np
 import torch
 from accelerate import Accelerator
@@ -158,6 +159,40 @@ def compute_loss_weighting_for_anima(weighting_scheme: str, sigmas: torch.Tensor
     return weighting
 
 
+def show_timesteps(args):
+    """Visualize the actual sampled-timestep / loss-weighting distribution for the current Anima settings, then return.
+
+    Anima reuses ``flux_train_utils.get_noisy_model_input_and_timesteps`` for sampling but has its own loss weighting.
+    """
+    from library import flux_train_utils, sd3_train_utils, timestep_visualization
+
+    num_train_timesteps = 1000
+    noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=num_train_timesteps, shift=args.discrete_flow_shift
+    )
+    h, w = flux_train_utils.parse_show_timesteps_latent_size(args)  # latent size for the assumed image resolution
+    device, dtype = device_utils.get_preferred_device(), torch.float32
+
+    def sample_timesteps(bsz):
+        latents = torch.zeros(bsz, 16, h, w, dtype=dtype, device=device)
+        noise = torch.ones_like(latents)
+        _, timesteps, _ = flux_train_utils.get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype)
+        return timesteps
+
+    def compute_weighting(timesteps):
+        sigmas = timesteps / num_train_timesteps
+        return compute_loss_weighting_for_anima(args.weighting_scheme, sigmas)
+
+    header = (
+        "Timestep distribution / タイムステップ分布:\n  "
+        + flux_train_utils.get_timestep_sampling_info(args)
+        + f", resolution={args.show_timesteps_resolution} (latent {h}x{w})"
+    )
+    timestep_visualization.show_timestep_distribution(
+        args.show_timesteps, sample_timesteps, compute_weighting, num_train_timesteps=num_train_timesteps, header=header
+    )
+
+
 # Parameter groups (6 groups with separate LRs)
 def get_anima_param_groups(
     dit,

library/args.py

@@ -686,6 +686,26 @@ def add_dit_training_arguments(parser: argparse.ArgumentParser):
         default=1.29,
         help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme` / モード重み付けスキームのスケール",
     )
+    parser.add_argument(
+        "--show_timesteps",
+        type=str,
+        default=None,
+        choices=["console", "image"],
+        help="visualize the actual sampled timestep distribution and loss weighting for the current settings, then exit. "
+        "'console' prints an ASCII histogram, 'image' shows a matplotlib plot. "
+        "Resolution-dependent sampling (e.g. flux_shift) assumes ~1024px. "
+        " / 現在の設定で実際にサンプリングされるタイムステップ分布とloss weightingを可視化して終了する。"
+        "'console'はASCIIヒストグラム、'image'はmatplotlibで表示。flux_shift等の解像度依存サンプリングは約1024px想定。",
+    )
+    parser.add_argument(
+        "--show_timesteps_resolution",
+        type=str,
+        default="1024",
+        help="image resolution (pixels) assumed by --show_timesteps for resolution-dependent sampling (e.g. flux_shift). "
+        "Comma-separated: a single value is used for both H and W, two values are H,W (W,H gives the same result). Default: 1024. "
+        " / --show_timesteps が解像度依存サンプリング（flux_shift等）で想定する画像解像度（ピクセル）。"
+        "カンマ区切りで、数値が1つならH/W両方に、2つならH,Wに使用（W,Hでも結果は同じ）。デフォルト: 1024。",
+    )
 
     # offloading
     parser.add_argument(

library/flux_train_utils.py

@@ -1,6 +1,7 @@
 import argparse
 import math
 import os
+from library import device_utils
 import numpy as np
 import toml
 import json
@@ -530,6 +531,94 @@ def get_noisy_model_input_and_timesteps(
     return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas
 
 
+# timestep_sampling values whose distribution actually depends on --discrete_flow_shift.
+# "shift" uses it explicitly; "sigma" uses it via the shifted scheduler.timesteps. The others
+# ("uniform", "sigmoid", "flux_shift") ignore discrete_flow_shift entirely. This is shared by both
+# FLUX and Anima since Anima reuses get_noisy_model_input_and_timesteps below.
+_SHIFT_AWARE_TIMESTEP_SAMPLING = ("sigma", "shift")
+
+
+def get_timestep_sampling_info(args) -> str:
+    """One-line, human-readable summary of the timestep sampling config.
+
+    Makes it explicit whether ``--discrete_flow_shift`` actually affects the distribution for the chosen
+    ``--timestep_sampling`` (the core confusion behind issue #2383).
+    """
+    sampling = args.timestep_sampling
+    parts = [f"timestep_sampling={sampling}"]
+    if sampling in _SHIFT_AWARE_TIMESTEP_SAMPLING:
+        parts.append(f"discrete_flow_shift={args.discrete_flow_shift} (applied)")
+    else:
+        parts.append(
+            f"discrete_flow_shift={args.discrete_flow_shift} (IGNORED for timestep_sampling='{sampling}'; "
+            "only 'sigma' and 'shift' use it)"
+        )
+    if sampling in ("sigmoid", "shift", "flux_shift"):
+        parts.append(f"sigmoid_scale={args.sigmoid_scale}")
+    if sampling == "sigma":
+        parts.append(f"weighting_scheme={args.weighting_scheme}")
+    return ", ".join(parts)
+
+
+def log_timestep_sampling_info(args):
+    """Log the timestep sampling config once at training start (shared by FLUX and Anima)."""
+    logger.info(f"Timestep sampling: {get_timestep_sampling_info(args)}")
+
+
+def parse_show_timesteps_latent_size(args, vae_compression: int = 8) -> Tuple[int, int]:
+    """Parse ``--show_timesteps_resolution`` (``H`` or ``H,W`` pixels) into latent (h, w) dims.
+
+    A single value is used for both H and W; two values are H,W (W,H gives the same result since
+    flux_shift only depends on the token count). The VAE spatial compression is 8 for both FLUX and Anima.
+    """
+    res = getattr(args, "show_timesteps_resolution", None) or "1024"
+    vals = [int(v.strip()) for v in str(res).split(",") if v.strip() != ""]
+    if len(vals) == 1:
+        height = width = vals[0]
+    elif len(vals) == 2:
+        height, width = vals[0], vals[1]
+    else:
+        raise ValueError(f"--show_timesteps_resolution must be 'H' or 'H,W', got: {res}")
+    return height // vae_compression, width // vae_compression
+
+
+def show_timesteps(args):
+    """Visualize the actual sampled-timestep / loss-weighting distribution for the current FLUX settings, then return.
+
+    Builds a weight-free scheduler and repeatedly calls ``get_noisy_model_input_and_timesteps`` with dummy
+    latents so the distribution reflects ``--timestep_sampling`` / ``--discrete_flow_shift`` / ``--weighting_scheme`` etc.
+    """
+    from library import sd3_train_utils, timestep_visualization
+
+    num_train_timesteps = 1000
+    noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=num_train_timesteps, shift=args.discrete_flow_shift
+    )
+    # latent size for the assumed image resolution (flux_shift reads h, w for the packed-size shift)
+    h, w = parse_show_timesteps_latent_size(args)
+    device, dtype = device_utils.get_preferred_device(), torch.float32
+
+    def sample_timesteps(bsz):
+        latents = torch.zeros(bsz, 16, h, w, dtype=dtype, device=device)
+        noise = torch.ones_like(latents)
+        _, timesteps, _ = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype)
+        return timesteps
+
+    def compute_weighting(timesteps):
+        if args.model_prediction_type != "sigma_scaled":
+            return None
+        sigmas = timesteps / num_train_timesteps
+        return compute_loss_weighting_for_sd3(args.weighting_scheme, sigmas)
+    header = (
+        "Timestep distribution / タイムステップ分布:\n  "
+        + get_timestep_sampling_info(args)
+        + f", resolution={args.show_timesteps_resolution} (latent {h}x{w}), model_prediction_type={args.model_prediction_type}"
+    )
+    timestep_visualization.show_timestep_distribution(
+        args.show_timesteps, sample_timesteps, compute_weighting, num_train_timesteps=num_train_timesteps, header=header
+    )
+
+
 def apply_model_prediction_type(args, model_pred, noisy_model_input, sigmas):
     weighting = None
     if args.model_prediction_type == "raw":

library/timestep_visualization.py

@@ -0,0 +1,126 @@
+"""Visualize the actual sampled-timestep / loss-weighting distribution for Flow Matching training.
+
+This module is a *leaf* utility: it only knows how to render a distribution given
+sampling and weighting callables. Model-specific assembly (which scheduler to build,
+which ``get_noisy_model_input_and_timesteps`` to call) lives in the per-model
+``*_train_utils`` modules so this file has no model dependencies.
+
+Typical use from a training script::
+
+    if args.show_timesteps:
+        flux_train_utils.show_timesteps(args)  # builds the callables and calls show_timestep_distribution
+    else:
+        trainer.train(args)
+"""
+
+import logging
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+CONSOLE_WIDTH = 64
+N_TIMESTEPS_PER_LINE = 25
+
+
+def show_timestep_distribution(
+    mode: str,
+    sample_timesteps: Callable[[int], torch.Tensor],
+    compute_weighting: Optional[Callable[[torch.Tensor], Optional[torch.Tensor]]] = None,
+    num_train_timesteps: int = 1000,
+    n_try: int = 100000,
+    batch_size: int = 64,
+    header: Optional[str] = None,
+):
+    """Sample timesteps many times and render their distribution (and loss weighting).
+
+    Args:
+        mode: ``"console"`` for an ASCII histogram, ``"image"`` for a matplotlib plot.
+        sample_timesteps: ``fn(batch_size) -> Tensor`` returning sampled timesteps in
+            ``[0, num_train_timesteps]`` (any shape; it is flattened).
+        compute_weighting: ``fn(timesteps) -> Tensor | None`` returning the loss
+            weighting for a 1-D tensor of integer timesteps ``[1 .. num_train_timesteps]``.
+            ``None`` (or an all-``None`` return) is treated as uniform weighting of 1.0.
+        num_train_timesteps: scheduler resolution (number of discrete timesteps).
+        n_try: total number of timesteps to sample.
+        batch_size: per-iteration sample count.
+        header: optional multi-line text printed before the histogram (e.g. the args used).
+    """
+    if header:
+        print(header)
+
+    # --- sample the timestep distribution ---
+    sampled_timesteps = np.zeros(num_train_timesteps, dtype=np.float64)
+    for _ in tqdm(range(max(1, n_try // batch_size)), desc="sampling timesteps"):
+        timesteps = sample_timesteps(batch_size).detach().flatten().float().cpu().numpy()
+        indices = np.clip(timesteps.astype(np.int64), 0, num_train_timesteps - 1)
+        np.add.at(sampled_timesteps, indices, 1.0)
+
+    # --- compute the loss weighting per timestep ---
+    sampled_weighting = np.ones(num_train_timesteps, dtype=np.float64)
+    if compute_weighting is not None:
+        all_timesteps = torch.arange(1, num_train_timesteps + 1, dtype=torch.float32)
+        weighting = compute_weighting(all_timesteps)
+        if weighting is not None:
+            weighting = weighting.detach().float()
+            # guard against inf / nan from schemes like sigma_sqrt at sigma==0
+            weighting = torch.nan_to_num(weighting, nan=1.0, posinf=1.0, neginf=1.0)
+            sampled_weighting = weighting.cpu().numpy().reshape(-1)
+
+    if mode == "image":
+        _show_image(sampled_timesteps, sampled_weighting)
+    else:
+        _show_console(sampled_timesteps, sampled_weighting, num_train_timesteps)
+
+
+def _show_console(sampled_timesteps: np.ndarray, sampled_weighting: np.ndarray, num_train_timesteps: int):
+    # average per line so the histogram fits the console
+    n_lines = num_train_timesteps // N_TIMESTEPS_PER_LINE
+    usable = n_lines * N_TIMESTEPS_PER_LINE
+    ts = sampled_timesteps[:usable].reshape(-1, N_TIMESTEPS_PER_LINE).mean(axis=1)
+    wt = sampled_weighting[:usable].reshape(-1, N_TIMESTEPS_PER_LINE).mean(axis=1)
+
+    # display high -> low timestep (noisiest first), since t=num_train_timesteps is pure noise and t=0 is clean
+    max_count = max(ts.max(), 1e-8)
+    print(f"Sampled timesteps (top=noisy {num_train_timesteps}, bottom=clean 0): max count={max_count:.1f}")
+    for i in reversed(range(len(ts))):
+        line = f"{i * N_TIMESTEPS_PER_LINE:4d}-{(i + 1) * N_TIMESTEPS_PER_LINE - 1:4d}: "
+        line += "#" * int(ts[i] / max_count * CONSOLE_WIDTH)
+        print(line)
+
+    max_weighting = max(wt.max(), 1e-8)
+    print(f"Sampled loss weighting (top=noisy {num_train_timesteps}, bottom=clean 0): max weighting={max_weighting:.2f}")
+    for i in reversed(range(len(wt))):
+        line = f"{i * N_TIMESTEPS_PER_LINE:4d}-{(i + 1) * N_TIMESTEPS_PER_LINE - 1:4d}: {wt[i]:8.2f} "
+        line += "#" * int(wt[i] / max_weighting * CONSOLE_WIDTH)
+        print(line)
+
+
+def _show_image(sampled_timesteps: np.ndarray, sampled_weighting: np.ndarray):
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        logger.error("matplotlib is not installed; falling back to console output. / matplotlibが見つかりません。コンソール出力にフォールバックします。")
+        _show_console(sampled_timesteps, sampled_weighting, len(sampled_timesteps))
+        return
+
+    plt.figure(figsize=(10, 5))
+    plt.subplot(1, 2, 1)
+    plt.bar(range(len(sampled_timesteps)), sampled_timesteps, width=1.0)
+    plt.title("Sampled timesteps")
+    plt.xlabel("Timestep (left=noisy, right=clean)")
+    plt.ylabel("Count")
+    plt.gca().invert_xaxis()  # noisiest timestep on the left
+
+    plt.subplot(1, 2, 2)
+    plt.bar(range(len(sampled_weighting)), sampled_weighting, width=1.0)
+    plt.title("Sampled loss weighting")
+    plt.xlabel("Timestep (left=noisy, right=clean)")
+    plt.ylabel("Weighting")
+    plt.gca().invert_xaxis()  # noisiest timestep on the left
+
+    plt.tight_layout()
+    plt.show()