Shrink and Perturb Weight Initialization

[melo-gonzo]

PhysicsNeMo Pull Request

Description

Adds physicsnemo.nn.shrink_and_perturb_

In-place shrink-and-perturb weight re-initialization, reference paper:

θ ← shrink·θ + perturb·ε

• θ: pretrained model weights
• shrink: retention factor in [0, 1), how much of each pretrained weight to keep
• perturb: noise level, scaled per-layer by the pretrained weight's std (default) or by a fresh init's std
• ε: random noise (Gaussian by default)

Why. Warm-starting from pretrained weights can converge to a worse loss asymptote
than training from scratch (a plasticity loss associated with grown weight norms).
Shrinking each weight toward zero and adding noise is intended to restore the scale
statistics and plasticity of a fresh init while keeping the pretrained direction,
helping recover that warm-start handicap.

Use.

from physicsnemo.nn import shrink_and_perturb_
model.load_state_dict(pretrained)
shrink_and_perturb_(model, shrink=0.5, perturb=0.1)

Model-agnostic (any torch.nn.Module, including physicsnemo.Module); in place; buffers
untouched; tied params handled.

Noise (noise=). "scaled_normal" (default): ε = std(θ)·N(0,1), scaled per-tensor
by the pretrained weight's own spread. "normal": unit Gaussian. Or a callable handed
fresh scratch (shape/dtype/device), so torch.randn_like, torch.nn.init.*_, or a
custom sampler all work, e.g. scaling noise by a fresh init's std instead of the
pretrained one.

Applications. Transfer/fine-tuning from a pretrained checkpoint, continual and
active learning (to help re-inject plasticity between rounds), and other warm-start
settings where a fresh baseline might otherwise win.

Limitations. Intended as a repair for a bad warm start rather than a way to add
signal: it tends to interpolate between fine-tune and scratch. scaled_normal adds no noise
to scalar or constant tensors (std=0), and shrink/perturb generally need tuning.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Summary

This PR adds an in-place shrink-and-perturb utility for warm-started models. The main changes are:

• New physicsnemo.nn.shrink_and_perturb_ public helper.
• Noise presets and callable noise support for parameter perturbation.
• Include filtering and tied-parameter handling.
• Unit tests for update behavior, validation, warnings, and reproducibility.
• Changelog entry for the new API.

Important Files Changed

Filename	Overview
physicsnemo/nn/module/utils/weight_init.py	Adds shrink_and_perturb_ with noise resolution, noise validation, include filtering, tied-parameter deduplication, and in-place parameter updates.
physicsnemo/nn/init.py	Re-exports shrink_and_perturb_ from the top-level physicsnemo.nn package.
physicsnemo/nn/module/utils/init.py	Re-exports shrink_and_perturb_ from the module utilities package.
test/nn/module/test_weight_init.py	Adds tests for shrink-only updates, noise modes, callable noise, validation errors, include behavior, tied parameters, buffers, scalars, warnings, and RNG preservation.
CHANGELOG.md	Documents the new shrink-and-perturb helper in the Added section.

_{Reviews (1): Last reviewed commit: "docs: fix merge conflict with changelog" | Re-trigger Greptile}

[laserkelvin]

Few minor things to address

[laserkelvin]

This formatting looks kind of weird to me - is this actually ruff'd?

[laserkelvin]

Could you not just do randn_like?

[laserkelvin]

Do you know how torch.manual_seed propagates to this?

[laserkelvin]

We don't validate the generator being on the same device as module, if it is passed

[laserkelvin]

Complete reference somewhere?