Skip to content

feat: Implement linear and cosine beta schedules#22

Open
KrishanYadav333 wants to merge 1 commit intoML4SCI:mainfrom
KrishanYadav333:feat/noise-scheduler
Open

feat: Implement linear and cosine beta schedules#22
KrishanYadav333 wants to merge 1 commit intoML4SCI:mainfrom
KrishanYadav333:feat/noise-scheduler

Conversation

@KrishanYadav333
Copy link
Copy Markdown

Pre-GSoC groundwork for the "Denoising Astronomical Observations
of Protoplanetary Disks" project (GSoC 2026).

Third PR in the groundwork series, following the config system and
model skeleton. This one is purely math — no model code touched.

Implemented the two beta schedules that control how noise is added
during the forward diffusion process:

  • Linear schedule — simple linspace from beta_start to beta_end,
    as used in the original DDPM paper (Ho et al., 2020)
  • Cosine schedule — smoother alternative from Nichol & Dhariwal 2021,
    avoids adding too much noise too quickly at the start

Also precomputes all the derived tensors needed downstream
(alphas, alphas_cumprod, sqrt variants) and implements q_sample —
the closed-form forward diffusion step:
x_t = sqrt(ᾱ_t) * x_0 + sqrt(1 - ᾱ_t) * ε

p_sample and p_sample_loop are still stubs — those need the U-Net
which is coming in the next PR.

36 unit tests covering schedule shapes, value ranges,
monotonicity, q_sample output, and invalid input validation.

Files changed:

  • src/models/noise_scheduler.py
  • tests/test_noise_scheduler.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant