Skip to content

Commit b6bd5ed

Browse files
authored
Entropic time scheduler (#1024)
### Description Generates an entropic time schedule by remapping time based on the cumulative information gain provided by a predictor function. Inspired by the work from **Dejan Stancevic, Florian Handke, & Luca Ambrogioni. (2025).** Entropic Time Schedulers for Generative Diffusion Models (https://arxiv.org/abs/2504.13612) This entropy rate is used to create an optimized, data-dependent time-stepping schedule for the generative process. ### Type of changes <!-- Mark the relevant option with an [x] --> - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Refactor - [ ] Documentation update - [ ] Other (please describe): ### CI Pipeline Configuration Configure CI behavior by applying the relevant labels: - [SKIP_CI](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#skip_ci) - Skip all continuous integration tests - [INCLUDE_NOTEBOOKS_TESTS](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#include_notebooks_tests) - Execute notebook validation tests in pytest - [INCLUDE_SLOW_TESTS](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#include_slow_tests) - Execute tests labelled as slow in pytest for extensive testing > [!NOTE] > By default, the notebooks validation tests are skipped unless explicitly enabled. #### Authorizing CI Runs We use [copy-pr-bot](https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/#automation) to manage authorization of CI runs on NVIDIA's compute resources. * If a pull request is opened by a trusted user and contains only trusted changes, the pull request's code will automatically be copied to a pull-request/ prefixed branch in the source repository (e.g. pull-request/123) * If a pull request is opened by an untrusted user or contains untrusted changes, an NVIDIA org member must leave an `/ok to test` comment on the pull request to trigger CI. This will need to be done for each new commit. ### Usage <!--- How does a user interact with the changed code --> ```python from bionemo.moco.schedules.inference_time_schedules import ( EntropicInferenceSchedule, TimeDirection #Optional ) def predictor_fn(t: Tensor, x: Tensor) -> Tensor: # Ensure t is broadcastable for concatenation if t.ndim == 1: t = t.unsqueeze(-1) if t.shape[0] != x.shape[0]: t = t.expand(x.shape[0], -1) model_input = torch.cat([x, t], dim=-1) return model(model_input) def x_0_sampler_fn(n_samples: int) -> Tensor: return #Noise Sampler def x_1_sampler_fn(n_samples: int) -> Tensor: return #Target Sampler inference_sched = EntropicInferenceSchedule( predictor=predictor_fn, x_0_sampler=x_0_sampler_fn, x_1_sampler=x_1_sampler_fn, nsteps=_FLOW_STEPS, n_schedule_points=10, # More points -> more accurate schedule, but slower to generate batch_size=batch_size, direction=TimeDirection.UNIFIED, device=DEVICE, ) schedule = inference_sched.generate_schedule().to(DEVICE) dts = inference_sched.discretize().to(DEVICE) ``` ### Pre-submit Checklist <!--- Ensure all items are completed before submitting --> - [x] I have tested these changes locally - [x] I have updated the documentation accordingly - [] I have added/updated tests as needed - [] All existing tests pass successfully --------- Signed-off-by: btrentini <brunoxtrentini@gmail.com>
1 parent 77fdb9a commit b6bd5ed

3 files changed

Lines changed: 768 additions & 1 deletion

File tree

0 commit comments

Comments
 (0)