Skip to content

feat: Implement linear and cosine beta schedules#22

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

feat: Implement linear and cosine beta schedules#22
KrishanYadav333 wants to merge 1 commit into
ML4SCI: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

KrishanYadav333 added a commit to KrishanYadav333/EXXA that referenced this pull request Jun 19, 2026
Add a "Week-4 Scaled Upgrade" subsection (16 cells) after the 64px/17M baseline:
bigger conditional DDPM (DiffusionUNet ch=128, 5-level, 71.4M params) trained on
128x128 patches with multi-GPU DataParallel. Includes forward-diffusion q(xt|x0)
visualization over T=1000, linear-vs-cosine beta-schedule plots (NoiseScheduler,
PR ML4SCI#22), an automatic CUDA-OOM fallback builder (halve batch -> drop to 64px),
and the Week-4 deliverables: training loss curve + first DDIM sample outputs +
PSNR/SSIM/MSE. The 64px/17M baseline is kept for the scaling-story record.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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