Commit 0083cf9
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
0 commit comments