Skip to content

feat(scalarization): Implement GradNormScalarizer#733

Open
powerofaisinstudy-debug wants to merge 5 commits into
SimplexLab:mainfrom
powerofaisinstudy-debug:feat/gradnorm-scalarizer
Open

feat(scalarization): Implement GradNormScalarizer#733
powerofaisinstudy-debug wants to merge 5 commits into
SimplexLab:mainfrom
powerofaisinstudy-debug:feat/gradnorm-scalarizer

Conversation

@powerofaisinstudy-debug

@powerofaisinstudy-debug powerofaisinstudy-debug commented Jun 11, 2026

Copy link
Copy Markdown

Summary
This PR adds the GradNormScalarizer to the scalarization package, enabling dynamic loss balancing based on gradient norms as proposed by the GradNorm algorithm.

Changes

Added GradNormScalarizer class inheriting from the Scalarizer base class.

Integrated the weighting logic directly into the forward pass to ensure compatibility with the existing TorchJD architecture.

Implemented _compute_gradient_norms to facilitate gradient tracking during the forward pass.

Testing

Verified the implementation against the library's Scalarizer interface.

Confirmed that the forward method correctly handles the required input tensors.

Note

This implementation addresses the feedback provided in the previous issue regarding the integration of the scalarizer logic.

@powerofaisinstudy-debug powerofaisinstudy-debug requested a review from a team as a code owner June 11, 2026 11:46
@PierreQuinton PierreQuinton added cc: feat Conventional commit type for new features. package: scalarization labels Jun 11, 2026
@github-actions github-actions Bot changed the title feat: implement GradNormScalarizer feat(scalarization): Implement GradNormScalarizer Jun 11, 2026
@PierreQuinton

PierreQuinton commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Thanks for the PR.

Please comply to the CONTRIBUTING.md file to install and run pre-commit.

Also please review the code you submit from your agent before submitting it, it contains comments that are directly addressed to you, we are not supposed to resolve these.

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

Labels

cc: feat Conventional commit type for new features. package: scalarization

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants