Commit b6bd5ed
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
- sub-packages/bionemo-moco
- examples
- src/bionemo/moco/schedules
- tests/bionemo/moco/schedules
0 commit comments