Soft-JEPA-Flow (baseline_run_regularizer_cosim)

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