Skip to main content

Advanced Research Techniques

The Research section collects advanced techniques that go beyond standard supervised training. Each page is a research-grade, from-scratch implementation in JAX and Flax NNX — with the math, runnable code, and the design decisions that matter. These guides assume you already know how to train a model; they focus on the non-standard ideas layered on top.

Prerequisites

These pages move fast and assume solid fundamentals. Make sure you are comfortable with a simple training loop and the training best practices before diving in.

Start here​

New to this section? Begin with Custom Training Loops. Almost every technique below builds on the explicit TrainState and custom training-step patterns it introduces.

Representation and model efficiency​

Learn powerful representations and shrink models without losing accuracy.

  • Contrastive Learning — SimCLR self-supervised representation learning: the NT-Xent loss, augmentation pipeline, and linear-probe evaluation, all without labels.
  • Knowledge Distillation — Train small, fast student models from large teachers using temperature-scaled soft targets.
  • Metric Learning — Learn embeddings where same-class inputs cluster, using Siamese networks and triplet loss with mining.

Learning paradigms​

Move beyond a single fixed dataset and objective.

  • Meta-Learning — MAML / learning-to-learn: find an initialization that adapts to new tasks in a few gradient steps.
  • Reinforcement Learning — Deep Q-Networks (DQN) in JAX with experience replay, target networks, and JAX-native gymnax environments.
  • Curriculum Learning — Easy-to-hard training schedules with difficulty scoring, pacing functions, and self-paced learning.

Robustness and rigor​

Make models trustworthy and results reproducible.

  • Adversarial Training — Build robustness against adversarial examples using FGSM and the min-max robust optimization objective.
  • Experiment Reproducibility — Deterministic, reproducible runs via explicit PRNG keys, config management, and XLA determinism.
  • Interpretability & Saliency — Explain predictions with vanilla-gradient saliency, Integrated Gradients, and Grad-CAM.
  • Uncertainty Estimation — Quantify what the model doesn't know with MC-dropout and deep ensembles.

Search and control​

Take fine-grained control of the architecture and the training step itself.

  • Custom Training Loops — Fine-grained control over the training step: extend TrainState, add gradient clipping, EMA, and auxiliary losses.
  • Neural Architecture Search — DARTS / automating architecture design via continuous relaxation and bilevel optimization.