| data: | |
| dataset: zebra | |
| sizes: | |
| - 5x5 | |
| repack: true | |
| pin_answer: false | |
| model: | |
| hidden_size: 384 | |
| cond_dim: 128 | |
| n_blocks: 12 | |
| n_heads: 12 | |
| dropout: 0.1 | |
| self_cond: true | |
| train: | |
| lr: 0.0003 | |
| weight_decay: 0.3 | |
| ema_decay: 0.9999 | |
| local_batch_size: 128 | |
| warmup_steps: 0 | |
| log_every: 100 | |
| save_every: 5000 | |
| val_every: 5000 | |
| save_best_metric: val/solved | |
| pipeline: | |
| seed: 0 | |
| num_workers: 2 | |
| stage_a_steps: 200000 | |
| stage_b_steps: 200000 | |
| stage_b_weight_decay: 0.0 | |
| warm_start: best | |
| loss: | |
| self_cond_prob: 0.5 | |
| self_forcing_rollout: 0.5 | |
| sfr_t_min: 0.0 | |
| sfr_fixed_depth: 16 | |
| sfr_deep_sup_n: 4 | |
| validation: | |
| num_samples: 256 | |
| subset_seed: 0 | |
| batch_size: 256 | |
| sampler: euler | |
| steps: 128 | |
| self_cond: true | |
| evaluation: | |
| num_samples: 1024 | |
| subset_seed: 0 | |
| batch_size: 64 | |
| sampler: sde | |
| noise_scale: 0.2 | |
| steps: 256 | |
| self_cond: true | |