Entropic time scheduler

[btrentini]

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

• Bug fix (non-breaking change which fixes an issue)
• 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 - Skip all continuous integration tests
• INCLUDE_NOTEBOOKS_TESTS - Execute notebook validation tests in pytest
• 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 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

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

• I have tested these changes locally
• I have updated the documentation accordingly
• [] I have added/updated tests as needed
• [] All existing tests pass successfully

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[btrentini]

@nvdreidenbach FYI

In my next iteration I intend to modify _calculate_entropy_rate to make it more general

[nvdreidenbach]

lgtm. Please merge after running the notebook tests to make sure there are no memory issues. Also remove vscode settings

[nvdreidenbach]

/ok to test

[copy-pr-bot]

/ok to test

@nvdreidenbach, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[nvdreidenbach]

/ok to test

[copy-pr-bot]

/ok to test

@nvdreidenbach, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[nvdreidenbach]

/ok to test cd43133

[btrentini]

/ok to test 2878d66

[nvdreidenbach]

/ok to test 2878d66

[btrentini]

/ok to test 5c20730

[btrentini]

/ok to test 2ff04aa

[nvdreidenbach]

/ok to test 5c20730

[copy-pr-bot]

/ok to test 5c20730

@nvdreidenbach, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[nvdreidenbach]

/ok to test 2ff04aa

[youhanl-nvidia]

lgtm. Nice works!

[youhanl-nvidia]

Looks good to me. Approved.