Implemented Methods

MedLatents implements a broad set of generative-modeling, sampling, and post-training methods, each grounded in the literature. This page gives a high-level map of what is implemented and where; full citations are collected in References.

Generative model families

Family	Method	Module
Autoregressive	Causal transformer over discrete tokens; speculative and Medusa multi-head decoding for acceleration.	medlatents.autoregressive
Masked parallel	MaskGIT bidirectional masked transformer with confidence-based decoding.	medlatents.maskgit
Discrete diffusion	D3PM (absorbing / uniform transitions), plus SEDD score-entropy and MDLM masked-diffusion objectives.	medlatents.diffusion
Flow matching	Discrete mixture-path flow matching and continuous rectified flow, with optimal-transport coupling and shortcut models.	medlatents.flow_matching
Bayesian flow	Bayesian Flow Networks for discrete data with entropy encoding, score-guided and particle sampling, and higher-order solvers.	medlatents.bayesian_flow
Continuous latent diffusion	Continuous Gaussian diffusion paired with a continuous DiT backbone.	medlatents.diffusion, medlatents.networks

Sampling and inference techniques

• Halton scheduling – low-discrepancy, spatially dispersed unmasking for MaskGIT (medlatents.sampling.MaskGITScheduler).

• KLASS early stopping – KL-adaptive stability sampling that halts MaskGIT and D3PM once the prediction stabilizes (medlatents.sampling.KLASSGenerator, medlatents.diffusion.D3PMKLASS).

• Running confidence remasking – confidence-driven re-masking for masked decoders (medlatents.sampling.RunningConfidenceRemasker).

• Classifier-free guidance – constant, linear, cosine, and triangular time-dependent schedules, plus CFG-Zero* rescaling (medlatents.sampling.GuidedSampler, medlatents.sampling.cfg_zero_star_guidance()).

• Decoupled straight-through estimators – decoupled ST-Gumbel-Softmax and ReinMax for low-variance discrete gradients (medlatents.sampling.DecoupledSTGumbelSoftmax).

• Higher-order ODE solvers – Euler / Heun / RK4 integration for rectified flow, and DPM-Solver / DDIM / Euler schedulers for diffusion (medlatents.sampling).

• Speculative and Medusa decoding – draft-model and multi-head speculative decoding for autoregressive acceleration (medlatents.autoregressive.SpeculativeDecoder, medlatents.autoregressive.MedusaModel).

Training techniques

• Zero-terminal-SNR schedules for diffusion (medlatents.diffusion.enforce_zero_terminal_snr(), medlatents.diffusion.min_snr_weighting()).

• Representation alignment (REPA / REPA-E) – align internal features to pretrained encoders, including end-to-end VAE + DiT training (medlatents.training.REPALoss, medlatents.training.REPAETrainer).

• Contrastive flow-matching regularizers (medlatents.training.ContrastiveFlowMatchingLoss).

• Curriculum learning over masking ratio, noise level, sequence length, and token difficulty (medlatents.training).

• Unified conditioning via medlatents.conditioning.ConditioningBundle and frozen pretrained encoders (DINOv2, SigLIP, MedSigLIP, NeuroVFM).

Post-training and alignment

These methods are documented in depth in Post-Training Guide.

• Preference optimization – DPO and per-architecture variants (autoregressive, MaskGIT, D3PM, flow), with sigmoid/hinge/IPO/KTO losses, and step-level SPO.

• Reinforcement learning – DDPO, GRPO, and GARDO for reward-driven fine-tuning.

• Distillation – reflow and consistency distillation for few-step sampling.

• Self-play – SPIN and rejection fine-tuning (RFT).

Spatial-to-sequence conversion

• Space-filling curves – Hilbert and Z-order rasterization that preserve spatial locality when flattening grids to sequences (medlatents.rasterization). See Token Interface for the layout contract.

Research process

New methods are integrated through a deliberate triage / reproduce / ablate / integrate gate rather than added speculatively. Open an issue to propose a method before implementing it.