Skip to content

fix pred_std scaling - #523

Closed
Debadri-das wants to merge 1 commit into
mllam:mainfrom
Debadri-das:fix-pred-std-scaling
Closed

fix pred_std scaling#523
Debadri-das wants to merge 1 commit into
mllam:mainfrom
Debadri-das:fix-pred-std-scaling

Conversation

@Debadri-das

@Debadri-das Debadri-das commented Mar 27, 2026

Copy link
Copy Markdown

Describe your changes

Immediately after softplus, scale the predictor standard deviation (pred_std) using the one-step difference statistics (self.diff_std).

Instead of beginning around ~0.69, this first scales pred_std to the empirical step differences. This keeps measurements like NLL and CRPS from experiencing logarithmic loss explosions in the early stages of training. This offers a better initialization scale, as was discussed with @joeloskarsson. We consider this to be the main initialization step; training curve evaluation is necessary to determine whether the suggested / softplus(0.) scaling multiplier will be needed as a follow-up.
None.

Issue Link

linked to #347

Type of change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section reflecting type of change, please add where missing.

@Debadri-das
Debadri-das force-pushed the fix-pred-std-scaling branch 2 times, most recently from 99fb1ec to f67dd44 Compare March 27, 2026 06:31
@Debadri-das

Debadri-das commented Mar 27, 2026

Copy link
Copy Markdown
Author

Hi @sadamov @joeloskarsson , I noticed this issue regarding pred_std scaling and have created a preliminary PR adjusting the initialization to self.diff_std. Could you please review this?

@Debadri-das Debadri-das changed the title Fix pred std scaling [fix] pred std scaling Mar 27, 2026
@Debadri-das Debadri-das changed the title [fix] pred std scaling fix pred std scaling Mar 27, 2026
@Debadri-das Debadri-das changed the title fix pred std scaling fix pred_std scaling Mar 27, 2026
@sadamov sadamov added the bug Something isn't working label Apr 13, 2026
@sadamov sadamov self-assigned this Apr 13, 2026
@sadamov
sadamov self-requested a review April 17, 2026 19:31
@sadamov

sadamov commented Apr 23, 2026

Copy link
Copy Markdown
Collaborator

thanks @Debadri-das. Did you see that in the linked issue @joeloskarsson asked for training curves and offered a even better implementation option? It would be great to see these curves before we merge any fix.

sadamov

This comment was marked as low quality.

Sign in to view

@Debadri-das

Debadri-das commented Apr 25, 2026

Copy link
Copy Markdown
Author

Absolutely @sadamov , I'll check that linked issue and generate the training curves with @joeloskarsson 's suggested approach. Let me run those and share the results before we merge.

@sadamov
sadamov force-pushed the fix-pred-std-scaling branch from f67dd44 to db586bd Compare June 5, 2026 07:27
sadamov added a commit to Debadri-das/neural-lam that referenced this pull request Jun 5, 2026
@sadamov @Debadri-das @claude
fix: Scale pred_std by diff_std at init in BaseGraphModel (mllam#523)
db586bd 
Re-applies @Debadri-das's PR mllam#523 onto current main. The target file
`base_graph_model.py` was moved to `models/step_predictors/graph/base.py`
by mllam#208; the surrounding code is otherwise unchanged. Expanded the
NOTE-comment to explain WHY the scaling is needed (initial pred_std
otherwise sits at softplus(0) = ln(2), independent of variable scale,
which blows up NLL/CRPS in early training for any variable whose
physical step-diff std is much larger).

Co-Authored-By: Debadri Das <Debadri-das@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@sadamov @Debadri-das @claude
fix: Scale pred_std by diff_std at init in BaseGraphModel (mllam#523)
b302892 
Re-applies @Debadri-das's PR mllam#523 onto current main. The target file
`base_graph_model.py` was moved to `models/step_predictors/graph/base.py`
by mllam#208; the surrounding code is otherwise unchanged. Expanded the
NOTE-comment to explain WHY the scaling is needed (initial pred_std
otherwise sits at softplus(0) = ln(2), independent of variable scale,
which blows up NLL/CRPS in early training for any variable whose
physical step-diff std is much larger).

Co-Authored-By: Debadri Das <Debadri-das@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@sadamov
sadamov force-pushed the fix-pred-std-scaling branch from db586bd to b302892 Compare June 5, 2026 07:28
@sadamov

sadamov commented Jun 5, 2026

Copy link
Copy Markdown
Collaborator

@Debadri-das pinging this - rebased onto current main (file moved to models/step_predictors/graph/base.py via #208) and expanded the inline comment.

Still up for the training curves from April, or want me to close?

@sadamov
sadamov self-requested a review June 5, 2026 07:30
sadamov added a commit to Ayushhgit/neural-lam that referenced this pull request Jun 6, 2026
@sadamov @Ayushhgit @claude
fix: Guard per_var_std against zero feature weights (mllam#253)
ae8550d 
Re-applies @Ayushhgit's PR mllam#253 onto current main, scoped down to the
clear-cut feature-weights fix and rebased onto the post-mllam#208 layout:

- The original PR also added `+ 1e-6` to `softplus(pred_std_raw)` to
  guard against float32 underflow at very negative `pred_std_raw`.
  That overlapped with mllam#523's `* self.diff_std` scaling on the same
  line and is largely redundant once pred_std starts on the empirical
  scale, so dropped here to avoid coupling the two PRs.

- Target file `models/ar_model.py` is gone since mllam#208; fix is applied
  at `models/module.py:120` next to the existing `per_var_std`
  registration.

- Use `torch.finfo(torch.float32).eps` instead of a hard-coded `1e-8`
  for the sqrt epsilon, matching the eps choice used a few lines below
  for `state_std` / `forcing_std`.

- Slimmed tests/test_numerical_stability.py to the 3 feature-weights
  assertions (finite output, eps does not perturb non-zero weights,
  warning fires with the zero index). Dropped the softplus / wmae /
  wmse integration tests since the softplus piece is no longer in
  scope.

Refs mllam#526.

Co-Authored-By: Ayush <Ayushhgit@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This was referenced Jun 6, 2026
Open
Open
@sadamov

sadamov commented Jun 6, 2026

Copy link
Copy Markdown
Collaborator

Closing this in favour of #649, which consolidates the probabilistic metrics + ensemble plumbing track for v0.8.0. The pred_std post-softplus scaling fix will be folded into the consolidated PR when the probabilistic model arrives. Thanks for the contribution @Debadri-das.

@sadamov sadamov closed this Jun 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sadamov sadamov Awaiting requested review from sadamov

Assignees

@sadamov sadamov

Labels

bug Something isn't working

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Debadri-das @sadamov