feat(dpa4): preserve cutoff smoothness in edge feature branches

[OutisLi]

What

• Add an edge_norm descriptor option to control normalization on feature branches that vanish at the cutoff, with matching PyTorch and DPModel implementations.
• When disabled, remove hidden radial-MLP RMSNorm, bypass environment-seed FiLM and cross-focus competition norms, and use a unit variance floor for post-SO(2) residual scaling.
• Serialize the new option and its nested module settings, expose it through input validation, and enable the cutoff-smooth mode in the DPA4 water example.
• Add cutoff-smoothness, serialization, configuration-propagation, and PyTorch/DPModel parity coverage.

Why

DPA4 radial features carry a C3 cutoff envelope and should decay to zero at rcut. Per-edge RMSNorm divides out that envelope until the feature variance reaches the norm epsilon floor, which can introduce a narrow, non-physical curvature and force feature just inside the cutoff. The same small-signal amplification applies to other normalization sites driven by cutoff-vanishing inputs.

The new mode preserves the cutoff envelope: FiLM modulation returns smoothly to identity, cross-focus weights approach uniform, and post-SO(2) residual messages remain linear near zero.

Impact

• edge_norm=True remains the default, so existing configurations retain their current behavior.
• Setting edge_norm=False opts into the recommended cutoff-smooth behavior.
• PyTorch and DPModel serialization and numerical behavior remain aligned in both modes.
• Normalization on persistent node states and FFN branches is unchanged.

Checks

• Added a near-cutoff radial sweep regression that verifies the curvature spike is suppressed when radial RMSNorm is disabled.
• Added structure and serialization round-trip coverage for both normalization modes.
• Added block- and descriptor-level PyTorch/DPModel parity coverage for standard and unit-floor post-SO(2) scaling.
• Tests were run separately by the author; they were not rerun in this checkout per request.

Summary by CodeRabbit

• New Features
  • Added edge_norm for DPA4 and SeZM descriptors to control edge-dependent normalization behavior.
  • Added radial_norm for radial embeddings and focus_norm for SO(2) focus competition normalization.
  • Added so2_post_norm_eps to tune SO(2 post-normalization stability.
  • Updated model save/restore so these normalization settings are preserved.
• Documentation / Examples
  • Updated the water DPA4 example input to set edge_norm explicitly.
• Tests
  • Added tests covering the new switches (including serialization) and parity/consistency across implementations.

[Copilot]

Pull request overview

This PR adds an edge_norm switch to the DPA4/SeZM descriptor to disable normalization on feature branches whose inputs vanish at the cutoff, preserving the intended cutoff envelope (and avoiding eps-floor amplification artifacts) while keeping the default behavior unchanged (edge_norm=True). The change is implemented in both PyTorch and DPModel codepaths, wired through config/serialization and input validation, and covered by new regression and parity tests; the water DPA4 example is updated to enable the cutoff-smooth mode.

Changes:

• Add edge_norm descriptor option and propagate it to radial hidden RMSNorm, env-seed FiLM norms, cross-focus competition norm, and post-SO(2) variance floor selection (PyTorch + DPModel).
• Extend serialization/config plumbing for the new switches (edge_norm, radial_norm, focus_norm, so2_post_norm_eps) and expose edge_norm in argument validation.
• Add tests for cutoff smoothness regression, serialization round-trips, and PT/DPModel parity across both normalization modes.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
source/tests/pt/model/test_dpa4_dpmodel_parity.py	Extends parity coverage for eps handling and adds edge_norm parity assertions.
source/tests/pt/model/test_descriptor_sezm.py	Adds cutoff-smoothness regression and structure/serialization checks for edge_norm/radial_norm gating.
examples/water/dpa4/input.json	Enables edge_norm: false in the DPA4 water example.
deepmd/utils/argcheck.py	Adds edge_norm to the dpa4/SeZM descriptor argument schema and documentation.
deepmd/pt/model/descriptor/sezm.py	Adds edge_norm descriptor config, propagates it to submodules, and serializes it.
deepmd/pt/model/descriptor/sezm_nn/so2.py	Adds focus_norm option to gate competition-input normalization and serializes it.
deepmd/pt/model/descriptor/sezm_nn/radial.py	Adds radial_norm option to drop hidden RMSNorm in RadialMLP and serializes it.
deepmd/pt/model/descriptor/sezm_nn/block.py	Adds so2_post_norm_eps and focus_norm plumbing into interaction blocks and serialization.
deepmd/dpmodel/descriptor/dpa4.py	DPModel-side edge_norm propagation: FiLM norm gating, radial_norm, focus_norm, post-SO(2) eps selection, and serialization/variable handling.
deepmd/dpmodel/descriptor/dpa4_nn/so2.py	DPModel SO2Convolution gains focus_norm gating + serialization updates.
deepmd/dpmodel/descriptor/dpa4_nn/radial.py	DPModel RadialMLP gains radial_norm gating + serialization updates.
deepmd/dpmodel/descriptor/dpa4_nn/block.py	DPModel interaction blocks gain focus_norm and so2_post_norm_eps plumbing + serialization updates.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: d950eab5-6b0c-4a3b-b622-1e82144939cd

📥 Commits

Reviewing files that changed from the base of the PR and between 2370fc1 and c8a2646.

📒 Files selected for processing (1)

• source/tests/pt/model/test_dpa4_dpmodel_parity.py

🚧 Files skipped from review as they are similar to previous changes (1)

• source/tests/pt/model/test_dpa4_dpmodel_parity.py

---

📝 Walkthrough

Walkthrough

DPA4 and SeZM now expose an edge_norm option controlling radial, FiLM, focus-competition, and SO(2) normalization. The setting propagates across both backends, is serialized and validated, is configured in an example, and is covered by behavior and parity tests.

Changes

Edge normalization controls

Layer / File(s)	Summary
DPmodel normalization controls deepmd/dpmodel/descriptor/dpa4.py, deepmd/dpmodel/descriptor/dpa4_nn/*	DPA4 conditionally applies FiLM, radial, focus, and SO(2) normalization and preserves the settings in serialized state and configuration.
PyTorch normalization controls deepmd/pt/model/descriptor/sezm.py, deepmd/pt/model/descriptor/sezm_nn/*	SeZM adds matching normalization flags, wires them into radial and interaction modules, adjusts SO(2) epsilon values, and serializes the configuration.
Descriptor configuration wiring deepmd/utils/argcheck.py, examples/water/dpa4/input.json	The descriptor argument schema exposes edge_norm, and the water DPA4 example disables it.
Normalization and parity validation source/tests/pt/model/test_descriptor_sezm.py, source/tests/pt/model/test_dpa4_dpmodel_parity.py	Tests cover cutoff behavior, serialization, epsilon handling, and PT/dpmodel parity for both normalization modes.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Descriptor
  participant RadialMLP
  participant InteractionBlock
  participant SO2Convolution
  Descriptor->>RadialMLP: configure radial_norm from edge_norm
  Descriptor->>InteractionBlock: configure focus_norm and post-norm epsilon
  InteractionBlock->>SO2Convolution: configure focus_norm
  Descriptor->>Descriptor: conditionally normalize FiLM logits

Loading

Possibly related PRs

• deepmodeling/deepmd-kit#5599: Touches the same DPA4/SeZM neural-network components and their constructor and serialization wiring.

Suggested reviewers: wanghan-iapcm, njzjz, iProzd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 46.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: preserving cutoff smoothness in edge feature branches.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.56%. Comparing base (98c3ab0) to head (c8a2646).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5772      +/-   ##
==========================================
- Coverage   79.70%   79.56%   -0.14%     
==========================================
  Files        1020     1020              
  Lines      116359   116386      +27     
  Branches     4305     4303       -2     
==========================================
- Hits        92742    92602     -140     
- Misses      22076    22230     +154     
- Partials     1541     1554      +13     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.