From fb03cf1e8bb8b78365ed614e8cd04dfcd3975839 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:43:10 -0400 Subject: [PATCH 1/9] refactor(trainer): the Objective seam, one trainer for any learning problem The train step hardcoded diffusion: normalize, VAE encode, CFG dropout, sample timesteps, corrupt, apply, weight. Nothing about sharding, EMA, checkpointing or the loops needs to know any of that, and JEPA needs a different loss over several parameter subtrees with an EMA that tracks only one of them. An Objective now owns what is being learned - init_params, loss, an EMA spec and the validation artifacts - and the trainer owns the mechanics. The loss always runs under has_aux, so an objective can report per-component diagnostics; the train loop forwards them to wandb. The EMA spec carries a step-indexed decay schedule and a subtree path, because I-JEPA ramps its momentum and averages only its context encoder. DiffusionObjective is the existing objective moved behind that interface, verified bit-identical: params, EMA and adam state after five real steps match the inlined implementation exactly, and that fingerprint is now a test. Two fixes fell out of the move. The video check read len(sample_data_shape) == 5 on a shape that never carries the batch axis, so it was always false and video validation sampled with sequence_length=None. And unconditional ModulatedBlocks now exist: conditioning=None gives a plain pre-norm residual block with learned affine norms, which is the ViT block a JEPA encoder needs. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/dit_common.py | 43 ++- flaxdiff/trainer/general_diffusion_trainer.py | 341 +++++------------- flaxdiff/trainer/objectives.py | 198 ++++++++++ flaxdiff/trainer/simple_trainer.py | 3 +- tests/test_objectives.py | 142 ++++++++ 5 files changed, 463 insertions(+), 264 deletions(-) create mode 100644 flaxdiff/trainer/objectives.py create mode 100644 tests/test_objectives.py diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py index 848d1b9..a3bd4d7 100644 --- a/flaxdiff/models/dit_common.py +++ b/flaxdiff/models/dit_common.py @@ -177,11 +177,16 @@ class ModulatedBlock(nn.Module): mixer='ssm' replaces attention with a bidirectional S5 scan, optionally followed by Spatial-Mamba style 2D state fusion. freqs_cis is unused by the SSM mixer but kept in the interface so blocks are interchangeable. + + modulated=False drops the adaLN-Zero conditioning path entirely, leaving a + plain pre-norm residual block with learned affine norms - the ViT block a + JEPA encoder needs, where there is no timestep to condition on. """ features: int num_heads: int rope_emb: RotaryEmbedding = None mixer: str = 'attention' + modulated: bool = True mlp_ratio: int = 4 dropout_rate: float = 0.0 dtype: Optional[Dtype] = None @@ -201,13 +206,17 @@ def setup(self): assert self.mixer in ('attention', 'ssm'), f"Unknown mixer {self.mixer}" hidden_features = int(self.features * self.mlp_ratio) - self.ada_params_module = AdaLNParams( - self.features, dtype=self.dtype, precision=self.precision) + if self.modulated: + self.ada_params_module = AdaLNParams( + self.features, dtype=self.dtype, precision=self.precision) + # Without modulation the norms carry their own affine, since there is + # no conditioning vector left to supply the shift and scale + affine = not self.modulated self.norm1 = nn.LayerNorm( - epsilon=self.norm_epsilon, use_scale=False, use_bias=False, + epsilon=self.norm_epsilon, use_scale=affine, use_bias=affine, dtype=self.dtype, name="norm1") self.norm2 = nn.LayerNorm( - epsilon=self.norm_epsilon, use_scale=False, use_bias=False, + epsilon=self.norm_epsilon, use_scale=affine, use_bias=affine, dtype=self.dtype, name="norm2") if self.mixer == 'attention': @@ -281,9 +290,16 @@ def _apply_2d_fusion(self, ssm_output): @nn.compact def __call__(self, x, conditioning, freqs_cis, train: bool = False): - scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( - self.ada_params_module(conditioning), 6, axis=-1 - ) + if self.modulated: + scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( + self.ada_params_module(conditioning), 6, axis=-1 + ) + else: + assert conditioning is None, "an unmodulated block takes no conditioning" + # Identity modulation - the block body below collapses to a plain + # pre-norm residual block without branching on the mode + scale_mlp = shift_mlp = scale_attn = shift_attn = 0.0 + gate_mlp = gate_attn = 1.0 # --- Mixer path (attention or SSM) --- residual = x @@ -314,12 +330,17 @@ def __call__(self, x, conditioning, freqs_cis, train: bool = False): return x +def identity_rope_freqs(seq_len: int, dim: int, dtype: Dtype = jnp.float32): + """RoPE frequencies that rotate by nothing, for sequences whose 1D index is + not a position (scan orders other than raster, masked token subsets).""" + shape = (seq_len, dim // 2) + return jnp.ones(shape, dtype=dtype), jnp.zeros(shape, dtype=dtype) + + def neutralized_rope_freqs(rope: RotaryEmbedding, seq_len: int, scan_order: str): """RoPE frequencies for the sequence, neutralized to identity for hilbert/zigzag scans where the 1D index is not a 2D position (the 2D sincos embedding already carries position there).""" - freqs_cos, freqs_sin = rope(seq_len=seq_len) if scan_order != 'raster': - freqs_cos = jnp.ones_like(freqs_cos) - freqs_sin = jnp.zeros_like(freqs_sin) - return freqs_cos, freqs_sin + return identity_rope_freqs(seq_len, rope.dim) + return rope(seq_len=seq_len) diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index b651d9d..835c435 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -10,32 +10,48 @@ from jax.sharding import Mesh, PartitionSpec as P from jax.experimental.shard_map import shard_map -from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple +from ..schedulers import NoiseScheduler from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform -from ..samplers.common import DiffusionSampler -from ..samplers.ddim import DDIMSampler from flaxdiff.utils import RandomMarkovState, serialize_model, get_latest_checkpoint from flaxdiff.inputs import ConditioningEncoder, ConditionalInputConfig, DiffusionInputConfig from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics, convert_to_global_tree +from .objectives import Objective, DiffusionObjective from flaxdiff.models.autoencoder.autoencoder import AutoEncoder from flax.training import dynamic_scale as dynamic_scale_lib import shutil +def _replace_subtree(tree, path, value): + if not path: + return value + return {**tree, path[0]: _replace_subtree(tree[path[0]], path[1:], value)} + + +def _subtree(tree, path): + for key in path: + tree = tree[key] + return tree + + class TrainState(SimpleTrainState): rngs: jax.random.PRNGKey ema_params: dict - def apply_ema(self, decay: float = 0.999): - new_ema_params = jax.tree_util.tree_map( + def apply_ema(self, decay: float = 0.999, path: Tuple[str, ...] = ()): + """EMA over a subtree of the parameters (the whole tree by default). + + JEPA's target encoder is the EMA of the context encoder alone, so the + predictor's slice of the tree must be left out of the average. + """ + new_subtree = jax.tree_util.tree_map( lambda ema, param: decay * ema + (1 - decay) * param, - self.ema_params, - self.params, + _subtree(self.ema_params, path), + _subtree(self.params, path), ) - return self.replace(ema_params=new_ema_params) + return self.replace(ema_params=_replace_subtree(self.ema_params, path, new_subtree)) from flaxdiff.metrics.common import EvaluationMetric @@ -90,9 +106,10 @@ class GeneralDiffusionTrainer(SimpleTrainer): def __init__(self, model: nn.Module, optimizer: optax.GradientTransformation, - noise_schedule: NoiseScheduler, input_config: DiffusionInputConfig, rngs: jax.random.PRNGKey, + noise_schedule: NoiseScheduler = None, + objective: Objective = None, unconditional_prob: float = 0.12, name: str = "GeneralDiffusion", model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(), @@ -103,17 +120,19 @@ def __init__(self, eval_metrics: List[EvaluationMetric] = None, best_tracker_metric: str = "train/best_loss", ema_decay: float = 0.999, + loss_fn: Callable = optax.l2_loss, **kwargs ): """ Initialize the general diffusion trainer. - + Args: model: Neural network model optimizer: Optimization algorithm - noise_schedule: Noise scheduler for diffusion process input_config: Configuration for input data, including keys, shapes and conditioning inputs rngs: Random number generator keys + noise_schedule: Noise scheduler for the default diffusion objective + objective: What to optimize. Defaults to diffusion over `model`. unconditional_prob: Probability of training with unconditional samples name: Name of this trainer model_output_transform: Transform for model predictions @@ -167,6 +186,21 @@ def __init__(self, native_resolution = native_resolution * autoencoder.downscale_factor self.native_resolution = native_resolution + if objective is None: + objective = DiffusionObjective( + model=model, + noise_schedule=noise_schedule, + model_output_transform=model_output_transform, + input_config=input_config, + input_shapes=input_shapes, + autoencoder=autoencoder, + unconditional_prob=unconditional_prob, + loss_fn=loss_fn, + ema_decay=ema_decay, + native_resolution=native_resolution, + ) + self.objective = objective + super().__init__( model=model, input_shapes=input_shapes, @@ -174,6 +208,7 @@ def __init__(self, rngs=rngs, name=name, wandb_config=wandb_config, + loss_fn=loss_fn, **kwargs ) @@ -184,8 +219,6 @@ def __init__(self, # List of conditional inputs self.conditional_inputs = input_config.conditions - # Determine if we're working with video or images - self.is_video = self._is_video_data() self.best_val_metrics = {} # val/ -> True if higher is better (default lower-is-better) @@ -204,8 +237,7 @@ def generate_states( print("Generating states for DiffusionTrainer") rngs, subkey = jax.random.split(rngs) - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) + params = self.objective.init_params(subkey) state = TrainState.create( apply_fn=model.apply, @@ -218,121 +250,44 @@ def generate_states( ) return state, state - def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): + def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, **validation_step_args): local_batch_size = data['local_batch_size'] - validation_step_args = { - "sampler_class": sampler_class, - "sampling_noise_schedule": sampling_noise_schedule, - } return super().fit( - data, - train_steps_per_epoch=training_steps_per_epoch, - epochs=epochs, - train_step_args={"batch_size": local_batch_size}, + data, + train_steps_per_epoch=training_steps_per_epoch, + epochs=epochs, + train_step_args={"batch_size": local_batch_size}, val_steps_per_epoch=val_steps_per_epoch, validation_step_args=validation_step_args, ) - def _is_video_data(self): - sample_data_shape = self.input_config.sample_data_shape - return len(sample_data_shape) == 5 - def _define_train_step(self, batch_size): """ - Define the training step function for both image and video diffusion. - Optimized for efficient sharding and JIT compilation. + Define the training step: the objective supplies the loss, the trainer + supplies gradients, sharding, EMA and the state update. """ - # Access class variables once for JIT optimization - noise_schedule = self.noise_schedule - model = self.model - model_output_transform = self.model_output_transform - loss_fn = self.loss_fn + objective = self.objective distributed_training = self.distributed_training - autoencoder = self.autoencoder - unconditional_prob = self.unconditional_prob - - input_config = self.input_config - sample_data_key = input_config.sample_data_key - - # JIT-optimized function for processing conditional inputs - # @functools.partial(jax.jit, static_argnums=(2,)) - def process_conditioning(batch, uncond_mask): - return input_config.process_conditioning( - batch, - uncond_mask=uncond_mask, - ) + ema = objective.ema - # Main training step function - optimized for JIT compilation and sharding def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index): """Training step optimized for distributed execution.""" - # Random key handling rng_state, key_fold = rng_state.get_random_key() folded_key = jax.random.fold_in(key_fold, local_device_index.reshape()) - local_rng_state = RandomMarkovState(folded_key) - - # Extract and normalize data (works for both images and videos) - data = batch[sample_data_key] - local_batch_size = data.shape[0] - data = (jnp.asarray(data, dtype=jnp.float32) - 127.5) / 127.5 - - # Autoencoder step (handles both image and video data) - if autoencoder is not None: - local_rng_state, enc_key = local_rng_state.get_random_key() - data = autoencoder.encode(data, enc_key) - - # Determine number of unconditional samples per mini batch randomly - local_rng_state, uncond_key = local_rng_state.get_random_key() - # Determine unconditional samples - uncond_mask = jax.random.bernoulli( - uncond_key, - shape=(local_batch_size,), - p=unconditional_prob - ) - - # Process conditioning - all_conditional_inputs = process_conditioning(batch, uncond_mask) - - # Generate diffusion timesteps - noise_level, local_rng_state = noise_schedule.generate_timesteps(local_batch_size, local_rng_state) - - # Generate noise - local_rng_state, noise_key = local_rng_state.get_random_key() - noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32) - local_rng_state, dropout_key = local_rng_state.get_random_key() - - # Forward diffusion process - rates = noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data)) - noisy_data, c_in, expected_output = model_output_transform.forward_diffusion(data, noise, rates) - - # Loss function - def model_loss(params): - # Apply model - inputs = noise_schedule.transform_inputs(noisy_data * c_in, noise_level) - preds = model.apply( - params, *inputs, *all_conditional_inputs, - train=True, rngs={'dropout': dropout_key}, - ) - - # Transform predictions and calculate loss - preds = model_output_transform.pred_transform(noisy_data, preds, rates) - sample_losses = loss_fn(preds, expected_output) - - # Apply loss weighting - weights = noise_schedule.get_weights(noise_level, get_coeff_shapes_tuple(sample_losses)) - weighted_loss = sample_losses * weights - - return jnp.mean(weighted_loss) - - # Compute gradients and apply updates + def objective_loss(params): + return objective.loss(params, train_state.ema_params, batch, + folded_key, train_state.step) + if train_state.dynamic_scale is not None: # Mixed precision training with dynamic scale - grad_fn = train_state.dynamic_scale.value_and_grad(model_loss, axis_name="data") - dynamic_scale, is_finite, loss, grads = grad_fn(train_state.params) - + grad_fn = train_state.dynamic_scale.value_and_grad( + objective_loss, has_aux=True, axis_name="data") + dynamic_scale, is_finite, (loss, aux), grads = grad_fn(train_state.params) + train_state = train_state.replace(dynamic_scale=dynamic_scale) new_state = train_state.apply_gradients(grads=grads) - + # Handle NaN/Inf gradients select_fn = functools.partial(jnp.where, is_finite) new_state = new_state.replace( @@ -340,104 +295,37 @@ def model_loss(params): params=jax.tree.map(select_fn, new_state.params, train_state.params) ) else: - # Standard gradient computation - grad_fn = jax.value_and_grad(model_loss) - loss, grads = grad_fn(train_state.params) - + grad_fn = jax.value_and_grad(objective_loss, has_aux=True) + (loss, aux), grads = grad_fn(train_state.params) + if distributed_training: grads = jax.lax.pmean(grads, axis_name="data") - + new_state = train_state.apply_gradients(grads=grads) - - # Apply EMA update - new_state = new_state.apply_ema(self.ema_decay) - - # Average loss across devices if distributed + + new_state = new_state.apply_ema(ema.decay(train_state.step), ema.path) + if distributed_training: loss = jax.lax.pmean(loss, axis_name="data") - - return new_state, loss, rng_state + aux = jax.tree.map(lambda v: jax.lax.pmean(v, axis_name="data"), aux) + + return new_state, loss, aux, rng_state # Apply sharding for distributed training if distributed_training: train_step = shard_map( - train_step, - mesh=self.mesh, - in_specs=(P(), P(), P('data'), P('data')), - out_specs=(P(), P(), P()), + train_step, + mesh=self.mesh, + in_specs=(P(), P(), P('data'), P('data')), + out_specs=(P(), P(), P(), P()), ) - + # Apply JIT compilation train_step = jax.jit(train_step, donate_argnums=(2)) return train_step - def _define_validation_step(self, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): - """ - Define the validation step for both image and video diffusion models. - """ - # Setup for validation - model = self.model - autoencoder = self.autoencoder - input_config = self.input_config - conditional_inputs = self.conditional_inputs - is_video = self.is_video - - # Get necessary parameters - image_size = self._get_image_size() - - # Get sequence length only for video data - sequence_length = self._get_sequence_length() if is_video else None - - # Initialize the sampler - sampler = sampler_class( - model=model, - noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule, - model_output_transform=self.model_output_transform, - input_config=input_config, - autoencoder=autoencoder, - guidance_scale=3.0, - ) - - def generate_samples( - val_state: TrainState, - batch, - diffusion_steps: int, - ): - # Process all conditional inputs - model_conditioning_inputs = [cond_input(batch) for cond_input in conditional_inputs] - - # Determine batch size - batch_size = len(model_conditioning_inputs[0]) if model_conditioning_inputs else 4 - - # Generate samples - works for both images and videos - return sampler.generate_samples( - params=val_state.ema_params, - resolution=image_size, - num_samples=batch_size, - sequence_length=sequence_length, # Will be None for images - diffusion_steps=diffusion_steps, - end_step=0, - priors=None, - model_conditioning_inputs=tuple(model_conditioning_inputs), - ) - - return generate_samples - - def _get_image_size(self): - """Helper to determine image size from available information.""" - if self.native_resolution is not None: - return self.native_resolution - - sample_data_shape = self.input_config.sample_data_shape - return sample_data_shape[-2] # Assuming [..., H, W, C] format - - def _get_sequence_length(self): - """Helper to determine sequence length for video generation.""" - if not self.is_video: - return None - - sample_data_shape = self.input_config.sample_data_shape - return sample_data_shape[1] # Assuming [B,T,H,W,C] format + def _define_validation_step(self, **kwargs): + return self.objective.make_validation_step(**kwargs) def validation_loop( self, @@ -446,10 +334,9 @@ def validation_loop( val_ds, val_steps_per_epoch, current_step, - diffusion_steps=200, ): """ - Run validation and log samples for both image and video diffusion. + Score the objective's validation artifacts and let it visualize them. """ global_device_count = jax.device_count() local_device_count = jax.local_device_count() @@ -468,13 +355,8 @@ def validation_loop( batch = next(val_ds) if self.distributed_training and global_device_count > 1: batch = convert_to_global_tree(self.mesh, batch) - # Generate samples - samples = generate_samples( - val_state, - batch, - diffusion_steps, - ) - + samples = generate_samples(val_state, batch) + if self.eval_metrics is not None: for metric in self.eval_metrics: try: @@ -489,16 +371,10 @@ def validation_loop( if i == 0: print(f"Evaluation started for process index {process_index}") - # Log samples to wandb if self.wandb is not None and self.wandb: - import numpy as np - - # Process samples differently based on dimensionality - if len(samples.shape) == 5: # [B,T,H,W,C] - Video data - self._log_video_samples(samples, current_step) - else: # [B,H,W,C] - Image data - self._log_image_samples(samples, current_step) - + self.objective.log_validation_artifacts( + self.wandb, samples, current_step) + # Flatten the metrics if metrics: metrics = {k: np.mean(v) for k, v in metrics.items()} @@ -538,46 +414,7 @@ def validation_loop( import traceback traceback.print_exc() - - def _log_video_samples(self, samples, current_step): - """Helper to log video samples to wandb.""" - import numpy as np - from wandb import Video as wandbVideo - - for i in range(samples.shape[0]): - # Convert to numpy, denormalize and clip - sample = np.array(samples[i]) - sample = (sample + 1) * 127.5 - sample = np.clip(sample, 0, 255).astype(np.uint8) - - # Log as video - self.wandb.log({ - f"video_sample_{i}": wandbVideo( - sample, - fps=10, - caption=f"Video Sample {i} at step {current_step}" - ) - }, step=current_step) - - def _log_image_samples(self, samples, current_step): - """Helper to log image samples to wandb.""" - import numpy as np - from wandb import Image as wandbImage - - for i in range(samples.shape[0]): - # Convert to numpy, denormalize and clip - sample = np.array(samples[i]) - sample = (sample + 1) * 127.5 - sample = np.clip(sample, 0, 255).astype(np.uint8) - - # Log as image - self.wandb.log({ - f"sample_{i}": wandbImage( - sample, - caption=f"Sample {i} at step {current_step}" - ) - }, step=current_step) - + def push_to_registry( self, registry_name: str = 'wandb-registry-model', diff --git a/flaxdiff/trainer/objectives.py b/flaxdiff/trainer/objectives.py new file mode 100644 index 0000000..b9edfb1 --- /dev/null +++ b/flaxdiff/trainer/objectives.py @@ -0,0 +1,198 @@ +"""What the trainer is optimizing. + +The trainer owns training mechanics - sharding, EMA bookkeeping, checkpoints, +logging, the loops. An Objective owns what is being learned: the parameters it +holds, the loss it computes from a batch, the telemetry that loss reports, and +the artifacts validation scores. Swapping the objective swaps the research +question without touching any of the mechanics. + +The loss always returns auxiliary metrics alongside the scalar, so an objective +with several loss terms or with per-step diagnostics (JEPA's collapse +telemetry) can surface them without the trainer knowing what they mean. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Type + +import jax +import jax.numpy as jnp +import optax + +from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple +from ..predictors import DiffusionPredictionTransform +from ..samplers.common import DiffusionSampler +from ..samplers.ddim import DDIMSampler +from ..models.autoencoder.autoencoder import AutoEncoder +from ..inputs import DiffusionInputConfig +from ..utils import RandomMarkovState + + +@dataclass +class EMASpec: + """Which slice of the parameters the EMA copy tracks, and how fast. + + decay is a step-indexed schedule rather than a constant because momentum + ramps are load-bearing for some objectives (I-JEPA anneals 0.996 -> 1.0). + path selects a subtree; the empty path means the whole parameter tree. + """ + decay: Callable[[Any], Any] + path: Tuple[str, ...] = () + + +class Objective(ABC): + """The learning problem: parameters, loss, EMA policy, validation artifacts.""" + + ema: EMASpec + + @abstractmethod + def init_params(self, rng: jax.Array) -> Any: + """Build the full parameter tree, which may hold several sub-modules.""" + + @abstractmethod + def loss(self, params: Any, ema_params: Any, batch: Any, rng: jax.Array, + step: Any) -> Tuple[jax.Array, Dict[str, jax.Array]]: + """Scalar loss plus auxiliary metrics, differentiated w.r.t. params.""" + + @abstractmethod + def make_validation_step(self, **kwargs) -> Callable[[Any, Any], Any]: + """Build (val_state, batch) -> artifacts, which eval metrics score.""" + + def log_validation_artifacts(self, wandb, artifacts, step: int): + """Visualize the artifacts. Nothing to draw unless an objective says so.""" + + +class DiffusionObjective(Objective): + """Denoising diffusion: sample a noise level, corrupt, predict, weight.""" + + def __init__( + self, + model, + noise_schedule: NoiseScheduler, + model_output_transform: DiffusionPredictionTransform, + input_config: DiffusionInputConfig, + input_shapes: Dict[str, Tuple[int, ...]], + autoencoder: AutoEncoder = None, + unconditional_prob: float = 0.12, + loss_fn: Callable = optax.l2_loss, + ema_decay: float = 0.999, + native_resolution: int = None, + diffusion_steps: int = 200, + ): + self.model = model + self.noise_schedule = noise_schedule + self.model_output_transform = model_output_transform + self.input_config = input_config + self.input_shapes = input_shapes + self.autoencoder = autoencoder + self.unconditional_prob = unconditional_prob + self.loss_fn = loss_fn + self.native_resolution = native_resolution + self.diffusion_steps = diffusion_steps + self.ema = EMASpec(decay=optax.constant_schedule(ema_decay)) + + def init_params(self, rng): + input_ones = {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()} + return self.model.init(rng, **input_ones) + + def loss(self, params, ema_params, batch, rng, step): + local_rng_state = RandomMarkovState(rng) + + # Extract and normalize data (works for both images and videos) + data = batch[self.input_config.sample_data_key] + local_batch_size = data.shape[0] + data = (jnp.asarray(data, dtype=jnp.float32) - 127.5) / 127.5 + + if self.autoencoder is not None: + local_rng_state, enc_key = local_rng_state.get_random_key() + data = self.autoencoder.encode(data, enc_key) + + local_rng_state, uncond_key = local_rng_state.get_random_key() + uncond_mask = jax.random.bernoulli( + uncond_key, shape=(local_batch_size,), p=self.unconditional_prob) + all_conditional_inputs = self.input_config.process_conditioning( + batch, uncond_mask=uncond_mask) + + noise_level, local_rng_state = self.noise_schedule.generate_timesteps( + local_batch_size, local_rng_state) + + local_rng_state, noise_key = local_rng_state.get_random_key() + noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32) + + local_rng_state, dropout_key = local_rng_state.get_random_key() + + rates = self.noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data)) + noisy_data, c_in, expected_output = self.model_output_transform.forward_diffusion( + data, noise, rates) + + inputs = self.noise_schedule.transform_inputs(noisy_data * c_in, noise_level) + preds = self.model.apply( + params, *inputs, *all_conditional_inputs, + train=True, rngs={'dropout': dropout_key}, + ) + + preds = self.model_output_transform.pred_transform(noisy_data, preds, rates) + sample_losses = self.loss_fn(preds, expected_output) + + weights = self.noise_schedule.get_weights( + noise_level, get_coeff_shapes_tuple(sample_losses)) + return jnp.mean(sample_losses * weights), {} + + def make_validation_step( + self, + sampler_class: Type[DiffusionSampler] = DDIMSampler, + sampling_noise_schedule: NoiseScheduler = None, + ): + sampler = sampler_class( + model=self.model, + noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule, + model_output_transform=self.model_output_transform, + input_config=self.input_config, + autoencoder=self.autoencoder, + guidance_scale=3.0, + ) + conditional_inputs = self.input_config.conditions + image_size = self._image_size() + sequence_length = self._sequence_length() + + def generate_samples(val_state, batch): + model_conditioning_inputs = [cond_input(batch) for cond_input in conditional_inputs] + batch_size = len(model_conditioning_inputs[0]) if model_conditioning_inputs else 4 + return sampler.generate_samples( + params=val_state.ema_params, + resolution=image_size, + num_samples=batch_size, + sequence_length=sequence_length, # None for images + diffusion_steps=self.diffusion_steps, + end_step=0, + priors=None, + model_conditioning_inputs=tuple(model_conditioning_inputs), + ) + + return generate_samples + + def log_validation_artifacts(self, wandb, artifacts, step: int): + import numpy as np + + is_video = len(artifacts.shape) == 5 + for i in range(artifacts.shape[0]): + sample = np.clip((np.array(artifacts[i]) + 1) * 127.5, 0, 255).astype(np.uint8) + if is_video: + from wandb import Video as wandbVideo + wandb.log({f"video_sample_{i}": wandbVideo( + sample, fps=10, caption=f"Video Sample {i} at step {step}")}, step=step) + else: + from wandb import Image as wandbImage + wandb.log({f"sample_{i}": wandbImage( + sample, caption=f"Sample {i} at step {step}")}, step=step) + + def _is_video(self): + return len(self.input_config.sample_data_shape) == 4 + + def _image_size(self): + if self.native_resolution is not None: + return self.native_resolution + return self.input_config.sample_data_shape[-2] + + def _sequence_length(self): + return self.input_config.sample_data_shape[0] if self._is_video() else None diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index e9ff541..c42b18b 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -439,7 +439,7 @@ def train_loop( if self.distributed_training and global_device_count > 1: # # Convert the local device batches to a unified global jax.Array batch = convert_to_global_tree(self.mesh, batch) - train_state, loss, rng_state = train_step_fn(train_state, rng_state, batch, global_device_indexes) + train_state, loss, aux, rng_state = train_step_fn(train_state, rng_state, batch, global_device_indexes) if i == 0: print(f"Training started for process index {process_index} at step {current_step}") @@ -470,6 +470,7 @@ def train_loop( self.wandb.log({ "train/step" : current_step, "train/loss": loss, + **{f"train/{k}": v for k, v in aux.items()}, }, step=current_step) # Save the model every few steps if save_every and i % save_every == 0 and i > 0: diff --git a/tests/test_objectives.py b/tests/test_objectives.py new file mode 100644 index 0000000..fd6b39d --- /dev/null +++ b/tests/test_objectives.py @@ -0,0 +1,142 @@ +"""The Objective seam: the trainer owns mechanics, the objective owns the loss. + +The diffusion objective was lifted out of the trainer's train step, so the +first test here is a golden fingerprint of the parameters, EMA and optimizer +state after five real steps, captured from the implementation that inlined the +objective. It must not move. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.predictors import get_diffusion_preset +from flaxdiff.trainer import GeneralDiffusionTrainer +from flaxdiff.trainer.general_diffusion_trainer import TrainState +from flaxdiff.trainer.objectives import DiffusionObjective, EMASpec, Objective + +RES = 8 + + +def tree_fingerprint(tree): + # per-leaf sums accumulated in python floats, so the golden values below + # do not depend on float32 reduction order + return sum(float(jnp.sum(leaf)) for leaf in jax.tree.leaves(tree) + if jnp.issubdtype(jnp.asarray(leaf).dtype, jnp.floating)) + + +def make_trainer(tmp_path, **kwargs): + train_schedule, _, transform = get_diffusion_preset("edm") + return GeneralDiffusionTrainer( + model=SimpleDiT(patch_size=4, emb_features=16, num_layers=1, num_heads=2, mlp_ratio=1), + optimizer=optax.adam(1e-3), + noise_schedule=train_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig( + sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]), + rngs=jax.random.PRNGKey(0), + name="objective-parity", + wandb_config=None, + distributed_training=False, + checkpoint_base_path=str(tmp_path), + **kwargs, + ) + + +def batch_iterator(): + images = np.tile(np.linspace(0, 255, RES, dtype=np.float32)[None, :, None, None], + (8, 1, RES, 3)) + while True: + yield {"image": jnp.asarray(images)} + + +def test_diffusion_objective_reproduces_the_inlined_train_step(tmp_path): + trainer = make_trainer(tmp_path) + data = {"train": batch_iterator, "train_len": 32, "local_batch_size": 8} + state = trainer.fit(data, training_steps_per_epoch=5, epochs=1, val_steps_per_epoch=0) + + assert tree_fingerprint(state.params) == pytest.approx(8.4679913880548, abs=1e-9) + assert tree_fingerprint(state.ema_params) == pytest.approx(8.220795648946352, abs=1e-9) + assert tree_fingerprint(state.opt_state) == pytest.approx(-0.004871756572780005, abs=1e-12) + + +def test_train_step_returns_auxiliary_metrics(tmp_path): + """value_and_grad runs with has_aux so objectives can report diagnostics.""" + trainer = make_trainer(tmp_path) + train_step = trainer._define_train_step(batch_size=8) + batch = next(batch_iterator()) + + new_state, loss, aux, rng_state = train_step( + trainer.state, trainer.rngstate, batch, jnp.asarray(0)) + assert jnp.isfinite(loss) + assert isinstance(aux, dict) + + +class ConstantObjective(Objective): + """Two independent parameter subtrees, so EMA scoping is observable.""" + + def __init__(self, ema): + self.ema = ema + + def init_params(self, rng): + return {"params": {"tracked": {"w": jnp.ones((2,))}, + "untracked": {"w": jnp.ones((2,))}}} + + def loss(self, params, ema_params, batch, rng, step): + total = sum(jnp.sum(leaf ** 2) for leaf in jax.tree.leaves(params)) + return total, {"probe": jnp.asarray(1.0)} + + def make_validation_step(self, **kwargs): + return lambda val_state, batch: None + + +def make_state(objective, params=None): + params = objective.init_params(jax.random.PRNGKey(0)) if params is None else params + return TrainState.create( + apply_fn=lambda *a, **k: None, params=params, ema_params=params, + tx=optax.sgd(0.0), rngs=jax.random.PRNGKey(0), metrics=None, dynamic_scale=None) + + +def test_ema_over_a_subtree_leaves_its_siblings_alone(): + objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.5), + path=("params", "tracked"))) + state = make_state(objective) + state = state.replace(params=jax.tree.map(lambda p: p * 3.0, state.params)) + + updated = state.apply_ema(0.5, objective.ema.path) + assert jnp.allclose(updated.ema_params["params"]["tracked"]["w"], 2.0) + assert jnp.allclose(updated.ema_params["params"]["untracked"]["w"], 1.0) + + +def test_ema_over_the_whole_tree_is_the_default(): + state = make_state(ConstantObjective(EMASpec(decay=optax.constant_schedule(0.5)))) + state = state.replace(params=jax.tree.map(lambda p: p * 3.0, state.params)) + + updated = state.apply_ema(0.5) + assert jnp.allclose(updated.ema_params["params"]["tracked"]["w"], 2.0) + assert jnp.allclose(updated.ema_params["params"]["untracked"]["w"], 2.0) + + +def test_ema_decay_follows_the_step_schedule(): + """I-JEPA's momentum ramp: the trainer must read decay at the current step.""" + ema = EMASpec(decay=optax.linear_schedule(0.996, 1.0, transition_steps=100)) + assert float(ema.decay(0)) == pytest.approx(0.996) + assert float(ema.decay(100)) == pytest.approx(1.0) + assert float(ema.decay(50)) == pytest.approx(0.998) + + +def test_trainer_drives_an_arbitrary_objective(tmp_path): + """The seam is real: a non-diffusion objective trains through the same loop.""" + objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.9), + path=("params", "tracked"))) + trainer = make_trainer(tmp_path, objective=objective) + data = {"train": batch_iterator, "train_len": 32, "local_batch_size": 8} + state = trainer.fit(data, training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0) + + assert int(state.step) == 3 + # sum of squares under gradient descent must shrink + assert tree_fingerprint(state.params) < tree_fingerprint(objective.init_params(None)) From 06887ce14321d8637d41def0674adb4feaca7fc0 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:44 -0400 Subject: [PATCH 2/9] feat(schedulers): flow matching, min-SNR-gamma weighting, EDM2 defaults Three changes to the training objective layer, sharing one contract. Flow matching slots into the existing scheduler/predictor seams rather than next to them: FlowMatchingScheduler is the linear path alpha = 1 - t, sigma = t with SD3's logit-normal timesteps and Flux's resolution shift, and FlowMatchPredictionTransform predicts the velocity u = eps - x0. Because the path is linear, the unmodified DDIM and Euler updates already are the flow Euler step; there is a test proving the identity exactly. min-SNR-gamma (Hang et al. 2023) is defined on the x_0 loss, so converting it into the space the trainer actually computes the loss in needs the parameterization. get_weights becomes a template over get_schedule_weights and asks the transform for that conversion, which also makes the schedule/transform pairing explicit instead of implicit. No second weighting system. EDMNoiseScheduler's P_mean/P_std were hardcoded to EDM1's -1.2/1.2; they are now constructor arguments defaulting to EDM2's -0.4/1.0. Co-Authored-By: Claude Fable 5 --- flaxdiff/predictors/__init__.py | 68 ++++++++- flaxdiff/schedulers/__init__.py | 3 +- flaxdiff/schedulers/common.py | 30 +++- flaxdiff/schedulers/cosine.py | 2 +- flaxdiff/schedulers/discrete.py | 2 +- flaxdiff/schedulers/flow.py | 55 +++++++ flaxdiff/schedulers/karras.py | 18 ++- tests/test_flow_matching.py | 249 ++++++++++++++++++++++++++++++++ tests/test_schedulers.py | 87 ++++++++++- 9 files changed, 490 insertions(+), 24 deletions(-) create mode 100644 flaxdiff/schedulers/flow.py create mode 100644 tests/test_flow_matching.py diff --git a/flaxdiff/predictors/__init__.py b/flaxdiff/predictors/__init__.py index 68390f2..befce3b 100644 --- a/flaxdiff/predictors/__init__.py +++ b/flaxdiff/predictors/__init__.py @@ -7,6 +7,7 @@ CosineNoiseScheduler, KarrasVENoiseScheduler, EDMNoiseScheduler, + FlowMatchingScheduler, ) ############################################################################################################ @@ -39,6 +40,14 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray: return 1 + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + """||target error||^2 / ||x_0 error||^2 at the given SNR. + + Loss weights (min-SNR-gamma and friends) are defined on the x_0 loss; + dividing by this converts them into the space the model is trained in. + """ + return 1.0 + class EpsilonPredictionTransform(DiffusionPredictionTransform): def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: # preds is the predicted noise @@ -46,10 +55,13 @@ def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) signal_rates, noise_rates = rates x_0 = (x_t - epsilon * noise_rates) / signal_rates return x_0, epsilon - + def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: return epsilon + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + return snr + class DirectPredictionTransform(DiffusionPredictionTransform): def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: # Here the model predicts x_0 directly @@ -76,7 +88,27 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: v = signal_rate * epsilon - noise_rate * x_0 variance = signal_rate**2 + noise_rate**2 return v / jnp.sqrt(variance) - + + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + return snr + 1 + +class FlowMatchPredictionTransform(DiffusionPredictionTransform): + """Rectified flow velocity: the model predicts u = epsilon - x_0, the + constant velocity of the linear path, so both endpoints are one step away. + """ + def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: + signal_rate, noise_rate = rates + x_0 = x_t - noise_rate * preds + epsilon = x_t + signal_rate * preds + return x_0, epsilon + + def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: + return epsilon - x_0 + + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + # x_0 error is t times the velocity error, and t = 1 / (1 + sqrt(SNR)) + return (1 + jnp.sqrt(snr)) ** 2 + class KarrasPredictionTransform(DiffusionPredictionTransform): def __init__(self, sigma_data=0.5) -> None: super().__init__() @@ -102,6 +134,11 @@ def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon) return c_in + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + # x_0 error is c_out times the raw error, and alpha = 1 here so + # sigma^2 = 1 / SNR + return 1 / self.sigma_data ** 2 + snr + ############################################################################################################ # Noise schedule presets ############################################################################################################ @@ -113,6 +150,10 @@ def get_diffusion_preset( sigma_max: float = 80.0, rho: float = 7.0, sigma_data: float = 0.5, + P_mean: float = -0.4, + P_std: float = 1.0, + shift: float = 1.0, + min_snr_gamma: float = None, ) -> tuple[NoiseScheduler, NoiseScheduler, DiffusionPredictionTransform]: """Named (train schedule, sampling schedule, prediction transform) presets. @@ -120,18 +161,29 @@ def get_diffusion_preset( parameterization. Both training and inference build from here, so a model is always sampled with the same convention it was trained with. """ + # Only the training schedule ever produces loss weights, so only it needs + # to know the parameterization if name == 'edm': - train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) - sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) transform = KarrasPredictionTransform(sigma_data=sigma_data) - elif name == 'karras': - train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) + train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data, + P_mean=P_mean, P_std=P_std, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) + elif name == 'karras': transform = KarrasPredictionTransform(sigma_data=sigma_data) + train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) elif name == 'cosine': - train = CosineNoiseScheduler(timesteps, beta_end=1) - sample = CosineNoiseScheduler(timesteps, beta_end=1) transform = VPredictionTransform() + train = CosineNoiseScheduler(timesteps, beta_end=1, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = CosineNoiseScheduler(timesteps, beta_end=1) + elif name in ('flow', 'flow_matching'): + transform = FlowMatchPredictionTransform() + train = FlowMatchingScheduler(shift=shift, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = FlowMatchingScheduler(shift=shift) else: raise ValueError(f"Unknown noise schedule preset: {name}") return train, sample, transform diff --git a/flaxdiff/schedulers/__init__.py b/flaxdiff/schedulers/__init__.py index 32594e4..278bdb9 100644 --- a/flaxdiff/schedulers/__init__.py +++ b/flaxdiff/schedulers/__init__.py @@ -3,4 +3,5 @@ from .cosine import * from .linear import * from .sqrt import * -from .karras import * \ No newline at end of file +from .karras import * +from .flow import * diff --git a/flaxdiff/schedulers/common.py b/flaxdiff/schedulers/common.py index 3061542..e4c2459 100644 --- a/flaxdiff/schedulers/common.py +++ b/flaxdiff/schedulers/common.py @@ -18,11 +18,17 @@ def __init__(self, timesteps, dtype=jnp.float32, clip_min=-1.0, clip_max=1.0, + min_snr_gamma: float = None, + prediction_transform=None, *args, **kwargs): self.max_timesteps = timesteps self.dtype = dtype self.clip_min = clip_min self.clip_max = clip_max + if min_snr_gamma is not None and prediction_transform is None: + raise ValueError("min_snr_gamma needs the prediction transform it will be trained against") + self.min_snr_gamma = min_snr_gamma + self.prediction_transform = prediction_transform if type(timesteps) == int and timesteps > 1: timestep_generator = lambda rng, batch_size, max_timesteps = timesteps: jax.random.randint(rng, (batch_size,), 0, max_timesteps) else: @@ -33,10 +39,28 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n state, rng = state.get_random_key() timesteps = self.timestep_generator(rng, batch_size, self.max_timesteps) return timesteps, state - + + def get_snr(self, steps) -> jnp.ndarray: + signal_rates, noise_rates = self.get_rates(steps, shape=(-1,)) + return (signal_rates / noise_rates) ** 2 + def get_weights(self, steps, shape=(-1, 1, 1, 1)): + """Per-sample loss weight, in the target space of the paired parameterization. + + min-SNR-gamma (Hang et al. 2023) is defined as min(SNR, gamma) on the + x_0 loss; the transform converts that into whichever space the trainer + actually computes the loss in. It replaces the schedule's own weighting + rather than stacking on top of it. + """ + if self.min_snr_gamma is None: + return self.get_schedule_weights(steps, shape) + snr = self.get_snr(steps) + weights = jnp.minimum(snr, self.min_snr_gamma) / self.prediction_transform.target_error_scale(snr) + return weights.reshape(shape) + + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): raise NotImplementedError - + def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]: raise NotImplementedError @@ -77,7 +101,7 @@ def __init__(self, timesteps, sigma_min=0.002, sigma_max=80.0, sigma_data=1, *ar self.sigma_max = sigma_max self.sigma_data = sigma_data - def get_weights(self, steps, shape=(-1, 1, 1, 1)): + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): sigma = self.get_sigmas(steps) return (1 + (1 / (1 + ((1 - sigma ** 2)/(sigma ** 2)))) / (self.sigma_max ** 2)).reshape(shape) diff --git a/flaxdiff/schedulers/cosine.py b/flaxdiff/schedulers/cosine.py index 00de7dd..e1f0bd9 100644 --- a/flaxdiff/schedulers/cosine.py +++ b/flaxdiff/schedulers/cosine.py @@ -34,7 +34,7 @@ def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarra noise_rates = jnp.sin((jnp.pi * steps) / (2 * self.max_timesteps)) return reshape_rates((signal_rates, noise_rates), shape=shape) - def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: alpha, sigma = self.get_rates(steps, shape=shape) return 1 / (1 + (alpha ** 2 / sigma ** 2)) \ No newline at end of file diff --git a/flaxdiff/schedulers/discrete.py b/flaxdiff/schedulers/discrete.py index ee1a01c..a7af264 100644 --- a/flaxdiff/schedulers/discrete.py +++ b/flaxdiff/schedulers/discrete.py @@ -47,7 +47,7 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n def get_p2_weights(self, k, gamma): return (k + self.alpha_cumprod / (1 - self.alpha_cumprod)) ** -gamma - def get_weights(self, steps, shape=(-1, 1, 1, 1)): + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): steps = jnp.int16(steps) return self.p2_loss_weights[steps].reshape(shape) diff --git a/flaxdiff/schedulers/flow.py b/flaxdiff/schedulers/flow.py new file mode 100644 index 0000000..7909fd2 --- /dev/null +++ b/flaxdiff/schedulers/flow.py @@ -0,0 +1,55 @@ +import math +import jax +import jax.numpy as jnp +from ..utils import RandomMarkovState +from .continuous import ContinuousNoiseScheduler +from .common import reshape_rates + + +def compute_resolution_shift(sequence_length, base_seq_len=256, max_seq_len=4096, + base_shift=0.5, max_shift=1.15) -> float: + """Flux-style resolution dependent timestep shift. + + Longer token sequences carry more redundancy, so the trajectory has to + spend more of its budget at high noise for the global structure to settle. + mu is interpolated linearly in sequence length and the shift is exp(mu). + """ + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + mu = base_shift + slope * (sequence_length - base_seq_len) + return math.exp(mu) + + +class FlowMatchingScheduler(ContinuousNoiseScheduler): + """Rectified flow / conditional flow matching on the linear path. + + x_t = (1 - t) * x_0 + t * epsilon for t in [0, 1], so alpha + sigma = 1 and + the model input needs no scaling. Timesteps are drawn logit-normal as in + SD3, which concentrates training on the middle of the trajectory where the + velocity is hardest to predict. + """ + def __init__(self, shift: float = 1.0, logit_mean: float = 0.0, logit_std: float = 1.0, + *args, **kwargs): + super().__init__(*args, **kwargs) + self.shift = shift + self.logit_mean = logit_mean + self.logit_std = logit_std + + def shift_timesteps(self, steps) -> jnp.ndarray: + return self.shift * steps / (1 + (self.shift - 1) * steps) + + def generate_timesteps(self, batch_size, state: RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]: + state, rng = state.get_random_key() + normal = jax.random.normal(rng, (batch_size,), dtype=self.dtype) + return jax.nn.sigmoid(normal * self.logit_std + self.logit_mean), state + + def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]: + t = self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype)) + return reshape_rates((1 - t, t), shape=shape) + + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + return jnp.ones_like(jnp.asarray(steps, dtype=self.dtype)).reshape(shape) + + def transform_inputs(self, x, steps) -> tuple[jnp.ndarray, jnp.ndarray]: + # The Fourier time embedding is tuned for discrete-style timesteps, so + # the [0, 1] flow time is rescaled into that range + return x, self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype)) * 1000 diff --git a/flaxdiff/schedulers/karras.py b/flaxdiff/schedulers/karras.py index b9e09d5..078165e 100644 --- a/flaxdiff/schedulers/karras.py +++ b/flaxdiff/schedulers/karras.py @@ -17,7 +17,7 @@ def get_sigmas(self, steps) -> jnp.ndarray: sigmas = (self.max_inv_rho + ramp * (self.min_inv_rho - self.max_inv_rho)) ** self.rho return sigmas - def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: sigma = self.get_sigmas(steps) # EDM lambda(sigma) = (sigma^2 + sd^2) / (sigma * sd)^2, written in a # form that needs no epsilon guard (the old guard halved the weight at @@ -64,13 +64,21 @@ def get_sigmas(self, steps) -> jnp.ndarray: return self.sigmas[steps] class EDMNoiseScheduler(KarrasVENoiseScheduler): - def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs): + """Training sigmas drawn from exp(N(P_mean, P_std^2)). + + Defaults are EDM2's (Karras et al. 2024); EDM1's -1.2/1.2 concentrated too + much mass on low noise levels for larger models. Pass them explicitly to + reproduce an EDM1 run. + """ + def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, + P_mean=-0.4, P_std=1.0, *args, **kwargs): super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs) + self.P_mean = P_mean + self.P_std = P_std - def get_sigmas(self, steps, std=1.2, mean=-1.2) -> jnp.ndarray: + def get_sigmas(self, steps) -> jnp.ndarray: space = steps / self.max_timesteps - # space = jax.scipy.special.erfinv(self.erf_sigma_min + steps * (self.erf_sigma_max - self.erf_sigma_min)) - return jnp.exp(space * std + mean) + return jnp.exp(space * self.P_std + self.P_mean) def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]: state, rng = state.get_random_key() diff --git a/tests/test_flow_matching.py b/tests/test_flow_matching.py new file mode 100644 index 0000000..b4efa6f --- /dev/null +++ b/tests/test_flow_matching.py @@ -0,0 +1,249 @@ +"""Flow matching on the linear (rectified flow) path. + +Covers the schedule invariants, the exact velocity round-trip, the claim that +the existing DDIM/Euler samplers already integrate the flow ODE, and a toy +end-to-end run proving the objective actually learns a distribution. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest +from flax import linen as nn + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.predictors import FlowMatchPredictionTransform, get_diffusion_preset +from flaxdiff.samplers.ddim import DDIMSampler +from flaxdiff.samplers.euler import EulerSampler +from flaxdiff.schedulers import FlowMatchingScheduler +from flaxdiff.schedulers.flow import compute_resolution_shift +from flaxdiff.schedulers.common import get_coeff_shapes_tuple +from flaxdiff.utils import RandomMarkovState + +STEPS = jnp.array([0.05, 0.3, 0.6, 0.95]) + + +def test_linear_path_rates(): + schedule = FlowMatchingScheduler() + alpha, sigma = schedule.get_rates(STEPS, shape=(-1,)) + assert jnp.allclose(alpha + sigma, 1.0, atol=1e-6) + assert jnp.allclose(sigma, STEPS, atol=1e-6) + # No input preconditioning on the linear path + assert FlowMatchPredictionTransform().get_input_scale((alpha, sigma)) == 1 + + +def test_endpoints_are_data_and_noise(rng): + schedule = FlowMatchingScheduler() + key0, key1 = jax.random.split(rng) + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + assert jnp.allclose(schedule.add_noise(x0, noise, jnp.zeros((4,))), x0, atol=1e-6) + assert jnp.allclose(schedule.add_noise(x0, noise, jnp.ones((4,))), noise, atol=1e-6) + + +def test_timesteps_are_logit_normal(rng): + schedule = FlowMatchingScheduler(logit_mean=-0.3, logit_std=1.4) + steps, _ = schedule.generate_timesteps(50000, RandomMarkovState(rng)) + assert jnp.all((steps > 0) & (steps < 1)) + logits = jnp.log(steps) - jnp.log1p(-steps) + assert abs(float(jnp.mean(logits)) - (-0.3)) < 0.05 + assert abs(float(jnp.std(logits)) - 1.4) < 0.05 + + +def test_resolution_shift_is_identity_at_one(): + schedule = FlowMatchingScheduler(shift=1.0) + assert jnp.allclose(schedule.shift_timesteps(STEPS), STEPS, atol=1e-7) + + +@pytest.mark.parametrize("shift", [0.5, 1.0, 3.0]) +def test_resolution_shift_is_monotonic_and_fixes_endpoints(shift): + schedule = FlowMatchingScheduler(shift=shift) + t = jnp.linspace(0.0, 1.0, 101) + shifted = schedule.shift_timesteps(t) + assert jnp.all(jnp.diff(shifted) > 0) + assert float(shifted[0]) == pytest.approx(0.0, abs=1e-7) + assert float(shifted[-1]) == pytest.approx(1.0, abs=1e-7) + # A shift above 1 moves every interior timestep towards higher noise + if shift >= 1: + assert jnp.all(shifted >= t - 1e-7) + else: + assert jnp.all(shifted <= t + 1e-7) + + +def test_resolution_shift_grows_with_sequence_length(): + shifts = [compute_resolution_shift(n) for n in (256, 1024, 4096)] + assert shifts == sorted(shifts) + assert shifts[0] == pytest.approx(np.exp(0.5)) + assert shifts[-1] == pytest.approx(np.exp(1.15)) + + +def test_timestep_conditioning_is_scaled_to_the_embedding_range(): + schedule = FlowMatchingScheduler(shift=2.0) + x = jnp.zeros((4, 8, 8, 3)) + _, temb = schedule.transform_inputs(x, STEPS) + assert jnp.allclose(temb, schedule.shift_timesteps(STEPS) * 1000) + + +def test_velocity_roundtrip_is_exact(rng): + schedule = FlowMatchingScheduler() + transform = FlowMatchPredictionTransform() + key0, key1 = jax.random.split(rng) + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + rates = schedule.get_rates(STEPS, get_coeff_shapes_tuple(x0)) + + xt, _, target = transform.forward_diffusion(x0, noise, rates) + assert jnp.allclose(target, noise - x0, atol=1e-6) + + recovered_x0, recovered_noise = transform.backward_diffusion(xt, target, rates) + assert jnp.max(jnp.abs(recovered_x0 - x0)) < 1e-5 + assert jnp.max(jnp.abs(recovered_noise - noise)) < 1e-5 + + +def test_preset_wires_flow_matching(): + for name in ('flow', 'flow_matching'): + train, sample, transform = get_diffusion_preset(name, shift=2.0) + assert isinstance(train, FlowMatchingScheduler) + assert isinstance(sample, FlowMatchingScheduler) + assert isinstance(transform, FlowMatchPredictionTransform) + assert train.shift == 2.0 and sample.shift == 2.0 + + +############################################################################################################ +# The existing samplers already integrate the flow ODE +############################################################################################################ + +class ConstantVelocity(nn.Module): + """Stands in for a trained model; take_next_step never calls it.""" + @nn.compact + def __call__(self, x, temb): + return x + + +def _flow_sampler(sampler_class): + schedule = FlowMatchingScheduler() + return sampler_class( + model=ConstantVelocity(), + noise_schedule=schedule, + model_output_transform=FlowMatchPredictionTransform(), + input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(8, 8, 3), conditions=[]), + guidance_scale=0.0, + ) + + +@pytest.mark.parametrize("sampler_class", [EulerSampler, DDIMSampler]) +def test_sampler_step_is_the_flow_euler_step(sampler_class, rng): + """x_{t+dt} = x_t + u * dt exactly, for the unmodified samplers.""" + sampler = _flow_sampler(sampler_class) + transform = FlowMatchPredictionTransform() + schedule = sampler.noise_schedule + + key0, key1 = jax.random.split(rng) + x_t = jax.random.normal(key0, (4, 8, 8, 3)) + velocity = jax.random.normal(key1, (4, 8, 8, 3)) + current_step = jnp.full((4,), 0.8) + next_step = jnp.full((4,), 0.6) + + rates = schedule.get_rates(current_step, get_coeff_shapes_tuple(x_t)) + x0, eps = transform.backward_diffusion(x_t, velocity, rates) + + stepped, _ = sampler.take_next_step( + current_samples=x_t, + reconstructed_samples=x0, + model_conditioning_inputs=(), + pred_noise=eps, + current_step=current_step, + state=RandomMarkovState(rng), + sample_model_fn=None, + next_step=next_step, + ) + expected = x_t + velocity * (0.6 - 0.8) + assert jnp.max(jnp.abs(stepped - expected)) < 1e-5 + + +############################################################################################################ +# Toy end-to-end: a two-mode gaussian mixture in the plane +############################################################################################################ + +MODE_CENTERS = jnp.array([[-0.5, -0.5], [0.5, 0.5]]) +MODE_STD = 0.08 + + +def sample_mixture(key, n): + """Two well-separated modes in the leading two channels; the third channel + carries no mode information, so the sampler's fixed channel count does not + turn this into a harder problem.""" + mode_key, noise_key = jax.random.split(key) + modes = jax.random.bernoulli(mode_key, 0.5, (n,)).astype(jnp.int32) + centers = jnp.concatenate([MODE_CENTERS[modes], jnp.zeros((n, 1))], axis=-1) + return (centers + MODE_STD * jax.random.normal(noise_key, (n, 3))).reshape(n, 1, 1, 3) + + +class ToyVelocityMLP(nn.Module): + features: int = 128 + + @nn.compact + def __call__(self, x, temb): + t = jnp.reshape(temb, (-1, 1)) / 1000.0 + freqs = jnp.arange(1, 5, dtype=jnp.float32) * jnp.pi + h = jnp.concatenate([x.reshape(x.shape[0], -1), t, jnp.sin(t * freqs), jnp.cos(t * freqs)], axis=-1) + h = nn.swish(nn.Dense(self.features)(h)) + h = nn.swish(nn.Dense(self.features)(h)) + return nn.Dense(3)(h).reshape(x.shape) + + +def test_flow_matching_learns_a_two_mode_mixture(): + schedule, sampling_schedule, transform = get_diffusion_preset('flow') + model = ToyVelocityMLP() + key = jax.random.PRNGKey(0) + params = model.init(key, jnp.zeros((1, 1, 1, 3)), jnp.zeros((1,))) + optimizer = optax.adam(3e-3) + opt_state = optimizer.init(params) + + def loss_fn(params, x0, noise, steps): + rates = schedule.get_rates(steps, get_coeff_shapes_tuple(x0)) + x_t, c_in, target = transform.forward_diffusion(x0, noise, rates) + preds = model.apply(params, *schedule.transform_inputs(x_t * c_in, steps)) + weights = schedule.get_weights(steps, get_coeff_shapes_tuple(x0)) + return jnp.mean(weights * (preds - target) ** 2) + + @jax.jit + def train_step(params, opt_state, rng_state): + rng_state, data_key = rng_state.get_random_key() + rng_state, noise_key = rng_state.get_random_key() + x0 = sample_mixture(data_key, 512) + noise = jax.random.normal(noise_key, x0.shape) + steps, rng_state = schedule.generate_timesteps(512, rng_state) + loss, grads = jax.value_and_grad(loss_fn)(params, x0, noise, steps) + updates, opt_state = optimizer.update(grads, opt_state, params) + return optax.apply_updates(params, updates), opt_state, rng_state, loss + + rng_state = RandomMarkovState(jax.random.PRNGKey(1)) + for _ in range(1500): + params, opt_state, rng_state, loss = train_step(params, opt_state, rng_state) + assert float(loss) < 1.0, "flow matching loss did not come down" + + sampler = EulerSampler( + model=model, + noise_schedule=sampling_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(1, 1, 3), conditions=[]), + guidance_scale=0.0, + ) + samples = sampler.generate_samples( + params, num_samples=2048, resolution=1, diffusion_steps=64, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + ).reshape(-1, 3) + + assignment = samples[:, 0] > 0 + fraction = float(jnp.mean(assignment)) + assert 0.35 < fraction < 0.65, f"modes not balanced: {fraction:.2f}" + + for mode, center in enumerate(MODE_CENTERS): + members = samples[assignment == bool(mode)] + assert jnp.max(jnp.abs(jnp.mean(members[:, :2], axis=0) - center)) < 0.04 + assert abs(float(jnp.std(members[:, :2])) - MODE_STD) < 0.04 + # The third channel is a single zero-centred gaussian, not a mixture + assert abs(float(jnp.mean(samples[:, 2]))) < 0.04 + assert abs(float(jnp.std(samples[:, 2])) - MODE_STD) < 0.04 diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py index 44e05fc..67c745e 100644 --- a/tests/test_schedulers.py +++ b/tests/test_schedulers.py @@ -9,6 +9,11 @@ import jax.numpy as jnp import pytest +from flaxdiff.predictors import ( + EpsilonPredictionTransform, + VPredictionTransform, + get_diffusion_preset, +) from flaxdiff.schedulers import ( CosineNoiseScheduler, LinearNoiseSchedule, @@ -75,12 +80,84 @@ def test_karras_weights_at_sigma_min(): assert jnp.allclose(got, expected, rtol=1e-2) -def test_edm_lognormal_sigma_distribution(rng): - """EDM training sigmas must follow exp(N(-1.2, 1.2^2)) when timesteps=1.""" +@pytest.mark.parametrize("P_mean,P_std", [(-0.4, 1.0), (-1.2, 1.2)]) +def test_edm_lognormal_sigma_distribution(rng, P_mean, P_std): + """EDM training sigmas follow exp(N(P_mean, P_std^2)), defaulting to EDM2.""" from flaxdiff.utils import RandomMarkovState - schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5, P_mean=P_mean, P_std=P_std) steps, _ = schedule.generate_timesteps(20000, RandomMarkovState(rng)) log_sigma = jnp.log(schedule.get_sigmas(steps)) - assert abs(float(jnp.mean(log_sigma)) - (-1.2)) < 0.05 - assert abs(float(jnp.std(log_sigma)) - 1.2) < 0.05 + assert abs(float(jnp.mean(log_sigma)) - P_mean) < 0.05 + assert abs(float(jnp.std(log_sigma)) - P_std) < 0.05 + + +def test_edm_defaults_to_edm2_distribution(): + schedule = EDMNoiseScheduler(1) + assert (schedule.P_mean, schedule.P_std) == (-0.4, 1.0) + + +############################################################################################################ +# min-SNR-gamma loss weighting (Hang et al. 2023) +############################################################################################################ + +MIN_SNR_STEPS = jnp.array([10, 200, 400, 600, 800, 990]) + + +def make_min_snr_schedule(transform, gamma): + return CosineNoiseScheduler(1000, min_snr_gamma=gamma, prediction_transform=transform) + + +def test_min_snr_needs_a_parameterization(): + with pytest.raises(ValueError): + CosineNoiseScheduler(1000, min_snr_gamma=5.0) + + +def test_min_snr_weights_keep_the_requested_shape(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1, 1, 1, 1)).shape == (len(MIN_SNR_STEPS), 1, 1, 1) + assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)).shape == (len(MIN_SNR_STEPS),) + + +def test_min_snr_epsilon_weights_match_the_paper(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + snr = schedule.get_snr(MIN_SNR_STEPS) + expected = jnp.minimum(snr, 5.0) / snr + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5) + + +def test_min_snr_v_weights_match_the_paper(): + schedule = make_min_snr_schedule(VPredictionTransform(), 5.0) + snr = schedule.get_snr(MIN_SNR_STEPS) + expected = jnp.minimum(snr, 5.0) / (snr + 1) + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5) + + +def test_min_snr_weights_are_capped_and_non_increasing_in_snr(): + """The whole point: high-SNR (low noise) steps stop dominating the gradient.""" + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + # ascending timesteps are descending SNR, so weights must be non-decreasing + weights = schedule.get_weights(jnp.arange(1, 1000, 10), shape=(-1,)) + assert jnp.all(jnp.diff(weights) >= -1e-6) + assert jnp.all(weights <= 1.0 + 1e-6) + + +def test_min_snr_gamma_infinity_is_the_unweighted_case(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), float('inf')) + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), 1.0, atol=1e-6) + + +def test_min_snr_is_off_by_default(): + schedule = CosineNoiseScheduler(1000) + assert jnp.allclose( + schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), + schedule.get_schedule_weights(MIN_SNR_STEPS, shape=(-1,)), + ) + + +@pytest.mark.parametrize("name", ['cosine', 'edm', 'karras', 'flow']) +def test_preset_wires_min_snr_into_the_training_schedule_only(name): + train, sample, transform = get_diffusion_preset(name, min_snr_gamma=5.0) + assert train.min_snr_gamma == 5.0 + assert train.prediction_transform is transform + assert sample.min_snr_gamma is None From 8ad7ac9e1d33ac5c8cb5a71f0236b4e3a69f8aaf Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:52 -0400 Subject: [PATCH 3/9] feat(samplers): interval-limited classifier-free guidance Guidance costs diversity at high noise and buys nothing at low noise (Kynkaanniemi et al. 2024). guidance_start/guidance_stop bound the fraction of the trajectory over which it applies; outside the interval the scale drops to 1, which is exactly the plain conditional prediction. The default interval is the full range, so existing samplers are unchanged. Co-Authored-By: Claude Fable 5 --- flaxdiff/samplers/common.py | 27 ++++++-- tests/test_samplers.py | 119 +++++++++++++++++++++++++++++++++++- 2 files changed, 141 insertions(+), 5 deletions(-) diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py index 910202c..7caab39 100644 --- a/flaxdiff/samplers/common.py +++ b/flaxdiff/samplers/common.py @@ -22,31 +22,49 @@ def __init__( model_output_transform: DiffusionPredictionTransform, input_config: DiffusionInputConfig, guidance_scale: float = 0.0, + guidance_start: float = 0.0, + guidance_stop: float = 1.0, autoencoder: AutoEncoder = None, ): """Initialize the diffusion sampler. - + Args: model: Neural network model params: Model parameters noise_schedule: Noise scheduler model_output_transform: Transform for model predictions guidance_scale: Scale for classifier-free guidance (0.0 means disabled) + guidance_start: Fraction of the trajectory after which guidance turns on + guidance_stop: Fraction of the trajectory after which guidance turns off autoencoder: Optional autoencoder for latent diffusion """ self.model = model self.noise_schedule = noise_schedule self.model_output_transform = model_output_transform self.guidance_scale = guidance_scale + self.guidance_start = guidance_start + self.guidance_stop = guidance_stop self.autoencoder = autoencoder self.input_config = input_config - + unconditionals = input_config.get_unconditionals() - + if self.guidance_scale > 0: # Classifier free guidance print("Using classifier-free guidance") + def guidance_at(t): + """Interval-limited guidance (Kynkaanniemi et al. 2024). + + Guidance hurts at high noise and buys nothing at low noise, so + outside [guidance_start, guidance_stop] the scale drops to 1, + which is exactly the plain conditional prediction. Progress runs + from 0 at pure noise to 1 at the clean sample. + """ + progress = 1.0 - t / self.noise_schedule.max_timesteps + inside = (progress >= guidance_start) & (progress <= guidance_stop) + return jnp.where(inside, guidance_scale, 1.0) + def sample_model(params, x_t, t, *conditioning_inputs): # Concatenate unconditional and conditional inputs x_t_cat = jnp.concatenate([x_t] * 2, axis=0) @@ -71,7 +89,8 @@ def sample_model(params, x_t, t, *conditioning_inputs): # Split model output into unconditional and conditional parts model_output_cond, model_output_uncond = jnp.split(model_output, 2, axis=0) - model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond) + scale = guidance_at(t).reshape(get_coeff_shapes_tuple(model_output_cond)) + model_output = model_output_uncond + scale * (model_output_cond - model_output_uncond) x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule) return x_0, eps, model_output diff --git a/tests/test_samplers.py b/tests/test_samplers.py index 9156c91..3fa8507 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -7,12 +7,15 @@ shape bugs in one place, with no training involved. """ +from dataclasses import dataclass + import jax import jax.numpy as jnp import pytest from flax import linen as nn -from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.inputs import ConditionalInputConfig, DiffusionInputConfig +from flaxdiff.inputs.encoders import ConditioningEncoder from flaxdiff.predictors import EpsilonPredictionTransform, KarrasPredictionTransform from flaxdiff.samplers.ddim import DDIMSampler from flaxdiff.samplers.ddpm import DDPMSampler, SimpleDDPMSampler @@ -186,3 +189,117 @@ def test_multistep_dpm_reentrant(): assert jnp.allclose(first, second, atol=1e-5) +############################################################################################################ +# Interval-limited classifier-free guidance (Kynkaanniemi et al. 2024) +############################################################################################################ + +@dataclass +class LabelEncoder(ConditioningEncoder): + """Smallest conditioning seam that still gives CFG a real conditional and + unconditional branch to interpolate between.""" + @property + def key(self): + return "label" + + def tokenize(self, data): + return jnp.asarray(data, dtype=jnp.float32).reshape(-1, 1) + + def encode_from_tokens(self, tokens): + return tokens + + def serialize(self): + return {} + + @staticmethod + def deserialize(serialized_config): + raise NotImplementedError + + +conditional_input_config = DiffusionInputConfig( + sample_data_key="image", + sample_data_shape=(8, 8, 3), + conditions=[ + ConditionalInputConfig( + encoder=LabelEncoder(model=None, tokenizer=None), + conditioning_data_key="label", + unconditional_input=0.0, + ) + ], +) + + +class ConditionalVPOracle(nn.Module): + """VPOracle offset by the label, so the guided output reads back the scale + that was actually applied.""" + schedule: CosineNoiseScheduler + + @nn.compact + def __call__(self, x, temb, label): + alpha, sigma = self.schedule.get_rates(temb, get_coeff_shapes_tuple(x)) + eps = x * sigma / (alpha**2 * DATA_STD**2 + sigma**2) + return eps + label.reshape(get_coeff_shapes_tuple(x)) + + +def make_guided_sampler(**kwargs): + schedule = CosineNoiseScheduler(1000) + model = ConditionalVPOracle(schedule=schedule) + sampler = DDIMSampler( + model=model, + noise_schedule=schedule, + model_output_transform=EpsilonPredictionTransform(), + input_config=conditional_input_config, + **kwargs, + ) + params = model.init( + jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,)), jnp.ones((1, 1)) + ) + return params, sampler + + +@pytest.mark.parametrize("progress,inside", [(0.1, False), (0.5, True), (0.9, False)]) +def test_interval_cfg_applies_only_inside_the_interval(progress, inside): + params, full = make_guided_sampler(guidance_scale=3.0) + _, interval = make_guided_sampler(guidance_scale=3.0, guidance_start=0.4, guidance_stop=0.6) + _, unguided = make_guided_sampler(guidance_scale=0.0) + + x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3)) + labels = jnp.full((4, 1), 0.7) + t = jnp.full((4,), (1.0 - progress) * 1000) + + def output(sampler): + return sampler.sample_model(params, x_t, t, labels)[2] + + matches_full = bool(jnp.allclose(output(interval), output(full), atol=1e-5)) + matches_unguided = bool(jnp.allclose(output(interval), output(unguided), atol=1e-5)) + assert matches_full is inside + assert matches_unguided is not inside + + +def test_interval_cfg_defaults_to_the_full_range(): + params, default = make_guided_sampler(guidance_scale=3.0) + _, explicit = make_guided_sampler(guidance_scale=3.0, guidance_start=0.0, guidance_stop=1.0) + + x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3)) + labels = jnp.full((4, 1), 0.7) + t = jnp.full((4,), 500.0) + assert jnp.allclose( + default.sample_model(params, x_t, t, labels)[2], + explicit.sample_model(params, x_t, t, labels)[2], + atol=1e-6, + ) + + +def test_empty_guidance_interval_generates_the_unguided_samples(): + params, empty = make_guided_sampler(guidance_scale=3.0, guidance_start=0.9, guidance_stop=0.1) + _, unguided = make_guided_sampler(guidance_scale=0.0) + + def run(sampler): + return sampler.generate_samples( + params, num_samples=16, resolution=8, diffusion_steps=25, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + model_conditioning_inputs=(jnp.full((16, 1), 0.7),), + ) + + assert jnp.allclose(run(empty), run(unguided), atol=1e-5) + + From aadc65988e3cda9c3b8a5dc943a36ddb9f3deab2 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:52 -0400 Subject: [PATCH 4/9] feat(metrics): FID on the vendored InceptionV3 The jax-fid port in metrics/inception.py has been sitting unused. Wire it into an EvaluationMetric the same way metrics/images.py builds the CLIP ones, sharing a single cached copy of the weights. Per-batch FID is noisy and only meaningful as a trend across checkpoints, which the docstring says outright. Co-Authored-By: Claude Fable 5 --- flaxdiff/metrics/fid.py | 75 +++++++++++++++++++++++++++++++++++++++++ tests/test_metrics.py | 53 +++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+) create mode 100644 flaxdiff/metrics/fid.py create mode 100644 tests/test_metrics.py diff --git a/flaxdiff/metrics/fid.py b/flaxdiff/metrics/fid.py new file mode 100644 index 0000000..d11d13b --- /dev/null +++ b/flaxdiff/metrics/fid.py @@ -0,0 +1,75 @@ +from .common import EvaluationMetric +import jax +import jax.numpy as jnp +import numpy as np + + +# The FID InceptionV3 is ~90MB of weights; one copy is enough for every metric +# built from this module. +_inception_cache: dict = {} + + +def _get_inception(): + """Cached (model, params) for the pool3 feature extractor.""" + if 'inception' not in _inception_cache: + from .inception import InceptionV3 + print("[metrics] Loading InceptionV3 FID weights (cached for reuse)...") + model = InceptionV3(pretrained=True) + params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 299, 299, 3))) + _inception_cache['inception'] = (model, params) + return _inception_cache['inception'] + + +def frechet_distance(mu_a, sigma_a, mu_b, sigma_b, eps=1e-6) -> float: + """Frechet distance between two multivariate gaussians. + + Runs on the host through scipy: the matrix square root of the covariance + product has no jax equivalent, and FID is computed once per validation + batch so the transfer is irrelevant. + """ + from scipy import linalg + + mu_a, mu_b = np.atleast_1d(mu_a), np.atleast_1d(mu_b) + sigma_a, sigma_b = np.atleast_2d(sigma_a), np.atleast_2d(sigma_b) + + covmean, _ = linalg.sqrtm(sigma_a.dot(sigma_b), disp=False) + if not np.isfinite(covmean).all(): + # Singular product covariance, nudge the diagonal as in the reference + # implementations rather than returning a nan + offset = np.eye(sigma_a.shape[0]) * eps + covmean = linalg.sqrtm((sigma_a + offset).dot(sigma_b + offset)) + if np.iscomplexobj(covmean): + covmean = covmean.real + + diff = mu_a - mu_b + return float(diff.dot(diff) + np.trace(sigma_a) + np.trace(sigma_b) - 2 * np.trace(covmean)) + + +def _gaussian_stats(activations): + activations = np.asarray(activations, dtype=np.float64) + return activations.mean(axis=0), np.cov(activations, rowvar=False) + + +def get_fid_metric(): + """FID between the generated batch and the real batch, lower is better. + + Per-batch FID is noisy at typical validation batch sizes and is only + meaningful as a relative trend across checkpoints, not as a headline number + comparable to published FID-50k. + """ + model, params = _get_inception() + + @jax.jit + def activations(images): + # Inception wants [-1, 1] at 299x299; pool3 output is [B, 1, 1, 2048] + resized = jax.image.resize(images, (images.shape[0], 299, 299, 3), method='bilinear') + features = model.apply(params, resized, train=False) + return features.reshape(features.shape[0], -1) + + def fid_metric(generated: jnp.ndarray, batch): + original = (jnp.asarray(batch['image'], dtype=jnp.float32) - 127.5) / 127.5 + mu_gen, sigma_gen = _gaussian_stats(activations(generated)) + mu_real, sigma_real = _gaussian_stats(activations(original)) + return frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real) + + return EvaluationMetric(function=fid_metric, name='fid') diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..3801da9 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,53 @@ +"""Evaluation metric tests. + +The Frechet distance itself is checked against closed forms that need no +weights; the end-to-end InceptionV3 path downloads the FID checkpoint and is +network-marked. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from flaxdiff.metrics.common import EvaluationMetric +from flaxdiff.metrics.fid import frechet_distance, get_fid_metric + + +def test_frechet_distance_of_a_distribution_with_itself_is_zero(rng): + features = np.asarray(jax.random.normal(rng, (256, 16))) + mu, sigma = features.mean(axis=0), np.cov(features, rowvar=False) + assert frechet_distance(mu, sigma, mu, sigma) == pytest.approx(0.0, abs=1e-6) + + +def test_frechet_distance_of_shifted_gaussians_is_the_squared_mean_gap(): + sigma = np.eye(8) + mu_a = np.zeros(8) + mu_b = np.full(8, 0.5) + assert frechet_distance(mu_a, sigma, mu_b, sigma) == pytest.approx(8 * 0.25, abs=1e-6) + + +def test_frechet_distance_grows_with_covariance_mismatch(): + mu = np.zeros(8) + identity = np.eye(8) + near = frechet_distance(mu, identity, mu, identity * 1.5) + far = frechet_distance(mu, identity, mu, identity * 4.0) + assert 0 < near < far + + +@pytest.mark.network +def test_fid_metric_scores_real_images_better_than_noise(rng): + metric = get_fid_metric() + assert isinstance(metric, EvaluationMetric) + assert metric.name == 'fid' and metric.higher_is_better is False + + key_real, key_noise = jax.random.split(rng) + real = jax.random.randint(key_real, (8, 64, 64, 3), 0, 256, dtype=jnp.int32).astype(jnp.uint8) + batch = {'image': real} + + # Generated samples live in [-1, 1]; the same images should score far + # closer to the batch than unrelated noise does + matching = metric.function((jnp.asarray(real, jnp.float32) - 127.5) / 127.5, batch) + unrelated = metric.function(jax.random.normal(key_noise, (8, 64, 64, 3)), batch) + assert np.isfinite(matching) and np.isfinite(unrelated) + assert matching < unrelated From 29d1a2debb2e4731caa404c15aca5c9fce70dc01 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:58 -0400 Subject: [PATCH 5/9] feat(autoencoder): per-dataset latent normalization on the shared seam Latents are normalized as (z - latent_shift) * latent_scale and inverted on decode, the SD3 convention, so a caller can hand in a dataset's own mean and 1/std instead of relying on the VAE's single fixed constant. The SD VAE's scaling_factor now rides on that same seam rather than being applied inside its jitted frame functions, so there is one normalization path and the defaults reproduce the previous behavior exactly. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/autoencoder/autoencoder.py | 22 ++++++-- flaxdiff/models/autoencoder/diffusers.py | 19 ++++--- tests/test_vae.py | 58 ++++++++++++++++++++-- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/flaxdiff/models/autoencoder/autoencoder.py b/flaxdiff/models/autoencoder/autoencoder.py index a307da1..233eaec 100644 --- a/flaxdiff/models/autoencoder/autoencoder.py +++ b/flaxdiff/models/autoencoder/autoencoder.py @@ -10,11 +10,19 @@ @dataclass class AutoEncoder(ABC): """Base class for autoencoder models with video support. - + This class defines the interface for autoencoders and provides video handling functionality, allowing child classes to focus on implementing the core encoding/decoding for individual frames. + + Latents are normalized as (z - latent_shift) * latent_scale on the way out + and inverted on the way in, the SD3 convention. The defaults are the + identity; set them to the dataset's own latent mean and 1/std so the + diffusion model sees roughly unit-variance, zero-mean inputs. """ + latent_shift: float = 0.0 + latent_scale: float = 1.0 + @abstractmethod def __encode__(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray: """Abstract method for encoding a batch of images. @@ -72,13 +80,15 @@ def encode(self, x: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa # Encode all frames latent = self.__encode__(x_reshaped, key=key, **kwargs) - + # Reshape back to include temporal dimension [B, T, h, w, c] latent_shape = latent.shape - return latent.reshape(batch_size, seq_len, *latent_shape[1:]) + latent = latent.reshape(batch_size, seq_len, *latent_shape[1:]) else: # Standard image processing - return self.__encode__(x, key=key, **kwargs) + latent = self.__encode__(x, key=key, **kwargs) + + return (latent - self.latent_shift) * self.latent_scale def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs) -> jnp.ndarray: """Decode latent representations, with special handling for video data. @@ -95,9 +105,11 @@ def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa Returns: Decoded output with the same batch and temporal dimensions as input """ + z = z / self.latent_scale + self.latent_shift + # Check for video data (5D tensor) is_video = len(z.shape) == 5 - + if is_video: # Extract dimensions for reshaping batch_size, seq_len, height, width, channels = z.shape diff --git a/flaxdiff/models/autoencoder/diffusers.py b/flaxdiff/models/autoencoder/diffusers.py index 43ee03c..ab8015b 100644 --- a/flaxdiff/models/autoencoder/diffusers.py +++ b/flaxdiff/models/autoencoder/diffusers.py @@ -12,7 +12,8 @@ """ class StableDiffusionVAE(AutoEncoder): - def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16): + def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16, + latent_shift=None, latent_scale=None): pretrained = load_pretrained_vae(modelname) config = pretrained["config"] @@ -61,10 +62,14 @@ def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=dtype, ) - # Older VAE configs predate the scaling_factor key; 0.18215 is the SD default - scaling_factor = config.get("scaling_factor", 0.18215) - print(f"Scaling factor: {scaling_factor}") - + # The VAE's own latent normalization rides on the AutoEncoder seam, so a + # caller can override it with per-dataset statistics without a second + # scaling path. Older configs predate these keys; 0.0 and 0.18215 are + # the SD defaults. + self.latent_shift = config.get("shift_factor", 0.0) if latent_shift is None else latent_shift + self.latent_scale = config.get("scaling_factor", 0.18215) if latent_scale is None else latent_scale + print(f"Latent shift: {self.latent_shift}, latent scale: {self.latent_scale}") + def encode_single_frame(images, rngkey: jax.random.PRNGKey = None): latents = enc.apply({"params": params['encoder']}, images, deterministic=True) latents = quant_conv.apply({"params": params['quant_conv']}, latents) @@ -75,11 +80,9 @@ def encode_single_frame(images, rngkey: jax.random.PRNGKey = None): latents = mean + std * jax.random.normal(rngkey, mean.shape, dtype=mean.dtype) else: latents, _ = jnp.split(latents, 2, axis=-1) - latents *= scaling_factor return latents - + def decode_single_frame(latents): - latents = (1.0 / scaling_factor) * latents latents = post_quant_conv.apply({"params": params['post_quant_conv']}, latents) return dec.apply({"params": params['decoder']}, latents) diff --git a/tests/test_vae.py b/tests/test_vae.py index 96904ec..3b648e0 100644 --- a/tests/test_vae.py +++ b/tests/test_vae.py @@ -1,7 +1,7 @@ -"""Vendored Stable Diffusion VAE tests. +"""Vendored Stable Diffusion VAE tests, plus the latent normalization seam. -Marked as network tests: they download the pretrained weights from the -HuggingFace Hub on first run. Excluded in CI (-m "not network"). +Everything touching the pretrained weights is network-marked: they download +from the HuggingFace Hub on first run. Excluded in CI (-m "not network"). """ import jax @@ -9,7 +9,47 @@ import numpy as np import pytest -pytestmark = pytest.mark.network +from flaxdiff.models.autoencoder import AutoEncoder + + +class IdentityAutoEncoder(AutoEncoder): + """Latents are the input, so only the normalization seam is under test.""" + def __encode__(self, x, **kwargs): + return x + + def __decode__(self, z, **kwargs): + return z + + def serialize(self): + return {} + + +def test_latent_normalization_defaults_to_the_identity(rng): + autoencoder = IdentityAutoEncoder() + x = jax.random.normal(rng, (2, 8, 8, 4)) + assert jnp.allclose(autoencoder.encode(x), x) + assert jnp.allclose(autoencoder.decode(x), x) + + +@pytest.mark.parametrize("shape", [(2, 8, 8, 4), (2, 3, 8, 8, 4)]) +def test_latent_normalization_shifts_and_scales_roundtrip(rng, shape): + """SD3-style shift+scale: latents come out centred and rescaled, and + decoding inverts it exactly, for images and for video.""" + autoencoder = IdentityAutoEncoder(latent_shift=0.25, latent_scale=4.0) + x = jax.random.normal(rng, shape) + latent = autoencoder.encode(x) + assert jnp.allclose(latent, (x - 0.25) * 4.0, atol=1e-6) + assert jnp.allclose(autoencoder.decode(latent), x, atol=1e-5) + + +def test_latent_normalization_whitens_a_known_distribution(rng): + """The point of per-dataset stats: the diffusion model sees zero mean and + unit variance instead of whatever the encoder happens to produce.""" + x = 3.0 + 5.0 * jax.random.normal(rng, (4096, 1, 1, 4)) + autoencoder = IdentityAutoEncoder(latent_shift=float(jnp.mean(x)), latent_scale=1.0 / float(jnp.std(x))) + latent = autoencoder.encode(x) + assert abs(float(jnp.mean(latent))) < 1e-4 + assert abs(float(jnp.std(latent)) - 1.0) < 1e-4 @pytest.fixture(scope="module") @@ -18,11 +58,21 @@ def vae(): return StableDiffusionVAE(dtype=jnp.float32) +@pytest.mark.network def test_vae_shapes(vae): assert vae.downscale_factor == 8 assert vae.latent_channels == 4 +@pytest.mark.network +def test_vae_uses_the_latent_normalization_seam(vae): + """The SD scaling factor rides on the shared seam, so there is one + normalization path a caller can override with dataset statistics.""" + assert vae.latent_scale == pytest.approx(0.18215) + assert vae.latent_shift == 0.0 + + +@pytest.mark.network def test_vae_roundtrip_reconstructs(vae, rng): # A smooth image should survive the encode/decode roundtrip well ramp = jnp.linspace(-0.8, 0.8, 64) From 40f69d7149b84817338d6bc1a5b74b86849940d2 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:58 -0400 Subject: [PATCH 6/9] feat(training): expose the flow matching, min-SNR and FID options --noise_schedule gains flow/flow_matching, plus --min_snr_gamma, --flow_shift, and 'fid' as a validation metric. Co-Authored-By: Claude Fable 5 --- training.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/training.py b/training.py index 9ede2e7..b42bb43 100644 --- a/training.py +++ b/training.py @@ -101,7 +101,11 @@ def boolean_string(s): default='/home/mrwhite0racle/gcs_mount', help="Dataset location path") parser.add_argument('--noise_schedule', type=str, default='edm', - choices=['cosine', 'karras', 'edm'], help='Noise schedule') + choices=['cosine', 'karras', 'edm', 'flow', 'flow_matching'], help='Noise schedule') +parser.add_argument('--min_snr_gamma', type=float, default=None, + help='min-SNR-gamma loss weighting (Hang et al. 2023). 5.0 is the paper default; unset keeps the schedule own weighting.') +parser.add_argument('--flow_shift', type=float, default=1.0, + help='Resolution shift for the flow matching schedule. See flaxdiff.schedulers.flow.compute_resolution_shift.') # Any name from flaxdiff.models.registry, optionally with +2d/+hilbert/+zigzag # suffixes; validated by build_model against the registry itself @@ -469,6 +473,10 @@ def main(args): from flaxdiff.metrics.images import get_clip_score_metric print("Using CLIPScore (val/clip_score, higher is better) for validation") eval_metrics.append(get_clip_score_metric()) + if 'fid' in args.val_metrics: + from flaxdiff.metrics.fid import get_fid_metric + print("Using per-batch FID (val/fid) for validation") + eval_metrics.append(get_fid_metric()) CONFIG = { "model": model_config, @@ -493,7 +501,9 @@ def main(args): batches = batches if args.steps_per_epoch is None else args.steps_per_epoch - train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(args.noise_schedule) + train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset( + args.noise_schedule, shift=args.flow_shift, min_snr_gamma=args.min_snr_gamma, + ) if args.experiment_name is not None: experiment_name = args.experiment_name From f21bdd6c2fe990518603e88fc85f2d7c591b66f8 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 02:10:55 -0400 Subject: [PATCH 7/9] feat(jepa): I-JEPA and V-JEPA on the shared DiT blocks The encoders are arrangements, not new transformers. JepaEncoder is PatchSequenceEmbed plus a stack of unmodulated ModulatedBlocks; JepaVideoEncoder is the same blocks in VideoDiT's factorized spatial-temporal layout; the predictor is a narrower stack of the same blocks over context embeddings, mask tokens and target positions. mixer patterns carry straight through, so an encoder can be all-attention, all-S5 or any ratio between - the hybrid SSM encoder is a config string, not a second implementation. Masking follows the paper's multi-block strategy, resolved against the static patch grid at construction so every step produces identically-shaped masks under jit. The block area is the value inside the scale range with the most (h, w) factorizations in the aspect range, so shape and position stay random per sample; the context is a random subset of the complement of the target union, sized so it is always available. Video masks are tubelets: one spatial mask applied to every frame, which leaves the factorized layout intact. Position comes from the 2D sincos signal alone, with RoPE at identity: a masked subsequence has no meaningful 1D index, and sincos travels with the token when it is gathered. Time is untouched by masking, so the temporal half of the video stack keeps real RoPE. Collapse telemetry ships with the loss rather than after it. repr_std and repr_cov_offdiag go into the aux dict on every step, and the trainer already forwards aux to wandb, so a run that quietly agrees on a constant shows it in the curves. Both are tested against a degenerate encoder, not just asserted. Also here: frozen-encoder linear and k-NN probes as EvaluationMetrics that the existing validation loop runs over the target encoder's embeddings, the models in the registry, and training_jepa.py. Two fixes fell out. The validation loop's numpy import lived inside an `if i == 0 and wandb` branch while np.mean ran outside it, so any run with eval metrics and no wandb died in validation. And the wandb artifact name was hardcoded to "diffusion--res", which would have collided across objectives; it now takes the objective's tag. Co-Authored-By: Claude Fable 5 --- flaxdiff/data/sources/images.py | 2 + flaxdiff/jepa/__init__.py | 4 + flaxdiff/jepa/masking.py | 119 ++++++ flaxdiff/jepa/models.py | 316 +++++++++++++++ flaxdiff/jepa/objective.py | 147 +++++++ flaxdiff/jepa/probes.py | 79 ++++ flaxdiff/models/dit_common.py | 37 +- flaxdiff/models/registry.py | 4 + flaxdiff/models/ssm_dit.py | 20 +- flaxdiff/trainer/general_diffusion_trainer.py | 86 ++-- flaxdiff/trainer/objectives.py | 5 +- tests/test_jepa.py | 367 ++++++++++++++++++ tests/test_objectives.py | 2 +- training_jepa.py | 269 +++++++++++++ 14 files changed, 1364 insertions(+), 93 deletions(-) create mode 100644 flaxdiff/jepa/__init__.py create mode 100644 flaxdiff/jepa/masking.py create mode 100644 flaxdiff/jepa/models.py create mode 100644 flaxdiff/jepa/objective.py create mode 100644 flaxdiff/jepa/probes.py create mode 100644 tests/test_jepa.py create mode 100644 training_jepa.py diff --git a/flaxdiff/data/sources/images.py b/flaxdiff/data/sources/images.py index 569373c..5616471 100644 --- a/flaxdiff/data/sources/images.py +++ b/flaxdiff/data/sources/images.py @@ -187,6 +187,8 @@ def map(self, element) -> Dict[str, jnp.array]: results = self.tokenize(caption) return { "image": image, + # the class index, which the JEPA linear/kNN probes score against + "label": np.int32(element['label']), "text": { "input_ids": results['input_ids'][0], "attention_mask": results['attention_mask'][0], diff --git a/flaxdiff/jepa/__init__.py b/flaxdiff/jepa/__init__.py new file mode 100644 index 0000000..1f597a3 --- /dev/null +++ b/flaxdiff/jepa/__init__.py @@ -0,0 +1,4 @@ +from .masking import MultiBlockMask, multi_block_mask +from .models import JepaEncoder, JepaVideoEncoder, JepaPredictor, TokenStack, FactorizedTokenStack +from .objective import JepaObjective, representation_health +from .probes import get_linear_probe_metric, get_knn_probe_metric, linear_probe, knn_probe diff --git a/flaxdiff/jepa/masking.py b/flaxdiff/jepa/masking.py new file mode 100644 index 0000000..aaecc6a --- /dev/null +++ b/flaxdiff/jepa/masking.py @@ -0,0 +1,119 @@ +"""I-JEPA multi-block masking over a patch grid. + +The paper samples several target blocks with a random scale and aspect ratio +and takes the context to be what is left. Under jit every mask has to come out +the same shape on every step, so the geometry is resolved once, at +construction, against the static patch grid: + + - the block area is fixed to the value inside the configured scale range that + admits the most (height, width) factorizations within the aspect ratio + range, so each sample still draws a genuinely different block shape and + position, and every target block has exactly the same token count; + - the context is a uniformly random subset of the complement of the target + union, sized S - num_targets * area. That size is always available (the + union is at most num_targets disjoint blocks), and when the blocks happen + to overlap the context simply drops the surplus rather than changing shape. + +Indices refer to positions in the scan-ordered token sequence that +PatchSequenceEmbed produces, and come out sorted so that an SSM mixer scans +them in a meaningful order. +""" + +from dataclasses import dataclass +from typing import Tuple + +import jax +import jax.numpy as jnp + + +def _factorizations(area: int, grid: Tuple[int, int], aspect: Tuple[float, float]): + """(h, w) pairs of the given area that fit the grid and the aspect range.""" + H_P, W_P = grid + pairs = [] + for h in range(1, min(H_P, area) + 1): + if area % h: + continue + w = area // h + if w <= W_P and aspect[0] <= h / w <= aspect[1]: + pairs.append((h, w)) + return pairs + + +@dataclass(frozen=True) +class MultiBlockMask: + """Static mask geometry for one patch grid, plus the sampler over it.""" + grid: Tuple[int, int] + num_targets: int + block_shapes: Tuple[Tuple[int, int], ...] + num_context: int + + @property + def num_patches(self) -> int: + return self.grid[0] * self.grid[1] + + @property + def block_area(self) -> int: + h, w = self.block_shapes[0] + return h * w + + def sample(self, rng: jax.Array, batch_size: int): + """Draw (context_idx [B, num_context], target_idx [B, M, block_area]).""" + H_P, W_P = self.grid + S = self.num_patches + M = self.num_targets + + heights = jnp.asarray([h for h, _ in self.block_shapes]) + widths = jnp.asarray([w for _, w in self.block_shapes]) + # flat offsets of a block's tokens from its top-left corner, per shape + offsets = jnp.asarray([[(i // w) * W_P + (i % w) for i in range(h * w)] + for h, w in self.block_shapes]) + + shape_key, pos_key, context_key = jax.random.split(rng, 3) + choice = jax.random.randint(shape_key, (batch_size, M), 0, len(self.block_shapes)) + corner = jax.random.uniform(pos_key, (batch_size, M, 2)) + top = jnp.floor(corner[..., 0] * (H_P - heights[choice] + 1)).astype(jnp.int32) + left = jnp.floor(corner[..., 1] * (W_P - widths[choice] + 1)).astype(jnp.int32) + target_idx = (top * W_P + left)[..., None] + offsets[choice] + + is_target = jnp.zeros((batch_size, S), dtype=bool).at[ + jnp.arange(batch_size)[:, None], target_idx.reshape(batch_size, -1)].set(True) + # target tokens sort last, so the first num_context entries of a random + # ordering are a uniform sample of the complement + order = jax.random.uniform(context_key, (batch_size, S)) + is_target + context_idx = jnp.sort(jnp.argsort(order, axis=1)[:, :self.num_context], axis=1) + return context_idx, target_idx + + +def multi_block_mask( + grid: Tuple[int, int], + num_targets: int = 4, + scale: Tuple[float, float] = (0.15, 0.2), + aspect: Tuple[float, float] = (0.75, 1.5), +) -> MultiBlockMask: + """Resolve the I-JEPA mask geometry for a patch grid.""" + S = grid[0] * grid[1] + candidates = [ + (area, _factorizations(area, grid, aspect)) + for area in range(max(1, int(jnp.ceil(scale[0] * S))), int(scale[1] * S) + 1) + ] + candidates = [(area, shapes) for area, shapes in candidates if shapes] + if not candidates: + raise ValueError( + f"No block of scale {scale} on a {grid} grid has an aspect ratio in " + f"{aspect}; widen one of the ranges or use a finer patch grid.") + + midpoint = 0.5 * (scale[0] + scale[1]) * S + area, shapes = min(candidates, key=lambda c: (-len(c[1]), abs(c[0] - midpoint))) + + num_context = S - num_targets * area + if num_context <= 0: + raise ValueError( + f"{num_targets} target blocks of {area} tokens leave no context on a " + f"{grid} grid ({S} tokens).") + + return MultiBlockMask( + grid=grid, + num_targets=num_targets, + block_shapes=tuple(shapes), + num_context=num_context, + ) diff --git a/flaxdiff/jepa/models.py b/flaxdiff/jepa/models.py new file mode 100644 index 0000000..ca931b2 --- /dev/null +++ b/flaxdiff/jepa/models.py @@ -0,0 +1,316 @@ +"""JEPA encoders and predictors, arranged from the shared DiT machinery. + +A JEPA encoder is the DiT sandwich without the diffusion parts: patchify with +the 2D sincos signal, run unmodulated ModulatedBlocks over the tokens, norm. +There is no timestep to condition on, so the blocks run in their unmodulated +(plain pre-norm) mode, and the mixer is still pluggable - mixer patterns with +'ssm' give a linear-time S5 encoder. + +Position never comes from RoPE here. Both the encoder and the predictor work +on a masked subset of the sequence, where a token's index in the sequence is +not its position on the grid, so RoPE is left at identity and position is +carried entirely by the 2D sincos embedding that travels with each token. +""" + +import jax.numpy as jnp +from flax import linen as nn +from typing import Optional, Sequence, Tuple +from flax.typing import Dtype, PrecisionLike + +from ..models.dit_common import ( + PatchSequenceEmbed, ModulatedBlock, build_block_pattern, + identity_rope_freqs, scan_ordered_pos_embed, +) +from ..models.vit_common import RotaryEmbedding + + +def gather_tokens(tokens, idx): + """Select tokens per sample: [B, S, F] and [B, N] -> [B, N, F].""" + return jnp.take_along_axis(tokens, idx[..., None], axis=-2) + + +class TokenStack(nn.Module): + """A stack of unmodulated blocks over a token sequence.""" + features: int + num_layers: int + num_heads: int + mlp_ratio: int = 4 + ssm_attention_ratio: str = "all-attn" + block_pattern: Optional[Sequence[str]] = None + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + qk_norm: bool = False + attention_impl: Optional[str] = None + + def setup(self): + pattern = build_block_pattern( + self.num_layers, self.ssm_attention_ratio, self.block_pattern) + self.blocks = [ + ModulatedBlock( + features=self.features, + num_heads=self.num_heads, + mixer='ssm' if kind == 'ssm' else 'attention', + modulated=False, + mlp_ratio=self.mlp_ratio, + dropout_rate=self.dropout_rate, + dtype=self.dtype, + precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, + qk_norm=self.qk_norm, + attention_impl=self.attention_impl, + ssm_state_dim=self.ssm_state_dim, + bidirectional_ssm=self.bidirectional_ssm, + name=f"block_{i}", + ) for i, kind in enumerate(pattern) + ] + + def __call__(self, tokens, freqs_cis=None, train: bool = False): + if freqs_cis is None: + freqs_cis = identity_rope_freqs(tokens.shape[-2], self.features // self.num_heads) + for block in self.blocks: + tokens = block(tokens, conditioning=None, freqs_cis=freqs_cis, train=train) + return tokens + + +class FactorizedTokenStack(nn.Module): + """Spatial then temporal blocks over [B, T, N, F], as in VideoDiT. + + Time is a real 1D axis that masking never touches, so the temporal half + keeps genuine RoPE while the spatial half runs at identity. + """ + features: int + num_layers: int + num_heads: int + max_frames: int = 1024 + mlp_ratio: int = 4 + ssm_attention_ratio: str = "all-attn" + block_pattern: Optional[Sequence[str]] = None + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + qk_norm: bool = False + attention_impl: Optional[str] = None + + def setup(self): + def stack(name, num_layers, ratio): + return TokenStack( + features=self.features, num_layers=num_layers, num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, ssm_attention_ratio=ratio, + ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm, + dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, name=name) + + # one spatial and one temporal block per layer, built as single-block + # stacks so the two halves can be interleaved + pattern = build_block_pattern( + self.num_layers, self.ssm_attention_ratio, self.block_pattern) + self.spatial = [stack(f"spatial_{i}", 1, "all-ssm" if kind == 'ssm' else "all-attn") + for i, kind in enumerate(pattern)] + self.temporal = [stack(f"temporal_{i}", 1, "all-attn") for i in range(self.num_layers)] + self.temporal_rope = RotaryEmbedding( + dim=self.features // self.num_heads, max_seq_len=self.max_frames, + name="temporal_rope") + + def __call__(self, tokens, train: bool = False): + B, T, N, F = tokens.shape + freqs_temporal = self.temporal_rope(seq_len=T) + + tokens = tokens.reshape(B * T, N, F) + for spatial, temporal in zip(self.spatial, self.temporal): + tokens = spatial(tokens, train=train) + tokens = tokens.reshape(B, T, N, F).transpose(0, 2, 1, 3).reshape(B * N, T, F) + tokens = temporal(tokens, freqs_cis=freqs_temporal, train=train) + tokens = tokens.reshape(B, N, T, F).transpose(0, 2, 1, 3).reshape(B * T, N, F) + return tokens.reshape(B, T, N, F) + + +class JepaEncoder(nn.Module): + """ViT over an image, optionally restricted to a subset of its patches.""" + patch_size: int = 16 + emb_features: int = 384 + num_layers: int = 12 + num_heads: int = 6 + mlp_ratio: int = 4 + ssm_attention_ratio: str = "all-attn" + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + qk_norm: bool = False + attention_impl: Optional[str] = None + use_hilbert: bool = False + use_zigzag: bool = False + + @property + def scan_order(self): + assert not (self.use_hilbert and self.use_zigzag), \ + "use_hilbert and use_zigzag are mutually exclusive" + return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster' + + def setup(self): + self.embed = PatchSequenceEmbed( + patch_size=self.patch_size, + emb_features=self.emb_features, + scan_order=self.scan_order, + dtype=self.dtype, + precision=self.precision, + ) + self.stack = TokenStack( + features=self.emb_features, num_layers=self.num_layers, + num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, + ssm_attention_ratio=self.ssm_attention_ratio, + ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm, + dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, + ) + self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm") + + def __call__(self, x, token_idx=None, train: bool = False): + tokens, _ = self.embed(x) + if token_idx is not None: + tokens = gather_tokens(tokens, token_idx) + return self.norm(self.stack(tokens, train=train)) + + +class JepaVideoEncoder(nn.Module): + """Factorized spatial-temporal encoder over (B, T, H, W, C). + + token_idx selects a tubelet: the same patch positions in every frame, so + the factorized layout survives masking untouched. + """ + patch_size: int = 16 + emb_features: int = 384 + num_layers: int = 12 + num_heads: int = 6 + mlp_ratio: int = 4 + ssm_attention_ratio: str = "all-attn" + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + qk_norm: bool = False + attention_impl: Optional[str] = None + use_hilbert: bool = False + use_zigzag: bool = False + + @property + def scan_order(self): + assert not (self.use_hilbert and self.use_zigzag), \ + "use_hilbert and use_zigzag are mutually exclusive" + return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster' + + def setup(self): + self.embed = PatchSequenceEmbed( + patch_size=self.patch_size, + emb_features=self.emb_features, + scan_order=self.scan_order, + dtype=self.dtype, + precision=self.precision, + ) + self.stack = FactorizedTokenStack( + features=self.emb_features, num_layers=self.num_layers, + num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, + ssm_attention_ratio=self.ssm_attention_ratio, + ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm, + dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, + ) + self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm") + + def __call__(self, x, token_idx=None, train: bool = False): + B, T, H, W, C = x.shape + tokens, _ = self.embed(x.reshape(B * T, H, W, C)) + if token_idx is not None: + tokens = gather_tokens(tokens, jnp.repeat(token_idx, T, axis=0)) + tokens = tokens.reshape(B, T, tokens.shape[1], self.emb_features) + return self.norm(self.stack(tokens, train=train)) + + +class JepaPredictor(nn.Module): + """Narrow transformer from context embeddings to target embeddings. + + Context tokens are projected down, mask tokens stand in for the targets, + and both carry the sincos signal for the grid position they belong to. + """ + grid: Tuple[int, int] = (14, 14) + emb_features: int = 384 # encoder width, in and out + predictor_features: int = 192 + num_layers: int = 6 + num_heads: int = 6 + mlp_ratio: int = 4 + ssm_attention_ratio: str = "all-attn" + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + qk_norm: bool = False + attention_impl: Optional[str] = None + scan_order: str = 'raster' + factorized: bool = False # space-time blocks, for video + + def setup(self): + self.proj_in = nn.Dense(features=self.predictor_features, dtype=self.dtype, + precision=self.precision, name="proj_in") + self.mask_token = self.param( + "mask_token", nn.initializers.normal(0.02), (1, 1, self.predictor_features)) + stack_kwargs = dict( + features=self.predictor_features, num_layers=self.num_layers, + num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, + ssm_attention_ratio=self.ssm_attention_ratio, + ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm, + dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, + ) + self.stack = (FactorizedTokenStack(**stack_kwargs) if self.factorized + else TokenStack(**stack_kwargs)) + self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm") + self.proj_out = nn.Dense(features=self.emb_features, dtype=self.dtype, + precision=self.precision, name="proj_out") + + def __call__(self, context, context_idx, target_idx, train: bool = False): + """context: [B, (T,) N_ctx, F] -> predictions [B, (T,) N_tgt, F].""" + pos_embed = jnp.asarray( + scan_ordered_pos_embed(self.predictor_features, *self.grid, self.scan_order), + dtype=self.dtype or jnp.float32) + + def positions(idx): + pos = pos_embed[idx] # [B, N, P] + return pos[:, None] if self.factorized else pos + + num_targets = target_idx.shape[-1] + context = self.proj_in(context) + positions(context_idx) + targets = self.mask_token + positions(target_idx) + if self.factorized: + targets = jnp.broadcast_to(targets, (*context.shape[:-2], num_targets, + self.predictor_features)) + + tokens = jnp.concatenate([context, targets], axis=-2) + tokens = self.stack(tokens, train=train) + return self.proj_out(self.norm(tokens[..., -num_targets:, :])) diff --git a/flaxdiff/jepa/objective.py b/flaxdiff/jepa/objective.py new file mode 100644 index 0000000..2ba7a3e --- /dev/null +++ b/flaxdiff/jepa/objective.py @@ -0,0 +1,147 @@ +"""The I-JEPA / V-JEPA objective. + +Predict the representation of masked target blocks from the representation of +the visible context, in latent space. Three moving parts: + + - the context encoder sees only the context tokens and is trained; + - the target encoder sees the whole image and is not: it is the EMA of the + context encoder, so it lives in the trainer's ema_params rather than in a + parameter subtree of its own, and its branch is stop_gradient'd; + - the predictor maps context embeddings plus target positions to the target + representations. + +Targets are layer-normalized (no learned affine) before the L2 loss, which is +what keeps the scale of the prediction problem fixed as the encoder drifts. +The paper's L2 is used rather than the reference implementation's smooth-L1: +the LN already bounds the target scale, so there are no outliers for smooth-L1 +to protect against, and L2 keeps the loss directly comparable to the paper. + +The characteristic failure is silent: both encoders quietly agree on a +constant and the loss goes to zero. representation_health is reported on every +step so that collapse is visible in the training curves rather than at the end +of a probe run. +""" + +from typing import Dict, Tuple + +import jax +import jax.numpy as jnp +import optax + +from ..trainer.objectives import Objective, EMASpec +from .masking import MultiBlockMask + +CONTEXT_ENCODER = "context_encoder" +PREDICTOR = "predictor" + + +def representation_health(z) -> Dict[str, jax.Array]: + """Collapse telemetry for pooled embeddings [B, D]. + + repr_std is the per-dimension standard deviation across the batch: it goes + to zero exactly when the encoder stops distinguishing inputs. repr_cov_offdiag + is the RMS magnitude of the off-diagonal covariance, which rises when the + dimensions become redundant (dimensional collapse) even while repr_std holds. + """ + batch_size, dim = z.shape + centered = z - jnp.mean(z, axis=0, keepdims=True) + cov = (centered.T @ centered) / max(batch_size - 1, 1) + off_diagonal = cov * (1.0 - jnp.eye(dim, dtype=cov.dtype)) + return { + "repr_std": jnp.mean(jnp.std(z, axis=0)), + "repr_cov_offdiag": jnp.sqrt(jnp.sum(off_diagonal ** 2) / max(dim * (dim - 1), 1)), + } + + +def _layer_norm(x, epsilon: float = 1e-6): + """Feature-wise layer norm with no learned affine, as applied to targets.""" + mean = jnp.mean(x, axis=-1, keepdims=True) + variance = jnp.var(x, axis=-1, keepdims=True) + return (x - mean) * jax.lax.rsqrt(variance + epsilon) + + +class JepaObjective(Objective): + """Joint-embedding prediction over images (B,H,W,C) or video (B,T,H,W,C).""" + + tag = "jepa" + + def __init__( + self, + encoder, + predictor, + mask: MultiBlockMask, + sample_data_key: str, + sample_data_shape: Tuple[int, ...], + momentum: Tuple[float, float] = (0.996, 1.0), + momentum_steps: int = 100_000, + ): + self.encoder = encoder + self.predictor = predictor + self.mask = mask + self.sample_data_key = sample_data_key + self.sample_data_shape = tuple(sample_data_shape) + self.is_video = len(self.sample_data_shape) == 4 + self.ema = EMASpec( + decay=optax.linear_schedule(momentum[0], momentum[1], momentum_steps), + path=("params", CONTEXT_ENCODER), + ) + + def init_params(self, rng): + encoder_rng, predictor_rng = jax.random.split(rng) + sample = jnp.ones((1, *self.sample_data_shape)) + context_idx = jnp.arange(self.mask.num_context, dtype=jnp.int32)[None] + target_idx = jnp.arange(self.mask.block_area, dtype=jnp.int32)[None] + + encoder = self.encoder.init(encoder_rng, sample, context_idx) + context = self.encoder.apply(encoder, sample, context_idx) + predictor = self.predictor.init(predictor_rng, context, context_idx, target_idx) + return {"params": {CONTEXT_ENCODER: encoder["params"], + PREDICTOR: predictor["params"]}} + + def encode(self, encoder_params, data, token_idx=None, train=False, rngs=None): + return self.encoder.apply({"params": encoder_params}, data, token_idx, + train=train, rngs=rngs) + + def loss(self, params, ema_params, batch, rng, step): + data = (jnp.asarray(batch[self.sample_data_key], dtype=jnp.float32) - 127.5) / 127.5 + batch_size = data.shape[0] + mask_rng, dropout_rng = jax.random.split(rng) + context_idx, target_idx = self.mask.sample(mask_rng, batch_size) + num_targets = self.mask.num_targets + + # Target branch: the whole view through the EMA encoder, no gradient + full = _layer_norm(self.encode(ema_params["params"][CONTEXT_ENCODER], data)) + # [B, (T,) S, F] -> [B, M, (T,) n_tgt, F] + frame_axis = (1,) if self.is_video else () + gather_idx = target_idx.reshape(batch_size, num_targets, *frame_axis, -1, 1) + targets = jax.lax.stop_gradient( + jnp.take_along_axis(full[:, None], gather_idx, axis=-2)) + + context = self.encode( + params["params"][CONTEXT_ENCODER], data, context_idx, + train=True, rngs={"dropout": dropout_rng}) + + # Each target block is predicted from the same context: fold the block + # axis into the batch so one predictor call covers all M of them + repeated = jnp.repeat(context, num_targets, axis=0) + predictions = self.predictor.apply( + {"params": params["params"][PREDICTOR]}, + repeated, + jnp.repeat(context_idx, num_targets, axis=0), + target_idx.reshape(batch_size * num_targets, -1), + train=True, rngs={"dropout": dropout_rng}, + ).reshape(targets.shape) + + loss = jnp.mean((predictions - targets) ** 2) + pooled = jnp.mean(full, axis=tuple(range(1, full.ndim - 1))) + return loss, representation_health(pooled) + + def make_validation_step(self, **_): + """Frozen target-encoder embeddings, which the probes score.""" + def embed(val_state, batch): + data = (jnp.asarray(batch[self.sample_data_key], dtype=jnp.float32) - 127.5) / 127.5 + features = self.encode( + val_state.ema_params["params"][CONTEXT_ENCODER], data) + return jnp.mean(features, axis=tuple(range(1, features.ndim - 1))) + + return embed diff --git a/flaxdiff/jepa/probes.py b/flaxdiff/jepa/probes.py new file mode 100644 index 0000000..1adb196 --- /dev/null +++ b/flaxdiff/jepa/probes.py @@ -0,0 +1,79 @@ +"""Frozen-encoder probes, as evaluation metrics the validation loop can run. + +A JEPA run has no samples to look at, so representation quality is the only +signal that the training curve cannot give you. Both probes score the pooled +embeddings the objective's validation step produces against the batch labels, +fitting on the first half of the batch and scoring on the second - cheap +enough to run every epoch, and honest, since nothing is scored on data it was +fit on. They measure trend, not absolute transfer accuracy; a full-dataset +probe is a separate offline job. +""" + +import jax +import jax.numpy as jnp +import optax + +from ..metrics.common import EvaluationMetric + + +def _split(embeddings, labels): + fit = len(embeddings) // 2 + if fit == 0: + raise ValueError("probes need at least two samples per validation batch") + return embeddings[:fit], labels[:fit], embeddings[fit:], labels[fit:] + + +def linear_probe(embeddings, labels, num_classes: int, steps: int = 100, + learning_rate: float = 1e-2, weight_decay: float = 1e-4): + """Accuracy of a logistic regression fit on half the batch, scored on the rest.""" + fit_x, fit_y, test_x, test_y = _split(embeddings, labels) + mean, std = jnp.mean(fit_x, axis=0), jnp.std(fit_x, axis=0) + 1e-6 + fit_x, test_x = (fit_x - mean) / std, (test_x - mean) / std + + params = {"w": jnp.zeros((embeddings.shape[-1], num_classes)), + "b": jnp.zeros((num_classes,))} + optimizer = optax.adamw(learning_rate, weight_decay=weight_decay) + + def objective(p): + logits = fit_x @ p["w"] + p["b"] + return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits, fit_y)) + + def step(carry, _): + p, opt_state = carry + grads = jax.grad(objective)(p) + updates, opt_state = optimizer.update(grads, opt_state, p) + return (optax.apply_updates(p, updates), opt_state), None + + (params, _), _ = jax.lax.scan(step, (params, optimizer.init(params)), None, length=steps) + predicted = jnp.argmax(test_x @ params["w"] + params["b"], axis=-1) + return jnp.mean(predicted == test_y) + + +def knn_probe(embeddings, labels, num_classes: int, k: int = 20): + """Cosine k-NN accuracy, fit half against scored half.""" + fit_x, fit_y, test_x, test_y = _split(embeddings, labels) + fit_x = fit_x / (jnp.linalg.norm(fit_x, axis=-1, keepdims=True) + 1e-8) + test_x = test_x / (jnp.linalg.norm(test_x, axis=-1, keepdims=True) + 1e-8) + + similarity = test_x @ fit_x.T + neighbours = jnp.argsort(-similarity, axis=-1)[:, :min(k, fit_x.shape[0])] + votes = jnp.sum(jax.nn.one_hot(fit_y[neighbours], num_classes), axis=1) + return jnp.mean(jnp.argmax(votes, axis=-1) == test_y) + + +def get_linear_probe_metric(num_classes: int, label_key: str = "label", **kwargs): + return EvaluationMetric( + function=lambda embeddings, batch: linear_probe( + embeddings, jnp.asarray(batch[label_key]), num_classes, **kwargs), + name="linear_probe_accuracy", + higher_is_better=True, + ) + + +def get_knn_probe_metric(num_classes: int, label_key: str = "label", k: int = 20): + return EvaluationMetric( + function=lambda embeddings, batch: knn_probe( + embeddings, jnp.asarray(batch[label_key]), num_classes, k=k), + name="knn_probe_accuracy", + higher_is_better=True, + ) diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py index a3bd4d7..21dad0a 100644 --- a/flaxdiff/models/dit_common.py +++ b/flaxdiff/models/dit_common.py @@ -12,7 +12,7 @@ import jax import jax.numpy as jnp from flax import linen as nn -from typing import Optional +from typing import Optional, Sequence from flax.typing import Dtype, PrecisionLike from .common import FourierEmbedding, TimeProjection @@ -36,6 +36,28 @@ def scan_indices(scan_order: str, H_P: int, W_P: int): return None +def scan_ordered_pos_embed(emb_dim: int, H_P: int, W_P: int, scan_order: str): + """2D sincos position embedding permuted into the scan order, so token i of + the sequence carries the signal for the 2D position it actually came from.""" + pos_embed = build_2d_sincos_pos_embed(emb_dim, H_P, W_P) + idx = scan_indices(scan_order, H_P, W_P) + return pos_embed if idx is None else pos_embed[idx] + + +def build_block_pattern(num_layers: int, ssm_attention_ratio: str = "3:1", + block_pattern: Optional[Sequence[str]] = None): + """Per-layer mixer choice from a ratio string like '3:1', 'all-ssm', 'all-attn'.""" + if block_pattern is not None: + return list(block_pattern) + if ssm_attention_ratio == "all-ssm": + return ['ssm'] * num_layers + if ssm_attention_ratio == "all-attn": + return ['attn'] * num_layers + n_ssm, n_attn = (int(part) for part in ssm_attention_ratio.split(':')) + unit = ['ssm'] * n_ssm + ['attn'] * n_attn + return (unit * (num_layers // len(unit) + 1))[:num_layers] + + class PatchSequenceEmbed(nn.Module): """Patchify in raster/hilbert/zigzag order and add the 2D sincos signal. @@ -83,15 +105,10 @@ def __call__(self, x): else: tokens = self.patch_embed(x) - # Fixed 2D sincos position embedding (order-invariant spatial signal). - # For hilbert/zigzag, reorder the row-major embedding to match the scan - # so each token gets the sincos for its real 2D position. - pos_embed = build_2d_sincos_pos_embed(self.emb_features, H_P, W_P) - pos_embed = jnp.asarray(pos_embed, dtype=tokens.dtype) - idx = scan_indices(self.scan_order, H_P, W_P) - if idx is not None: - pos_embed = pos_embed[idx] - tokens = tokens + pos_embed[None, :, :] + # Fixed 2D sincos position embedding (order-invariant spatial signal) + pos_embed = scan_ordered_pos_embed( + self.emb_features, H_P, W_P, self.scan_order) + tokens = tokens + jnp.asarray(pos_embed, dtype=tokens.dtype)[None, :, :] return tokens, inv_idx diff --git a/flaxdiff/models/registry.py b/flaxdiff/models/registry.py index 2923b9a..c691712 100644 --- a/flaxdiff/models/registry.py +++ b/flaxdiff/models/registry.py @@ -18,6 +18,7 @@ from .ssm_dit import HybridSSMAttentionDiT from .video_dit import VideoDiT from .unet_3d import UNet3D +from ..jepa.models import JepaEncoder, JepaVideoEncoder, JepaPredictor MODEL_REGISTRY = { 'unet': Unet, @@ -29,6 +30,9 @@ 'hybrid_dit': HybridSSMAttentionDiT, 'video_dit': VideoDiT, 'unet_3d': UNet3D, + 'jepa_encoder': JepaEncoder, + 'jepa_video_encoder': JepaVideoEncoder, + 'jepa_predictor': JepaPredictor, } ARCHITECTURE_SUFFIX_FLAGS = { diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py index 749aaf0..c24211a 100644 --- a/flaxdiff/models/ssm_dit.py +++ b/flaxdiff/models/ssm_dit.py @@ -11,7 +11,7 @@ from .dit_common import ( PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, - ModulatedBlock, neutralized_rope_freqs, + ModulatedBlock, neutralized_rope_freqs, build_block_pattern, ) from .vit_common import RotaryEmbedding @@ -48,21 +48,6 @@ def scan_order(self): "use_hilbert and use_zigzag are mutually exclusive" return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster' - def _build_block_pattern(self): - """Generate block pattern from ratio string.""" - if self.block_pattern is not None: - pattern = list(self.block_pattern) - elif self.ssm_attention_ratio == "all-ssm": - pattern = ['ssm'] * self.num_layers - elif self.ssm_attention_ratio == "all-attn": - pattern = ['attn'] * self.num_layers - else: - parts = self.ssm_attention_ratio.split(':') - n_ssm, n_attn = int(parts[0]), int(parts[1]) - unit = ['ssm'] * n_ssm + ['attn'] * n_attn - pattern = (unit * (self.num_layers // len(unit) + 1))[:self.num_layers] - return pattern - def setup(self): self.embed = PatchSequenceEmbed( patch_size=self.patch_size, @@ -81,7 +66,8 @@ def setup(self): dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) - pattern = self._build_block_pattern() + pattern = build_block_pattern( + self.num_layers, self.ssm_attention_ratio, self.block_pattern) blocks = [] for i, block_type in enumerate(pattern): if block_type == 'ssm': diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index 835c435..683dfa3 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -1,8 +1,9 @@ import json +import numpy as np import flax from flax import linen as nn import jax -from typing import Callable, List, Dict, Tuple, Union, Any, Sequence, Type, Optional +from typing import Callable, List, Dict, Tuple, Union, Any, Sequence, Optional from dataclasses import field, dataclass import jax.numpy as jnp import optax @@ -56,43 +57,6 @@ def apply_ema(self, decay: float = 0.999, path: Tuple[str, ...] = ()): from flaxdiff.metrics.common import EvaluationMetric -def generate_modelname( - dataset_name: str, - noise_schedule_name: str, - architecture_name: str, - model: nn.Module, - input_config: DiffusionInputConfig, - autoencoder: AutoEncoder = None, - frames_per_sample: int = None, -) -> str: - """ - Generate a model name based on the configuration. - - Args: - config: Configuration dictionary. - - Returns: - A string representing the model name. - """ - import hashlib - import json - - # Extract key components for the name - - model_name = f"diffusion-{dataset_name}-res{input_config.sample_data_shape[-2]}" - - # model_name = f"diffusion-{dataset_name}-res{input_config.sample_data_shape[-2]}-{noise_schedule_name}-{architecture_name}" - - # if autoencoder is not None: - # model_name += f"-vae" - - # if frames_per_sample is not None: - # model_name += f"-frames_{frames_per_sample}" - - # model_name += f"-{'.'.join([cond.encoder.key for cond in input_config.conditions])}" - - return model_name - class GeneralDiffusionTrainer(SimpleTrainer): """ General trainer for diffusion models supporting both images and videos. @@ -147,32 +111,7 @@ def __init__(self, ) self.input_config = input_config self.eval_metrics = eval_metrics - - if wandb_config is not None: - # If input_config is not in wandb_config, add it - if 'input_config' not in wandb_config['config']: - wandb_config['config']['input_config'] = input_config.serialize() - # If model is not in wandb_config, add it - if 'model' not in wandb_config['config']: - wandb_config['config']['model'] = serialize_model(model) - if 'autoencoder' not in wandb_config['config'] and autoencoder is not None: - wandb_config['config']['autoencoder'] = autoencoder.name - wandb_config['config']['autoencoder_opts'] = json.dumps(autoencoder.serialize()) - - # Generate a model name based on the configuration - modelname = generate_modelname( - dataset_name=wandb_config['config']['arguments']['dataset'], - noise_schedule_name=wandb_config['config']['arguments']['noise_schedule'], - architecture_name=wandb_config['config']['arguments']['architecture'], - model=model, - input_config=input_config, - autoencoder=autoencoder, - frames_per_sample=frames_per_sample, - ) - print("Model name:", modelname) - self.modelname = modelname - wandb_config['config']['modelname'] = modelname - + self.noise_schedule = noise_schedule self.model_output_transform = model_output_transform self.unconditional_prob = unconditional_prob @@ -201,6 +140,25 @@ def __init__(self, ) self.objective = objective + if wandb_config is not None: + # If input_config is not in wandb_config, add it + if 'input_config' not in wandb_config['config']: + wandb_config['config']['input_config'] = input_config.serialize() + # If model is not in wandb_config, add it + if 'model' not in wandb_config['config']: + wandb_config['config']['model'] = serialize_model(model) + if 'autoencoder' not in wandb_config['config'] and autoencoder is not None: + wandb_config['config']['autoencoder'] = autoencoder.name + wandb_config['config']['autoencoder_opts'] = json.dumps(autoencoder.serialize()) + + # Names the wandb model artifact, so runs of different objectives on + # the same dataset and resolution do not collide + dataset_name = wandb_config['config']['arguments']['dataset'] + self.modelname = ( + f"{objective.tag}-{dataset_name}-res{input_config.sample_data_shape[-2]}") + print("Model name:", self.modelname) + wandb_config['config']['modelname'] = self.modelname + super().__init__( model=model, input_shapes=input_shapes, diff --git a/flaxdiff/trainer/objectives.py b/flaxdiff/trainer/objectives.py index b9edfb1..6c8fe4c 100644 --- a/flaxdiff/trainer/objectives.py +++ b/flaxdiff/trainer/objectives.py @@ -13,7 +13,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Tuple, Type +from typing import Any, Callable, Dict, Tuple, Type import jax import jax.numpy as jnp @@ -43,6 +43,7 @@ class EMASpec: class Objective(ABC): """The learning problem: parameters, loss, EMA policy, validation artifacts.""" + tag: str = "objective" # names the checkpoint artifact this run publishes ema: EMASpec @abstractmethod @@ -65,6 +66,8 @@ def log_validation_artifacts(self, wandb, artifacts, step: int): class DiffusionObjective(Objective): """Denoising diffusion: sample a noise level, corrupt, predict, weight.""" + tag = "diffusion" + def __init__( self, model, diff --git a/tests/test_jepa.py b/tests/test_jepa.py new file mode 100644 index 0000000..9f6b54d --- /dev/null +++ b/tests/test_jepa.py @@ -0,0 +1,367 @@ +"""I-JEPA and V-JEPA: masking, shapes, the objective, and collapse telemetry.""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.jepa import ( + JepaEncoder, JepaVideoEncoder, JepaObjective, multi_block_mask, + representation_health, linear_probe, knn_probe, +) +from flaxdiff.jepa.models import JepaPredictor +from flaxdiff.trainer import GeneralDiffusionTrainer + +RES = 32 +PATCH = 4 +GRID = (RES // PATCH, RES // PATCH) +FRAMES = 3 + + +@pytest.fixture +def mask(): + return multi_block_mask(GRID, num_targets=4, scale=(0.15, 0.2)) + + +def make_encoder(**kwargs): + return JepaEncoder(patch_size=PATCH, emb_features=32, num_layers=2, num_heads=2, + mlp_ratio=2, **kwargs) + + +def make_predictor(**kwargs): + return JepaPredictor(grid=GRID, emb_features=32, predictor_features=16, + num_layers=1, num_heads=2, mlp_ratio=2, **kwargs) + + +def make_objective(mask, **kwargs): + return JepaObjective(make_encoder(), make_predictor(), mask, + sample_data_key="image", sample_data_shape=(RES, RES, 3), **kwargs) + + +def images(seed=0, batch=4): + return jnp.asarray(np.random.RandomState(seed).uniform(0, 255, (batch, RES, RES, 3))) + + +def videos(seed=0, batch=2): + return jnp.asarray( + np.random.RandomState(seed).uniform(0, 255, (batch, FRAMES, RES, RES, 3))) + + +# --- masking --------------------------------------------------------------- + +def test_context_and_targets_are_disjoint(mask): + context_idx, target_idx = mask.sample(jax.random.PRNGKey(3), 8) + for b in range(8): + context = set(np.array(context_idx[b]).tolist()) + targets = set(np.array(target_idx[b]).reshape(-1).tolist()) + assert not (context & targets), "context encoder can see a target patch" + assert max(targets) < mask.num_patches and min(context) >= 0 + + +def test_target_coverage_sits_inside_the_configured_scale(): + scale = (0.15, 0.2) + mask = multi_block_mask(GRID, num_targets=4, scale=scale) + coverage = mask.block_area / mask.num_patches + assert scale[0] <= coverage <= scale[1] + for h, w in mask.block_shapes: + assert h * w == mask.block_area + assert 0.75 <= h / w <= 1.5 + + +def test_masks_are_reproducible_from_a_seed(mask): + a = mask.sample(jax.random.PRNGKey(11), 4) + b = mask.sample(jax.random.PRNGKey(11), 4) + c = mask.sample(jax.random.PRNGKey(12), 4) + assert all(jnp.array_equal(x, y) for x, y in zip(a, b)) + assert not all(jnp.array_equal(x, y) for x, y in zip(a, c)) + + +def test_block_shapes_and_positions_actually_vary(mask): + blocks = np.array(mask.sample(jax.random.PRNGKey(5), 64)[1]).reshape(-1, mask.block_area) + corners = {int(block.min()) for block in blocks} + widths = {len(np.unique(block % GRID[1])) for block in blocks} + assert len(corners) > 1, "every block landed in the same place" + assert len(widths) > 1, "every block came out the same shape" + + +def test_geometry_that_cannot_exist_is_rejected(): + with pytest.raises(ValueError, match="aspect ratio"): + multi_block_mask((4, 4), num_targets=2, scale=(0.18, 0.19)) + with pytest.raises(ValueError, match="no context"): + multi_block_mask(GRID, num_targets=6, scale=(0.15, 0.2)) + + +# --- forward shapes -------------------------------------------------------- + +def test_image_encoder_shapes(mask, rng): + encoder = make_encoder() + context_idx, _ = mask.sample(rng, 4) + x = images() + params = encoder.init(rng, x, context_idx) + + assert encoder.apply(params, x, context_idx).shape == (4, mask.num_context, 32) + assert encoder.apply(params, x).shape == (4, mask.num_patches, 32) + + +def test_image_predictor_shapes(mask, rng): + encoder, predictor = make_encoder(), make_predictor() + context_idx, target_idx = mask.sample(rng, 4) + x = images() + encoder_params = encoder.init(rng, x, context_idx) + context = encoder.apply(encoder_params, x, context_idx) + + block = target_idx[:, 0] + params = predictor.init(rng, context, context_idx, block) + assert predictor.apply(params, context, context_idx, block).shape == (4, mask.block_area, 32) + + +def test_video_encoder_and_predictor_shapes(mask, rng): + encoder = JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1, + num_heads=2, mlp_ratio=2) + predictor = make_predictor(factorized=True) + context_idx, target_idx = mask.sample(rng, 2) + x = videos() + + encoder_params = encoder.init(rng, x, context_idx) + context = encoder.apply(encoder_params, x, context_idx) + assert context.shape == (2, FRAMES, mask.num_context, 32) + assert encoder.apply(encoder_params, x).shape == (2, FRAMES, mask.num_patches, 32) + + block = target_idx[:, 0] + params = predictor.init(rng, context, context_idx, block) + out = predictor.apply(params, context, context_idx, block) + assert out.shape == (2, FRAMES, mask.block_area, 32) + + +def test_encoder_runs_on_the_ssm_mixer(mask, rng): + """The hybrid SSM encoder is the point of the shared block: same interface.""" + encoder = make_encoder(ssm_attention_ratio="3:1", ssm_state_dim=8) + context_idx, _ = mask.sample(rng, 4) + x = images() + params = encoder.init(rng, x, context_idx) + out = encoder.apply(params, x, context_idx) + assert out.shape == (4, mask.num_context, 32) + assert jnp.all(jnp.isfinite(out)) + + +# --- the objective --------------------------------------------------------- + +def test_fresh_loss_is_non_trivial_and_training_reduces_it(mask, rng): + objective = make_objective(mask) + params = objective.init_params(rng) + batch = {"image": images()} + + def loss_of(p): + return objective.loss(p, params, batch, jax.random.PRNGKey(7), 0)[0] + + initial = float(loss_of(params)) + assert initial > 0.1, "a fresh model already predicts the targets" + + optimizer = optax.adam(3e-3) + opt_state = optimizer.init(params) + for _ in range(20): + grads = jax.grad(loss_of)(params) + updates, opt_state = optimizer.update(grads, opt_state, params) + params = optax.apply_updates(params, updates) + + assert float(loss_of(params)) < initial * 0.9, "the objective does not train" + + +def test_video_objective_trains(mask, rng): + objective = JepaObjective( + JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1, + num_heads=2, mlp_ratio=2), + make_predictor(factorized=True), mask, + sample_data_key="video", sample_data_shape=(FRAMES, RES, RES, 3)) + params = objective.init_params(rng) + batch = {"video": videos()} + + def loss_of(p): + return objective.loss(p, params, batch, jax.random.PRNGKey(7), 0)[0] + + initial = float(loss_of(params)) + optimizer = optax.adam(3e-3) + opt_state = optimizer.init(params) + for _ in range(15): + grads = jax.grad(loss_of)(params) + updates, opt_state = optimizer.update(grads, opt_state, params) + params = optax.apply_updates(params, updates) + + assert float(loss_of(params)) < initial * 0.9 + + +def test_no_gradient_reaches_the_target_branch(mask, rng): + """stop_gradient on the target encoder is what stops the trivial solution.""" + objective = make_objective(mask) + params = objective.init_params(rng) + batch = {"image": images()} + + grads = jax.grad( + lambda ema: objective.loss(params, ema, batch, jax.random.PRNGKey(7), 0)[0] + )(params) + assert all(float(jnp.max(jnp.abs(g))) == 0.0 for g in jax.tree.leaves(grads)) + + +# --- collapse telemetry ---------------------------------------------------- + +def test_representation_std_detects_collapse(): + healthy = jax.random.normal(jax.random.PRNGKey(0), (32, 16)) + collapsed = jnp.tile(healthy[:1], (32, 1)) + + assert float(representation_health(healthy)["repr_std"]) > 0.5 + assert float(representation_health(collapsed)["repr_std"]) < 1e-6 + + +def test_collapse_telemetry_flows_through_a_degenerate_encoder(mask, rng): + """A constant encoder is the failure mode; the aux dict must show it.""" + class ConstantEncoder(JepaEncoder): + def __call__(self, x, token_idx=None, train: bool = False): + out = super().__call__(x, token_idx, train) + return jnp.zeros_like(out) + jnp.arange(out.shape[-1], dtype=out.dtype) + + batch = {"image": images()} + collapsed = JepaObjective( + ConstantEncoder(patch_size=PATCH, emb_features=32, num_layers=2, num_heads=2, + mlp_ratio=2), + make_predictor(), mask, "image", (RES, RES, 3)) + collapsed_params = collapsed.init_params(rng) + _, collapsed_aux = collapsed.loss( + collapsed_params, collapsed_params, batch, jax.random.PRNGKey(7), 0) + + healthy = make_objective(mask) + healthy_params = healthy.init_params(rng) + _, healthy_aux = healthy.loss( + healthy_params, healthy_params, batch, jax.random.PRNGKey(7), 0) + + assert float(collapsed_aux["repr_std"]) < 1e-5 + assert float(healthy_aux["repr_std"]) > 1e-3 + + +def test_offdiagonal_covariance_rises_with_redundant_dimensions(): + base = jax.random.normal(jax.random.PRNGKey(0), (64, 1)) + redundant = jnp.tile(base, (1, 8)) # every dim identical + independent = jax.random.normal(jax.random.PRNGKey(1), (64, 8)) + assert (float(representation_health(redundant)["repr_cov_offdiag"]) + > 5 * float(representation_health(independent)["repr_cov_offdiag"])) + + +# --- EMA target encoder ---------------------------------------------------- + +def test_momentum_schedule_endpoints(mask): + objective = make_objective(mask, momentum=(0.996, 1.0), momentum_steps=1000) + assert float(objective.ema.decay(0)) == pytest.approx(0.996) + assert float(objective.ema.decay(1000)) == pytest.approx(1.0) + assert objective.ema.path == ("params", "context_encoder") + + +def make_jepa_trainer(tmp_path, mask, **kwargs): + encoder = make_encoder() + return GeneralDiffusionTrainer( + model=encoder, + optimizer=optax.adam(1e-3), + input_config=DiffusionInputConfig( + sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]), + rngs=jax.random.PRNGKey(0), + objective=JepaObjective(encoder, make_predictor(), mask, "image", + (RES, RES, 3), **kwargs), + name="jepa-smoke", wandb_config=None, distributed_training=False, + checkpoint_base_path=str(tmp_path), + ) + + +def test_target_encoder_tracks_the_context_encoder(tmp_path, mask): + trainer = make_jepa_trainer(tmp_path, mask, momentum=(0.5, 0.5), momentum_steps=1) + initial = trainer.state.ema_params + + def batches(): + while True: + yield {"image": images()} + + state = trainer.fit({"train": batches, "train_len": 16, "local_batch_size": 4}, + training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0) + + context_moved = any( + not np.allclose(a, b) for a, b in zip( + jax.tree.leaves(state.ema_params["params"]["context_encoder"]), + jax.tree.leaves(initial["params"]["context_encoder"]))) + predictor_frozen = all( + np.allclose(a, b) for a, b in zip( + jax.tree.leaves(state.ema_params["params"]["predictor"]), + jax.tree.leaves(initial["params"]["predictor"]))) + assert context_moved, "the target encoder never followed the context encoder" + assert predictor_frozen, "EMA leaked outside the context encoder subtree" + + # and it followed rather than jumped: still between where it started and now + ema = jax.tree.leaves(state.ema_params["params"]["context_encoder"]) + live = jax.tree.leaves(state.params["params"]["context_encoder"]) + assert any(not np.allclose(a, b) for a, b in zip(ema, live)), "EMA is not lagging" + + +def test_validation_step_returns_pooled_embeddings(tmp_path, mask): + trainer = make_jepa_trainer(tmp_path, mask) + embed = trainer._define_validation_step() + out = embed(trainer.state, {"image": images()}) + assert out.shape == (4, 32) + + +# --- probes ---------------------------------------------------------------- + +def separable_embeddings(num_classes=4, per_class=8, dim=6, noise=0.05): + rng = np.random.RandomState(0) + centers = rng.normal(size=(num_classes, dim)) * 3 + labels = np.repeat(np.arange(num_classes), per_class) + x = centers[labels] + rng.normal(size=(len(labels), dim)) * noise + order = rng.permutation(len(labels)) + return jnp.asarray(x[order], dtype=jnp.float32), jnp.asarray(labels[order]) + + +def test_probes_separate_clustered_embeddings(): + x, y = separable_embeddings() + assert float(linear_probe(x, y, num_classes=4, steps=200)) > 0.9 + assert float(knn_probe(x, y, num_classes=4, k=3)) > 0.9 + + +def test_probes_are_at_chance_on_noise(): + x = jax.random.normal(jax.random.PRNGKey(0), (64, 6)) + y = jnp.asarray(np.random.RandomState(1).randint(0, 4, 64)) + assert float(knn_probe(x, y, num_classes=4, k=5)) < 0.6 + + +# --- registry and entrypoint ---------------------------------------------- + +@pytest.mark.parametrize("architecture", ["jepa_encoder", "jepa_video_encoder"]) +def test_registry_builds_the_jepa_models(architecture): + from flaxdiff.models.registry import build_model + model = build_model(f"{architecture}+hilbert", + {"patch_size": PATCH, "emb_features": 32, "num_layers": 1}) + assert model.emb_features == 32 and model.scan_order == 'hilbert' + + +def test_training_entrypoint_runs_end_to_end(tmp_path, monkeypatch): + """Registry -> mask -> objective -> trainer -> probes, as the CLI wires them.""" + import training_jepa + + monkeypatch.setenv("WANDB_MODE", "disabled") + classes = 4 + + def fake_dataset(*args, **kwargs): + def batches(): + rs = np.random.RandomState(0) + while True: + yield {"image": jnp.asarray(rs.uniform(0, 255, (4, RES, RES, 3))), + "label": jnp.asarray(rs.randint(0, classes, 4))} + return {"train": batches, "val": batches, "train_len": 16, "local_batch_size": 4} + + monkeypatch.setattr(training_jepa, "get_dataset_grain", fake_dataset) + args = training_jepa.parser.parse_args([ + '--image_size', str(RES), '--patch_size', str(PATCH), '--batch_size', '4', + '--emb_features', '32', '--num_layers', '1', '--num_heads', '2', '--mlp_ratio', '2', + '--predictor_features', '16', '--predictor_layers', '1', '--predictor_heads', '2', + '--epochs', '1', '--steps_per_epoch', '2', '--val_steps_per_epoch', '1', + '--probe_classes', str(classes), '--distributed_training', 'False', + '--ssm_attention_ratio', '3:1', '--ssm_state_dim', '8', + '--checkpoint_dir', str(tmp_path), + ]) + training_jepa.main(args) diff --git a/tests/test_objectives.py b/tests/test_objectives.py index fd6b39d..2fc3d1f 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -17,7 +17,7 @@ from flaxdiff.predictors import get_diffusion_preset from flaxdiff.trainer import GeneralDiffusionTrainer from flaxdiff.trainer.general_diffusion_trainer import TrainState -from flaxdiff.trainer.objectives import DiffusionObjective, EMASpec, Objective +from flaxdiff.trainer.objectives import EMASpec, Objective RES = 8 diff --git a/training_jepa.py b/training_jepa.py new file mode 100644 index 0000000..e88380c --- /dev/null +++ b/training_jepa.py @@ -0,0 +1,269 @@ +"""Train a JEPA encoder (I-JEPA over images, V-JEPA over video). + +A sibling of training.py rather than a flag on it: the two share the data +pipeline, the registry and the trainer, but nothing else. A JEPA run has no +noise schedule, no sampler, no text conditioning and no VAE, and folding an +--objective switch into training.py would mean threading None through all of +them. The Objective seam is what makes the same trainer serve both. +""" + +import argparse +import json +import os +import resource +import time +from datetime import datetime + +import jax +import optax + +from flaxdiff.data.dataloaders import get_dataset_grain, get_dataset_online +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.jepa import ( + JepaObjective, multi_block_mask, get_linear_probe_metric, get_knn_probe_metric, +) +from flaxdiff.models.registry import build_model, canonicalize_architecture +from flaxdiff.trainer import GeneralDiffusionTrainer + +os.environ['TOKENIZERS_PARALLELISM'] = "false" + +OPTIMIZER_MAP = {'adam': optax.adam, 'adamw': optax.adamw, 'lamb': optax.lamb} + + +def boolean_string(s): + return s if isinstance(s, bool) else s == 'True' + + +parser = argparse.ArgumentParser(description='Train a JEPA encoder') +parser.add_argument('--GRAIN_WORKER_COUNT', type=int, default=32) +parser.add_argument('--GRAIN_READ_THREAD_COUNT', type=int, default=140) +parser.add_argument('--GRAIN_READ_BUFFER_SIZE', type=int, default=96) +parser.add_argument('--GRAIN_WORKER_BUFFER_SIZE', type=int, default=100) + +parser.add_argument('--dataset', type=str, default='oxford_flowers102') +parser.add_argument('--dataset_path', type=str, default='/home/mrwhite0racle/gcs_mount') +parser.add_argument('--dataset_seed', type=int, default=0) +parser.add_argument('--batch_size', type=int, default=64) +parser.add_argument('--image_size', type=int, default=128) +parser.add_argument('--frames_per_sample', type=int, default=None, + help='Set for video (V-JEPA); leave unset for images (I-JEPA)') +parser.add_argument('--epochs', type=int, default=100) +parser.add_argument('--steps_per_epoch', type=int, default=None) +parser.add_argument('--val_steps_per_epoch', type=int, default=4) + +parser.add_argument('--architecture', type=str, default='jepa_encoder', + choices=['jepa_encoder', 'jepa_video_encoder']) +parser.add_argument('--patch_size', type=int, default=16) +parser.add_argument('--emb_features', type=int, default=384) +parser.add_argument('--num_layers', type=int, default=12) +parser.add_argument('--num_heads', type=int, default=6) +parser.add_argument('--mlp_ratio', type=int, default=4) +parser.add_argument('--ssm_attention_ratio', type=str, default='all-attn', + help='Mixer pattern per layer: "all-attn", "all-ssm", "3:1", ...') +parser.add_argument('--ssm_state_dim', type=int, default=64) +parser.add_argument('--use_hilbert', type=boolean_string, default=False) +parser.add_argument('--use_zigzag', type=boolean_string, default=False) +parser.add_argument('--dropout_rate', type=float, default=0.0) +parser.add_argument('--attention_impl', type=str, default=None, + choices=[None, 'xla', 'cudnn', 'tpu']) +parser.add_argument('--dtype', type=str, default=None) +parser.add_argument('--precision', type=str, default='default', + choices=['high', 'default', 'highest', 'None', None]) + +parser.add_argument('--predictor_features', type=int, default=192) +parser.add_argument('--predictor_layers', type=int, default=6) +parser.add_argument('--predictor_heads', type=int, default=6) + +parser.add_argument('--num_target_blocks', type=int, default=4) +parser.add_argument('--target_scale', type=float, nargs=2, default=[0.15, 0.2]) +parser.add_argument('--target_aspect', type=float, nargs=2, default=[0.75, 1.5]) +parser.add_argument('--momentum', type=float, nargs=2, default=[0.996, 1.0], + help='Target-encoder EMA momentum, ramped over --momentum_steps') +parser.add_argument('--momentum_steps', type=int, default=None, + help='Defaults to the full training run') + +parser.add_argument('--optimizer', type=str, default='adamw', choices=list(OPTIMIZER_MAP)) +parser.add_argument('--optimizer_opts', type=str, default='{}') +parser.add_argument('--learning_rate', type=float, default=1e-3) +parser.add_argument('--learning_rate_schedule', type=str, default=None, choices=[None, 'cosine']) +parser.add_argument('--learning_rate_peak', type=float, default=1.5e-3) +parser.add_argument('--learning_rate_end', type=float, default=1e-6) +parser.add_argument('--learning_rate_warmup_steps', type=int, default=10000) +parser.add_argument('--learning_rate_decay_epochs', type=int, default=1) +parser.add_argument('--clip_grads', type=float, default=0) + +parser.add_argument('--probe_classes', type=int, default=None, + help='Number of classes for the frozen-encoder probes') +parser.add_argument('--probe_label_key', type=str, default='label') +parser.add_argument('--knn_k', type=int, default=20) + +parser.add_argument('--distributed_training', type=boolean_string, default=True) +parser.add_argument('--experiment_name', type=str, default=None) +parser.add_argument('--load_from_checkpoint', type=str, default=None) +parser.add_argument('--resume_last_run', type=str, default=None) +parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints') +parser.add_argument('--checkpoint_fs', type=str, default='local', choices=['local', 'gcs']) +parser.add_argument('--max_checkpoints_to_keep', type=int, default=1) +parser.add_argument('--wandb_project', type=str, default='mlops-msml605-project') +parser.add_argument('--wandb_entity', type=str, default='umd-projects') + + +def main(args): + resource.setrlimit(resource.RLIMIT_CORE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY)) + resource.setrlimit(resource.RLIMIT_OFILE, (65535, 65535)) + + print("Initializing JAX") + try: + jax.distributed.initialize() + except Exception: + pass + print(f"Number of devices: {jax.device_count()}") + + checkpoint_dir = (f"gs://{args.checkpoint_dir}" if args.checkpoint_fs == 'gcs' + else args.checkpoint_dir) + + dataset_generator = (get_dataset_online if 'online' in args.dataset + else get_dataset_grain) + data = dataset_generator( + args.dataset, + batch_size=args.batch_size, image_scale=args.image_size, + worker_count=args.GRAIN_WORKER_COUNT, + read_thread_count=args.GRAIN_READ_THREAD_COUNT, + read_buffer_size=args.GRAIN_READ_BUFFER_SIZE, + worker_buffer_size=args.GRAIN_WORKER_BUFFER_SIZE, + seed=args.dataset_seed, + dataset_source=args.dataset_path, + ) + steps_per_epoch = args.steps_per_epoch or data['train_len'] // args.batch_size + total_steps = steps_per_epoch * args.epochs + + architecture, suffix_flags = canonicalize_architecture(args.architecture) + is_video = args.frames_per_sample is not None + if is_video != (architecture == 'jepa_video_encoder'): + raise ValueError( + "--frames_per_sample and --architecture=jepa_video_encoder go together") + + grid = (args.image_size // args.patch_size,) * 2 + mask = multi_block_mask( + grid, + num_targets=args.num_target_blocks, + scale=tuple(args.target_scale), + aspect=tuple(args.target_aspect), + ) + print(f"Mask geometry: {mask.block_area} tokens per target block " + f"({mask.block_shapes}), {mask.num_context} context tokens of {mask.num_patches}") + + shared = { + "emb_features": args.emb_features, + "num_heads": args.num_heads, + "mlp_ratio": args.mlp_ratio, + "ssm_attention_ratio": args.ssm_attention_ratio, + "ssm_state_dim": args.ssm_state_dim, + "dropout_rate": args.dropout_rate, + "attention_impl": args.attention_impl, + "dtype": args.dtype, + "precision": args.precision, + } + encoder_config = { + **shared, + "patch_size": args.patch_size, + "num_layers": args.num_layers, + "use_hilbert": args.use_hilbert or suffix_flags.get('use_hilbert', False), + "use_zigzag": args.use_zigzag or suffix_flags.get('use_zigzag', False), + } + predictor_config = { + **shared, + "grid": grid, + "predictor_features": args.predictor_features, + "num_layers": args.predictor_layers, + "num_heads": args.predictor_heads, + "factorized": is_video, + } + encoder = build_model(architecture, encoder_config) + predictor_config["scan_order"] = encoder.scan_order + predictor = build_model('jepa_predictor', predictor_config) + + sample_data_shape = ((args.frames_per_sample, args.image_size, args.image_size, 3) + if is_video else (args.image_size, args.image_size, 3)) + input_config = DiffusionInputConfig( + sample_data_key='video' if is_video else 'image', + sample_data_shape=sample_data_shape, + conditions=[], + ) + objective = JepaObjective( + encoder=encoder, + predictor=predictor, + mask=mask, + sample_data_key=input_config.sample_data_key, + sample_data_shape=sample_data_shape, + momentum=tuple(args.momentum), + momentum_steps=args.momentum_steps or total_steps, + ) + + eval_metrics = [] + if args.probe_classes: + eval_metrics = [ + get_linear_probe_metric(args.probe_classes, label_key=args.probe_label_key), + get_knn_probe_metric(args.probe_classes, label_key=args.probe_label_key, + k=args.knn_k), + ] + + learning_rate = args.learning_rate + if args.learning_rate_schedule == 'cosine': + learning_rate = optax.warmup_cosine_decay_schedule( + init_value=learning_rate, peak_value=args.learning_rate_peak, + warmup_steps=args.learning_rate_warmup_steps, + decay_steps=steps_per_epoch * args.learning_rate_decay_epochs, + end_value=args.learning_rate_end) + solver = OPTIMIZER_MAP[args.optimizer](learning_rate, **json.loads(args.optimizer_opts)) + if args.clip_grads > 0: + solver = optax.chain(optax.clip_by_global_norm(args.clip_grads), solver) + + experiment_name = args.experiment_name or ( + f"jepa-{args.dataset}/res-{args.image_size}/patch-{args.patch_size}/" + f"mixer-{args.ssm_attention_ratio}/emb-{args.emb_features}/" + f"lr-{args.learning_rate}/date-{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}") + print("Experiment_Name:", experiment_name) + + wandb_config = { + "project": args.wandb_project, + "entity": args.wandb_entity, + "name": experiment_name, + "config": { + "encoder": encoder_config, + "predictor": predictor_config, + "architecture": architecture, + "mask": {"grid": grid, "block_shapes": mask.block_shapes, + "block_area": mask.block_area, "num_context": mask.num_context}, + "dataset": {"name": args.dataset, "length": data['train_len']}, + "arguments": vars(args), + }, + } + if args.resume_last_run is not None: + wandb_config['id'] = args.resume_last_run + + trainer = GeneralDiffusionTrainer( + model=encoder, + optimizer=solver, + input_config=input_config, + rngs=jax.random.PRNGKey(4), + objective=objective, + name=experiment_name, + wandb_config=wandb_config, + distributed_training=args.distributed_training, + checkpoint_base_path=checkpoint_dir, + load_from_checkpoint=args.load_from_checkpoint, + max_checkpoints_to_keep=args.max_checkpoints_to_keep, + eval_metrics=eval_metrics, + best_tracker_metric="val/knn_probe_accuracy" if eval_metrics else "train/best_loss", + frames_per_sample=args.frames_per_sample, + ) + + start = time.time() + trainer.fit(data, training_steps_per_epoch=steps_per_epoch, epochs=args.epochs, + val_steps_per_epoch=args.val_steps_per_epoch) + print(f"Training finished in {time.time() - start:.0f}s") + + +if __name__ == '__main__': + main(parser.parse_args()) From 033ef3bde657f09c4626be5dee34e7e4160f9d50 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 10:03:19 -0400 Subject: [PATCH 8/9] feat(jepa): finish the objective wiring and widen test coverage Co-Authored-By: Claude Fable 5 --- flaxdiff/jepa/__init__.py | 2 +- flaxdiff/jepa/masking.py | 3 +- flaxdiff/jepa/models.py | 8 +- flaxdiff/jepa/objective.py | 11 ++- flaxdiff/trainer/__init__.py | 1 + flaxdiff/trainer/general_diffusion_trainer.py | 9 +-- flaxdiff/trainer/objectives.py | 3 +- tests/test_jepa.py | 76 ++++++++++++++++++- 8 files changed, 94 insertions(+), 19 deletions(-) diff --git a/flaxdiff/jepa/__init__.py b/flaxdiff/jepa/__init__.py index 1f597a3..1abef58 100644 --- a/flaxdiff/jepa/__init__.py +++ b/flaxdiff/jepa/__init__.py @@ -1,4 +1,4 @@ from .masking import MultiBlockMask, multi_block_mask from .models import JepaEncoder, JepaVideoEncoder, JepaPredictor, TokenStack, FactorizedTokenStack -from .objective import JepaObjective, representation_health +from .objective import JepaObjective, representation_health, normalize_targets from .probes import get_linear_probe_metric, get_knn_probe_metric, linear_probe, knn_probe diff --git a/flaxdiff/jepa/masking.py b/flaxdiff/jepa/masking.py index aaecc6a..042a048 100644 --- a/flaxdiff/jepa/masking.py +++ b/flaxdiff/jepa/masking.py @@ -19,6 +19,7 @@ them in a meaningful order. """ +import math from dataclasses import dataclass from typing import Tuple @@ -94,7 +95,7 @@ def multi_block_mask( S = grid[0] * grid[1] candidates = [ (area, _factorizations(area, grid, aspect)) - for area in range(max(1, int(jnp.ceil(scale[0] * S))), int(scale[1] * S) + 1) + for area in range(max(1, math.ceil(scale[0] * S)), math.floor(scale[1] * S) + 1) ] candidates = [(area, shapes) for area, shapes in candidates if shapes] if not candidates: diff --git a/flaxdiff/jepa/models.py b/flaxdiff/jepa/models.py index ca931b2..42e4e64 100644 --- a/flaxdiff/jepa/models.py +++ b/flaxdiff/jepa/models.py @@ -304,13 +304,13 @@ def positions(idx): pos = pos_embed[idx] # [B, N, P] return pos[:, None] if self.factorized else pos - num_targets = target_idx.shape[-1] + num_target_tokens = target_idx.shape[-1] context = self.proj_in(context) + positions(context_idx) targets = self.mask_token + positions(target_idx) if self.factorized: - targets = jnp.broadcast_to(targets, (*context.shape[:-2], num_targets, - self.predictor_features)) + targets = jnp.broadcast_to( + targets, (*context.shape[:-2], num_target_tokens, self.predictor_features)) tokens = jnp.concatenate([context, targets], axis=-2) tokens = self.stack(tokens, train=train) - return self.proj_out(self.norm(tokens[..., -num_targets:, :])) + return self.proj_out(self.norm(tokens[..., -num_target_tokens:, :])) diff --git a/flaxdiff/jepa/objective.py b/flaxdiff/jepa/objective.py index 2ba7a3e..0c37fde 100644 --- a/flaxdiff/jepa/objective.py +++ b/flaxdiff/jepa/objective.py @@ -53,8 +53,13 @@ def representation_health(z) -> Dict[str, jax.Array]: } -def _layer_norm(x, epsilon: float = 1e-6): - """Feature-wise layer norm with no learned affine, as applied to targets.""" +def normalize_targets(x, epsilon: float = 1e-6): + """Feature-wise layer norm with no learned affine. + + Applied to the target encoder's output so the prediction problem keeps a + fixed scale as the encoder drifts, and so shrinking the representation is + not a way to lower the loss. + """ mean = jnp.mean(x, axis=-1, keepdims=True) variance = jnp.var(x, axis=-1, keepdims=True) return (x - mean) * jax.lax.rsqrt(variance + epsilon) @@ -110,7 +115,7 @@ def loss(self, params, ema_params, batch, rng, step): num_targets = self.mask.num_targets # Target branch: the whole view through the EMA encoder, no gradient - full = _layer_norm(self.encode(ema_params["params"][CONTEXT_ENCODER], data)) + full = normalize_targets(self.encode(ema_params["params"][CONTEXT_ENCODER], data)) # [B, (T,) S, F] -> [B, M, (T,) n_tgt, F] frame_axis = (1,) if self.is_video else () gather_idx = target_idx.reshape(batch_size, num_targets, *frame_axis, -1, 1) diff --git a/flaxdiff/trainer/__init__.py b/flaxdiff/trainer/__init__.py index 871700a..cc7bf1c 100644 --- a/flaxdiff/trainer/__init__.py +++ b/flaxdiff/trainer/__init__.py @@ -1,2 +1,3 @@ from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics +from .objectives import Objective, EMASpec, DiffusionObjective from .general_diffusion_trainer import GeneralDiffusionTrainer, TrainState, ConditionalInputConfig diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index 683dfa3..681e087 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -299,8 +299,7 @@ def validation_loop( global_device_count = jax.device_count() local_device_count = jax.local_device_count() process_index = jax.process_index() - generate_samples = val_step_fn - + val_ds = iter(val_ds()) if val_ds else None print(f"Validation loop started for process index {process_index} with {global_device_count} devices.") # Evaluation step @@ -313,13 +312,13 @@ def validation_loop( batch = next(val_ds) if self.distributed_training and global_device_count > 1: batch = convert_to_global_tree(self.mesh, batch) - samples = generate_samples(val_state, batch) + artifacts = val_step_fn(val_state, batch) if self.eval_metrics is not None: for metric in self.eval_metrics: try: # Evaluate metrics - metric_val = metric.function(samples, batch) + metric_val = metric.function(artifacts, batch) metrics[metric.name].append(metric_val) except Exception as e: print("Error in evaluation metrics:", e) @@ -331,7 +330,7 @@ def validation_loop( print(f"Evaluation started for process index {process_index}") if self.wandb is not None and self.wandb: self.objective.log_validation_artifacts( - self.wandb, samples, current_step) + self.wandb, artifacts, current_step) # Flatten the metrics if metrics: diff --git a/flaxdiff/trainer/objectives.py b/flaxdiff/trainer/objectives.py index 6c8fe4c..a6a508e 100644 --- a/flaxdiff/trainer/objectives.py +++ b/flaxdiff/trainer/objectives.py @@ -17,6 +17,7 @@ import jax import jax.numpy as jnp +import numpy as np import optax from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple @@ -175,8 +176,6 @@ def generate_samples(val_state, batch): return generate_samples def log_validation_artifacts(self, wandb, artifacts, step: int): - import numpy as np - is_video = len(artifacts.shape) == 5 for i in range(artifacts.shape[0]): sample = np.clip((np.array(artifacts[i]) + 1) * 127.5, 0, 255).astype(np.uint8) diff --git a/tests/test_jepa.py b/tests/test_jepa.py index 9f6b54d..1d582a3 100644 --- a/tests/test_jepa.py +++ b/tests/test_jepa.py @@ -9,7 +9,7 @@ from flaxdiff.inputs import DiffusionInputConfig from flaxdiff.jepa import ( JepaEncoder, JepaVideoEncoder, JepaObjective, multi_block_mask, - representation_health, linear_probe, knn_probe, + representation_health, normalize_targets, linear_probe, knn_probe, ) from flaxdiff.jepa.models import JepaPredictor from flaxdiff.trainer import GeneralDiffusionTrainer @@ -135,6 +135,20 @@ def test_video_encoder_and_predictor_shapes(mask, rng): assert out.shape == (2, FRAMES, mask.block_area, 32) +def test_video_encoder_carries_a_temporal_signal(mask, rng): + """Tubelet masking keeps the spatial layout, so time is the only axis the + encoder can mix along: perturbing frame 0 must move frame 2's embedding.""" + encoder = JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1, + num_heads=2, mlp_ratio=2) + context_idx, _ = mask.sample(rng, 2) + x = videos() + params = jax.tree.map(lambda p: p + 0.02, encoder.init(rng, x, context_idx)) + + before = encoder.apply(params, x, context_idx) + after = encoder.apply(params, x.at[:, 0].add(1.0), context_idx) + assert not jnp.allclose(before[:, 2], after[:, 2]), "no information flow across frames" + + def test_encoder_runs_on_the_ssm_mixer(mask, rng): """The hybrid SSM encoder is the point of the shared block: same interface.""" encoder = make_encoder(ssm_attention_ratio="3:1", ssm_state_dim=8) @@ -192,6 +206,62 @@ def loss_of(p): assert float(loss_of(params)) < initial * 0.9 +def solid_colour_images(rs, n): + """One random colour per image, so a target block is fully determined by + anything else in the same image and by nothing in any other image.""" + colours = rs.uniform(20, 235, (n, 1, 1, 3)) + return jnp.asarray(np.clip(colours + rs.uniform(-15, 15, (n, RES, RES, 3)), 0, 255)) + + +def context_ablation(objective, params, ema, data, mask, rng): + """Prediction error from the true context vs. from another image's context.""" + n, blocks = data.shape[0], mask.num_targets + context_idx, target_idx = mask.sample(rng, n) + full = normalize_targets(objective.encode(ema["params"]["context_encoder"], data)) + targets = jnp.take_along_axis( + full[:, None], target_idx.reshape(n, blocks, -1, 1), axis=-2) + context = objective.encode(params["params"]["context_encoder"], data, context_idx) + + def error(ctx): + predictions = objective.predictor.apply( + {"params": params["params"]["predictor"]}, + jnp.repeat(ctx, blocks, axis=0), + jnp.repeat(context_idx, blocks, axis=0), + target_idx.reshape(n * blocks, -1), + ).reshape(targets.shape) + return float(jnp.mean((predictions - targets) ** 2)) + + return error(context), error(jnp.roll(context, 1, axis=0)) + + +def test_training_makes_the_prediction_depend_on_the_context(tmp_path, mask): + """The loss can be driven down by predicting the average target while + ignoring the context entirely - which would teach the encoder nothing. + Swapping in another image's context must cost real accuracy.""" + rs = np.random.RandomState(0) + train_x, test_x = solid_colour_images(rs, 128), solid_colour_images(rs, 32) + normalized_test = (test_x - 127.5) / 127.5 + + trainer = make_jepa_trainer(tmp_path, mask, learning_rate=3e-3, + momentum=(0.9, 0.99), momentum_steps=150) + objective = trainer.objective + before = context_ablation(objective, trainer.state.params, trainer.state.ema_params, + normalized_test, mask, jax.random.PRNGKey(9)) + assert before[1] / before[0] < 1.5, "a fresh predictor should not favour any context" + + def batches(): + while True: + yield {"image": train_x[rs.randint(0, len(train_x), 16)]} + + state = trainer.fit({"train": batches, "train_len": 128, "local_batch_size": 16}, + training_steps_per_epoch=150, epochs=1, val_steps_per_epoch=0) + + after = context_ablation(objective, state.params, state.ema_params, + normalized_test, mask, jax.random.PRNGKey(9)) + assert after[0] < before[0] / 2, "held-out prediction error did not improve" + assert after[1] / after[0] > 5.0, "the predictor still ignores its context" + + def test_no_gradient_reaches_the_target_branch(mask, rng): """stop_gradient on the target encoder is what stops the trivial solution.""" objective = make_objective(mask) @@ -256,11 +326,11 @@ def test_momentum_schedule_endpoints(mask): assert objective.ema.path == ("params", "context_encoder") -def make_jepa_trainer(tmp_path, mask, **kwargs): +def make_jepa_trainer(tmp_path, mask, learning_rate=1e-3, **kwargs): encoder = make_encoder() return GeneralDiffusionTrainer( model=encoder, - optimizer=optax.adam(1e-3), + optimizer=optax.adam(learning_rate), input_config=DiffusionInputConfig( sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]), rngs=jax.random.PRNGKey(0), From bc9d06be9a00413c760a417be4efe11d1ac7f736 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 11:33:33 -0400 Subject: [PATCH 9/9] test(jepa): prove the objective seam holds under FSDP and mixed precision JEPA over a 4x2 mesh: 21 of the objective's 59 parameter leaves land split over the fsdp axis, the EMA target encoder inherits their specs, and two real steps come back finite with the collapse telemetry intact. Compiling was never the claim - this is what shows the two halves of the merge compose. The mixed-precision branch of the step was rewritten to thread has_aux and had no coverage on either side of the merge, so it gets a test that an objective's aux survives it. Co-Authored-By: Claude Fable 5 --- tests/test_jepa.py | 47 ++++++++++++++++++++++++++++++++++++++-- tests/test_objectives.py | 16 ++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/tests/test_jepa.py b/tests/test_jepa.py index 86a35a9..9b94d26 100644 --- a/tests/test_jepa.py +++ b/tests/test_jepa.py @@ -13,6 +13,7 @@ ) from flaxdiff.jepa.models import JepaPredictor from flaxdiff.trainer import GeneralDiffusionTrainer +from flaxdiff.utils import DevicePrefetchIterator RES = 32 PATCH = 4 @@ -326,7 +327,7 @@ def test_momentum_schedule_endpoints(mask): assert objective.ema.path == ("params", "context_encoder") -def make_jepa_trainer(tmp_path, mask, learning_rate=1e-3, **kwargs): +def make_jepa_trainer(tmp_path, mask, learning_rate=1e-3, fsdp_size=1, **kwargs): encoder = make_encoder() return GeneralDiffusionTrainer( model=encoder, @@ -336,7 +337,12 @@ def make_jepa_trainer(tmp_path, mask, learning_rate=1e-3, **kwargs): rngs=jax.random.PRNGKey(0), objective=JepaObjective(encoder, make_predictor(), mask, "image", (RES, RES, 3), **kwargs), - name="jepa-smoke", wandb_config=None, distributed_training=False, + name="jepa-smoke", wandb_config=None, + distributed_training=fsdp_size > 1, + fsdp_size=fsdp_size, + # This encoder's parameters are far below the production shard + # threshold, so lower it or "FSDP on" would mean "all replicated" + fsdp_min_param_size=256, checkpoint_base_path=str(tmp_path), ) @@ -371,6 +377,43 @@ def batches(): assert any(not np.allclose(a, b) for a, b in zip(ema, live)), "EMA is not lagging" +def test_jepa_trains_under_fsdp(tmp_path, mask): + """Where the two halves of the trainer meet: an objective that owns a + multi-encoder parameter tree, run through the sharded, donating train step. + + Compiling is not the claim - the parameters have to be genuinely split + across the fsdp axis, the EMA target encoder has to follow their layout, + and the loss has to come back finite with its collapse telemetry intact. + """ + trainer = make_jepa_trainer(tmp_path, mask, fsdp_size=2) + specs = [p.sharding.spec for p in jax.tree.leaves(trainer.state.params)] + + sharded = [p for p in jax.tree.leaves(trainer.state.params) + if 'fsdp' in str(p.sharding.spec)] + assert sharded, "no JEPA parameter was sharded over the fsdp axis" + for param in sharded: + assert param.addressable_shards[0].data.size == param.size // 2 + + # The target encoder is a second copy of the same tree, so it must land on + # the mesh the same way rather than being gathered onto every device + assert specs == [p.sharding.spec for p in jax.tree.leaves(trainer.state.ema_params)] + + def batches(): + while True: + yield {"image": np.asarray(images(batch=jax.device_count()))} + + train_step = trainer._define_train_step(batch_size=jax.device_count()) + source = DevicePrefetchIterator(batches(), trainer.batch_sharding) + state, rng = trainer.state, trainer.rngstate + for _ in range(2): + state, loss, aux, rng, is_finite = train_step(state, rng, next(source)) + assert bool(is_finite) + assert float(aux["repr_std"]) > 0, "collapse telemetry was lost in the step" + + assert int(state.step) == 2 + assert specs == [p.sharding.spec for p in jax.tree.leaves(state.params)] + + def test_validation_step_returns_pooled_embeddings(tmp_path, mask): trainer = make_jepa_trainer(tmp_path, mask) embed = trainer._define_validation_step() diff --git a/tests/test_objectives.py b/tests/test_objectives.py index baeeb3c..29d6c7e 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -144,3 +144,19 @@ def test_trainer_drives_an_arbitrary_objective(tmp_path): assert int(state.step) == 3 # sum of squares under gradient descent must shrink assert tree_fingerprint(state.params) < tree_fingerprint(objective.init_params(None)) + + +def test_mixed_precision_carries_the_objective_aux(tmp_path): + """The dynamic-scale branch differentiates the same has_aux loss, so an + objective's telemetry has to survive mixed precision as well.""" + objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.9))) + trainer = make_trainer(tmp_path, objective=objective, use_dynamic_scale=True) + train_step = trainer._define_train_step(batch_size=8) + + state, loss, aux, _, is_finite = train_step( + trainer.state, trainer.rngstate, next(batch_iterator())) + # a live dynamic scale is exactly what selects that branch of the step + assert state.dynamic_scale is not None + assert bool(is_finite) + assert float(aux["probe"]) == 1.0 + assert int(state.step) == 1