maximeseince commited on
Commit
70edc47
·
verified ·
1 Parent(s): 86ec3e5

Document the conditioning-tutorial checkpoints

Browse files
Files changed (1) hide show
  1. README.md +16 -3
README.md CHANGED
@@ -12,7 +12,7 @@ tags:
12
 
13
  Small checkpoints used by the [`stip`](https://github.com/instadeepai/stip) tutorial
14
  notebooks, so that a tutorial can demonstrate sampling without spending ten
15
- minutes training first. They are toy models (a two-layer MLP, ~50k parameters,
16
  trained for 3000 steps on a 4-component 2D Gaussian mixture) and have no
17
  use outside the notebooks.
18
 
@@ -23,13 +23,13 @@ Checkpoints are [Orbax](https://orbax.readthedocs.io) directories written by
23
  ## `conditioning_and_guidance/`
24
 
25
  Used by `tutorials/notebooks/4.conditioning_and_guidance.ipynb`. Both models are
26
- `VelocityGenerativeModel`s with a `FlowMatchingOneSidedInterpolant`, but over
27
  different modalities:
28
 
29
  | Path | Model | Modalities | Role in the notebook |
30
  |---|---|---|---|
31
  | `conditioning_and_guidance/joint_model` | Unconditional cross-modal MLP | `coordinates` (continuous, 2D) and `index` (discrete, 4 categories) | Intrinsic guidance (Section 3): conditioning a model that was never trained to be conditional |
32
- | `conditioning_and_guidance/context_model` | The same MLP plus a label context path, trained with 50% context dropout | `coordinates` only; the corner label is passed as `context_data` instead of as a modality | Context conditioning and classifier-free guidance (Sections 4) |
33
 
34
  ### Loading
35
 
@@ -51,3 +51,16 @@ checkpointer = Checkpointer(
51
  )
52
  gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
53
  ```
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12
 
13
  Small checkpoints used by the [`stip`](https://github.com/instadeepai/stip) tutorial
14
  notebooks, so that a tutorial can demonstrate sampling without spending ten
15
+ minutes training first. They are toy models (a two-layer MLP, 18k-25k parameters,
16
  trained for 3000 steps on a 4-component 2D Gaussian mixture) and have no
17
  use outside the notebooks.
18
 
 
23
  ## `conditioning_and_guidance/`
24
 
25
  Used by `tutorials/notebooks/4.conditioning_and_guidance.ipynb`. Both models are
26
+ `VelocityOneSidedGenerativeModel`s with a `FlowMatchingOneSidedInterpolant`, but over
27
  different modalities:
28
 
29
  | Path | Model | Modalities | Role in the notebook |
30
  |---|---|---|---|
31
  | `conditioning_and_guidance/joint_model` | Unconditional cross-modal MLP | `coordinates` (continuous, 2D) and `index` (discrete, 4 categories) | Intrinsic guidance (Section 3): conditioning a model that was never trained to be conditional |
32
+ | `conditioning_and_guidance/context_model` | The same MLP plus a label context path, trained with 50% context dropout | `coordinates` only; the corner label is passed as `context_data` instead of as a modality | Context conditioning and classifier-free guidance (Sections 4-5) |
33
 
34
  ### Loading
35
 
 
51
  )
52
  gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
53
  ```
54
+
55
+ `restore_ema` reads only `ema_params` and `extra`, and applies the same bias
56
+ correction the training loop uses for evaluation.
57
+
58
+ ## Reproducing
59
+
60
+ ```bash
61
+ uv run python tutorials/scripts/train_conditioning_checkpoints.py
62
+ ```
63
+
64
+ The script mirrors the notebook's model definitions and PRNG chain, so it
65
+ reproduces these exact weights. A checkpoint pins the parameter structure: if a
66
+ notebook's network changes, re-run the script and re-upload.