Skip to content

Commit 0083cf9

Browse files
committed
feat: Adaptive Non-uniform Timestep Sampling (arXiv:2411.09998)
Implements the paper's method for adaptively sampling timesteps during diffusion model training by learning which timesteps yield the greatest reduction in the objective. Core module (library/adaptive_timestep_sampler.py): - TimestepSamplerNetwork: MLP that outputs Beta(a,b) parameters from latents - AdaptiveTimestepManager: Orchestrates Algorithm 1 (training) and Algorithm 2 (Delta approximation via queue + F-statistic feature selection) - REINFORCE policy gradient update with entropy regularization - Full-batch delta computation at |S| selected timesteps per paper line 7 Integration (train_network.py): - Initialize sampler when --adaptive_timestep_sampling is enabled - Use Beta-distributed timesteps instead of uniform random - Algorithm 2 hooks before/after optimizer step - Full batch delta computation via cache_batch_losses_at_S() CLI arguments (train_util.py): --adaptive_timestep_sampling, --adaptive_sampler_lr, --adaptive_sampler_entropy_coeff, --adaptive_sampler_update_freq (f_S=40), --adaptive_sampler_queue_size (|Q|=20), --adaptive_sampler_num_selected (|S|=3), --adaptive_sampler_hidden_channels (128), --adaptive_sampler_hidden_depth (2) Tests: 25 tests covering sampler network, F-statistic feature selection, manager lifecycle, delta approximation, policy gradient update, cached t REINFORCE fix, and full batch extension.
1 parent 9303a66 commit 0083cf9

4 files changed

Lines changed: 1280 additions & 1 deletion

File tree

0 commit comments

Comments
 (0)