---
license: apache-2.0
library_name: flax
tags:
- jax
- flax
- diffusion-transformer
- rectified-flow
- jepa
---

# Soft-JEPA-Flow (baseline_run_regularizer_cosim)

This repository stores checkpoints automatically uploaded from training.

## Model Architecture
- Backbone: DiT-style Transformer in JAX/Flax.
- Patch size: `2`
- Hidden size: `768`
- Depth: `12`
- Attention heads: `12`
- MLP ratio: `4.0`
- Latent shape: `32x32x4`
- Number of classes: `100`
- JEPA branch: disabled in baseline mode.

## Training Objective
- Mode: `baseline`
- Objective: L_total = L_gen (rectified-flow velocity prediction objective).
- Optimizer: `adamw` (lr=`0.0001`, beta1=`0.9`, beta2=`0.99`, weight_decay=`0.06`)
- Timestep schedule: `lognormal` (mean=`-0.4`, std=`1.0`)

## Metrics
- Best metric key: `quick_fid_4096`
- Best metric value: `32.018490`
- Best checkpoint step: `150000`

## Upload Timing (Automatic)
- Upload policy: upload whenever a **new best metric** is observed.
- Upload timestamp (UTC): `20260312-150258-UTC`
- Run name: `baseline_run_regularizer_cosim`
- Path for this artifact: `baseline_run_regularizer_cosim/best/step_150000_20260312-150258-UTC`

## Evaluation Notes
- FID pipeline uses latent decode with `stabilityai/sd-vae-ft-mse`, then InceptionV3 features.
- Inception input range is `[-1, 1]` after resize to `299x299`.

## Minimal Sampling/Inference Note
- The sampler uses the velocity head in `mode="baseline"` (including checkpoints trained with JEPA).

## Source
- HF repo id: `Bangchis/soft-jepa-flow`