📚 Tutorials¶

We provide a set of tutorials that cover the basics of the library so that one can quickly get started with a broad overview of its capabilities.

Each tutorial is a Jupyter notebook hosted in the repository; follow the Open the notebook link to view the rendered version on GitHub.

New to stix? We recommend reading the Introduction first — it maps the core mathematical objects onto the library’s API and will make the tutorials below easier to follow.

Training & sampling¶

Build, train, and sample from a stochastic-interpolant model end-to-end. A good first read to understand how the pieces of stix fit together.

What you’ll learn

  • How to define an interpolant and wrap it in a generative model.

  • How to configure and run the training loop.

  • How to draw samples with a solver and visualise the results.

Open the notebook →

Multimodal data loading¶

Feed your own data into stix with grain, a JAX-friendly, checkpointable data-loading library.

What you’ll learn

  • The minimum elements required to wire a data source into the stix pipeline.

  • How to plug in an external dataset and custom synthetic sources.

  • How to checkpoint and restore the data stream so an interrupted run resumes exactly where it stopped.

Open the notebook →

Building generative models¶

A detailed look at the GenerativeModel: how the choice of prediction target defines the loss, and how that same choice determines the sample-time conversions.

What you’ll learn

  • How to subclass GenerativeModel to build your own model, worked through three examples: a denoiser (predicts the clean data), a noise-prediction model, and an advanced hybrid model mixing a velocity target (continuous modalities) with a cross-entropy logits target (discrete modalities).

  • How the loss fixes what the network predicts, and how to convert that prediction into the sample-time generator via get_generator.

  • How to reuse the pre-baked one-sided-linear interpolant converters.

Open the notebook →

Conditioning and guidance¶

Learn how to use stix’s conditioning and guidance mechanisms to steer the generative model at sampling time. Understand both the code and theory behind the two conditioning channels and how to write your own guidance recipe.

What you’ll learn

  • The difference between the two conditioning channels : context (condition the network output based on provided context, requires training with context) and intrinsic (no specific training required, use standard data modality to build surrogate conitional score).

  • How to leverage intrinsic conditioning methods to condition on standard learned data modalities, including inpainting.

  • How to restore a context-conditional model and amplify it with classifier-free guidance.

  • How to write a custom GuidanceFn. In the notebook we go through the process of creating a novel recipe, in this case a combination of both intrinsic (inpainting) and classifier-free context guidance.

Open the notebook →

Coupling¶

Implement Coupling s in stix. A coupling pairs source and target samples together; choosing one that minimises the energy required to move between the two distributions can substantially reduce the number of solver steps needed at sampling time.

What you’ll learn

  • How the default (no coupling: independent pairing) compares to a mini-batch optimal-transport coupling.

  • How to implement your own mini-batch optimal-transport coupling, worked through two examples: on a discrete index modality and on a continuous coordinates modality.

  • How to couple at the dataset level instead of the batch level, reproducing rectified flow.

  • Observe the effect these different coupling approaches have on the training time and sampling performance.

Open the notebook →

Discrete models¶

Model discrete modalities (tokens, categories) as a continuous-time Markov chain (CTMC), and sample them jointly with the continuous O/SDEs. Covers masked and uniform diffusion — the two share a single convex mixture path and differ only in their source distribution.

What you’ll learn

Open the notebook →