| # Dataset | |
| data_path: "embedded-language-flows/openwebtext-t5" | |
| max_length: 1024 | |
| # Encoder | |
| encoder_model_name: t5-small | |
| encoder_checkpoint: "embedded-language-flows/t5_small_encoder_jax/t5_small_encoder_jax.pkl" | |
| latent_mean: 0.0 | |
| latent_std: 0.2 | |
| # Model architecture | |
| model: ELF-B # ELF-B, ELF-M | |
| bottleneck_dim: 128 | |
| num_time_tokens: 4 | |
| num_self_cond_cfg_tokens: 4 | |
| num_model_mode_tokens: 4 | |
| # Denoiser objective | |
| denoiser_p_mean: -1.5 | |
| denoiser_p_std: 0.8 | |
| denoiser_noise_scale: 2.0 | |
| t_eps: 0.05 | |
| time_schedule: "logit_normal" | |
| # Decoder objective | |
| decoder_prob: 0.2 | |
| decoder_noise_scale: 5.0 | |
| decoder_p_mean: 0.8 | |
| decoder_p_std: 0.8 | |
| # Conditioning / CFG | |
| self_cond_prob: 0.5 | |
| # Training (optimizer + schedule) | |
| epochs: 5 | |
| global_batch_size: 512 | |
| blr: 0.001 | |
| weight_decay: 0.0 | |
| warmup_epochs: 0.5 | |
| warmup_steps: -1 # -1 falls back to warmup_epochs | |
| optimizer: muon # "adamw" or "muon" | |
| # EMA | |
| ema_decay1: 0.9999 | |
| # Sampling | |
| sampling_configs_path: "configs/sampling_configs/uncond_sampling_configs.yml" | |
| num_samples: 1000 | |
| # Logging & Checkpointing | |
| log_freq: 100 | |
| save_freq: 1 | |
| save_last_freq: 1 | |
| eval_freq: 1 | |
| # Output | |
| output_dir: "outputs/elf_b-owt" | |
| resume: null | |
| # Wandb | |
| use_wandb: true | |
| wandb_project: elf | |
| wandb_entity: null | |
| wandb_run_name: elf_b-owt | |
| # Misc | |
| seed: 42 | |