Skip to content

Three center integration and numerical stabilities of trinity #366

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
floatingCatty wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Three center integration and numerical stabilities of trinity #366

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4478ae0
training stabilities enhance
floatingCatty Jul 4, 2026
7110ac1
add huberloss
floatingCatty Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 95 additions & 0 deletions docs/adr/0005-trinity-sota-enhancements.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,101 @@ P5. **Identity residuals** when irreps match; wider radial basis (16–32) for t
P6. **Parity-complete hidden irreps** (config) and, longer term, parity-restricted SO(2) maps for a
true O(3) prior.

### P2 + P3 — implemented (trainer hardening, opt-in, CPU-validated)

Shipped as `train_options` flags, all defaulting to the previous behaviour so existing configs and
checkpoints are unaffected:

- `per_group_lr: true` (**P2**) — `build_wrms_param_groups` (`dptb/utils/tools.py`) splits the
optimizer into per-block groups with lr scaled by each block's weight RMS (trust ratio at init,
clamped to [0.02, 1.0], referenced to the median group RMS). On a full-mode Trinity this puts the
small-init AtomicResNet heads (`edge_prediction_h2`, `edge_prediction_s`, |w|_rms≈0.05) at lr scale
≈0.077 while the O(1) embedding stack stays at 1.0 — a **13× measured spread**, exactly the
imbalance the audit flagged. Scales compose correctly with any scheduler (each group's base_lr is
scaled independently). Verified in `test_wrms_param_groups_*`.
- `grad_clip_norm: <float>` (**P3**) — global-norm gradient clipping each step
(`clip_grad_norm_`); off at 0.0. Tames the loss spikes typical of batch_size=1 Hamiltonian fitting.
- `ema_decay: <float>` (**P3**) — `ExponentialMovingAverage` (`dptb/nnops/ema.py`) of the weights,
with the standard `(1+t)/(10+t)` warmup on the effective decay. **Validation scores and Saver
checkpoints use the averaged weights** (`model_state_dict` = EMA), while the raw training weights
are stored under `raw_model_state_dict` for exact restart. This removes the single-iteration noise
from best-checkpoint selection (the failure mode where stage-2 `best.pth` was honestly worse than
`latest`). Verified end-to-end in `test_trainer_engages_hardening_and_checkpoints_ema`.
- `lr_scheduler.type: warmup_cos` (**P3**) — linear warmup then cosine to `eta_min`
(`get_lr_scheduler`), meant to be stepped per-iteration (`update_lr_per_iter: true`). Replaces the
RoP-to-floor collapse. Verified in `test_warmup_cosine_shape`.
- `allow_tf32: false` (**P3**, default) — TF32 matmuls are now explicitly disabled unless requested;
Hamiltonian targets need the precision.
Comment on lines +188 to +210

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Clarify that the TF32 default changes prior behavior.

“All defaulting to the previous behaviour” conflicts with Lines 209-210, which say TF32 is now explicitly disabled. Exclude allow_tf32 from that compatibility claim.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/adr/0005-trinity-sota-enhancements.md` around lines 188 - 210, Revise
the compatibility statement introducing the train_options flags so it excludes
allow_tf32 from the claim that all defaults preserve previous behavior. Keep the
allow_tf32 entry explicit that its default is false and therefore changes the
prior TF32 behavior.


Restart is backward/forward compatible: optimizer/scheduler/EMA state restore under try/except and
fall back to fresh init with a warning if the group structure changed (e.g. toggling `per_group_lr`).
Tests: `dptb/tests/test_trainer_hardening.py` (10 CPU tests). On-data A/B on Al66O36 in
`scratchpad/bench_p2p3.py`.

**Honest benchmark caveat (measured).** A short A/B on Al66O36 (2b+3b mode, 150 iters, CPU, seed 42,
lr 0.01) showed P2+P3 *slightly slower* than baseline in this regime: at iter 150, baseline
0.0942 train / 0.0915 eval vs P2+P3 0.1025 / 0.1028. This is **expected and not a regression**:
(i) `per_group_lr` *lowers* the lr on the small-init prediction heads, but in **2b+3b mode those
heads are the dominant learnable path**, so slowing them slows the transient — P2's designed benefit
is in **full mode**, where the audit measured the heads oscillating at a noise floor while the env
stack crawls, and is a *late-training* stability gain, not a first-150-iter speedup; (ii) EMA-on-eval
*lags* the raw weights during rapid initial descent (effective decay ≈0.94 by iter 150 averages the
last ~18 steps), so eval-on-EMA looks worse early and only wins once the loss flattens and raw-weight
noise dominates. The takeaway: **these are production-stability features whose payoff is at the
full-mode, long-horizon, near-convergence regime the user actually trains in** (days on GPU toward
5e-4), and they should be validated there — not on a short 2b+3b descent. They are OFF by default so
no existing run changes. Recommended production use: `grad_clip_norm` + `ema_decay` + `warmup_cos`
(low-risk, high-value at the final push) always; `per_group_lr` for full-mode runs, validated on GPU.

### P1 — implemented behind `spectral_balance` (anti spectral-collapse)

Shipped as the trinity embedding flag `spectral_balance: bool` (default **False**), affecting only the
env (message-passing) pathway. Two coupled changes attack the two measured drivers of the collapse:

- **Per-l normalization** in `SeperableLayerNorm` (`per_l=True`, `dptb/nn/norm.py`): each angular
momentum l is normalized by its OWN rotation-invariant RMS instead of the scalar-vs-pooled-l>0
split. The pooled split let high-l energy drain into low-l within the l>0 group across depth; per-l
forbids that. Uses variance-form eps `rsqrt(ms + eps²)`. **Adds no parameters** (the bucket
matrices are non-persistent buffers, rebuilt from the irreps).
- **Post-gate per-l gain** (`PerLGain`, `dptb/nn/norm.py`) after each `Gate` in `UpdateNode`/
`UpdateEdge`: one learnable scalar per l>0 (shared over mul and m). **Init 1.0 (identity)** — see the
workflow-compatibility note below. Equivariant (scalar per irrep). Adds one param per l>0 per message
block (6 params for a 3-layer model) — the only state_dict change, so old checkpoints load with
`strict=False` (gains default to 1.0) and nothing else is missing.

**Workflow compatibility — why the gain inits to identity, not 2.0.** The e3tb heads are calibrated by
`E3statistics`, which is *purely data-based*: it sets `head.scale = target-Hamiltonian per-irrep norm`
and `shift = target mean`, implicitly assuming the network features are ~unit at init. It does **not**
forward the model, so it cannot absorb a change in feature magnitude, and (critically) it runs only on
fresh-start and `init_model`, **never on restart**. An earlier gain init of 2.0 systematically doubled
the l>0 features → head output 2× target → a worse, *uncorrectable* init (re-running `E3statistics`
would not help — it re-derives the same data norms). Initializing the gain to identity keeps every
entry path (fresh / `init_model` / restart) self-consistent while leaving the gain as a learnable DOF
the optimizer can grow (most useful at the final layer, whose output has no downstream normalization).
Verified end-to-end across all three paths in `test_full_workflow_fresh_restart_initmodel`.

**Measured effect** (signal-propagation diagnostic at init, full mode, Al66O36,
`scratchpad/bench_p1_signal.py`; gain identity, so this isolates **per-l normalization**). Baseline
reproduces the collapse exactly (edge l=2 0.101→0.066 over 3 layers, node l=1 0.799→0.422, l=0 grows).
With `spectral_balance`, per-l normalization raises the **node** high-l RMS **~2× at every layer**
(node l=2 layer-0 0.34→0.79, node l=4 0.39→0.85) and modestly reduces the across-depth decay (node l=1
−47%→−41%); **edge** features are essentially unchanged. The high-l/l=0 balance on the node track
roughly doubles. Hermiticity and equivariance are preserved (`test_spectral_balance_*`).

It stays **off by default**: per-l normalization changes the node feature distribution the data-based
head calibration was tuned against, so the init loss is ~18% higher (0.42 vs 0.38 on Al66O36) — a
different, self-consistent starting point, not an inconsistency. A/B-validate on the real GPU horizon
(3-seed protocol) before making it a default.

**Honest convergence A/B** (`scratchpad/bench_p1_converge.py`, full mode, Al66O36, 200 iters, plain
Adam lr 5e-3, stats-init, everything else identical): P1 tracks baseline within noise — baseline
0.377→0.0273, P1 0.422→0.0302; P1 starts ~18% higher (per-l changes the node feature distribution vs
the data-based head calibration), briefly leads at iter 150 (0.0524 vs 0.0556), ends a hair behind.
Like P2/P3, **the mechanism is confirmed (2× higher node high-l signal) but the payoff is not a
short-horizon convergence win** — the collapse constrains the model's capacity to carry bond-anisotropy
through depth *at convergence*, which a 200-iter CPU probe on a 3-frame set does not exercise (early
loss is dominated by the larger low-l/scalar channels). Validate at the real horizon.

## Phase 3 (research directions)

- **Long-range electrostatics channel** for junctions/interfaces: onsite levels track the local
Expand Down
4 changes: 4 additions & 0 deletions dptb/nn/cutoff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ def cosine_cutoff(x: torch.Tensor, r_max: torch.Tensor, r_start_cos_ratio: float

Broadcasts over r_max.
"""
r_max = r_max.to(x.device) # tolerate an r_max tensor pinned to a different device than the data
r_max, x = torch.broadcast_tensors(r_max.unsqueeze(-1), x.unsqueeze(0))
r_decay: torch.Tensor = r_start_cos_ratio * r_max
# for x < r_decay, clamps to 1, for x > r_max, clamps to 0
Expand All @@ -31,6 +32,7 @@ def polynomial_cutoff(
Power used in envelope function
"""
assert p >= 2.0
r_max = r_max.to(x.device) # tolerate an r_max tensor pinned to a different device than the data
r_max, x = torch.broadcast_tensors(r_max.unsqueeze(-1), x.unsqueeze(0))
x = x / r_max

Expand Down Expand Up @@ -67,6 +69,8 @@ def boundary_envelope(r: torch.Tensor, r_max, onset: float = 0.95) -> torch.Tens
1 - (10 t^3 - 15 t^4 + 6 t^5) = (1-t)^3 (6 t^2 + 3 t + 1),
which is cancellation-free.
"""
if isinstance(r_max, torch.Tensor):
r_max = r_max.to(r.device) # tolerate an r_max tensor pinned to a different device than the data
t = ((r - onset * r_max) / ((1.0 - onset) * r_max)).clamp(0.0, 1.0)
omt = 1.0 - t
return omt * omt * omt * (6.0 * t * t + 3.0 * t + 1.0)
65 changes: 46 additions & 19 deletions dptb/nn/embedding/trinity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import math
from dptb.data.transforms import OrbitalMapper
from ..type_encode.one_hot import OneHotAtomEncoding
from dptb.nn.norm import SeperableLayerNorm
from dptb.nn.norm import SeperableLayerNorm, PerLGain
from dptb.data.AtomicDataDict import with_edge_vectors, with_batch
from dptb.nn.threecenter import ThreeCenterFactorized, EDGE_THREECENTER_KEY, NODE_THREECENTER_KEY
from math import ceil
Expand Down Expand Up @@ -64,6 +64,7 @@ def __init__(
freeze: Optional[Union[str, List[str]]] = None,
so2_gate: bool = False,
hermitian: bool = True,
spectral_balance: bool = False,
**kwargs,
):

Expand Down Expand Up @@ -164,6 +165,7 @@ def __init__(
r_start_cos_ratio=r_start_cos_ratio,
PolynomialCutoff_p=PolynomialCutoff_p,
cutoff_type=cutoff_type,
spectral_balance=spectral_balance,
device=device,
dtype=dtype,
)
Expand Down Expand Up @@ -210,6 +212,7 @@ def __init__(
res_update_ratios=res_update_ratios,
res_update_ratios_learnable=res_update_ratios_learnable,
so2_gate=so2_gate,
spectral_balance=spectral_balance,
dtype=dtype,
device=device,
)
Expand Down Expand Up @@ -525,6 +528,7 @@ def __init__(
r_start_cos_ratio: float = 0.8,
PolynomialCutoff_p: float = 6,
cutoff_type: str = "polynomial",
spectral_balance: bool = False,
device: Union[str, torch.device] = torch.device("cpu"),
dtype: Union[str, torch.dtype] = torch.float32,
):
Expand Down Expand Up @@ -578,20 +582,22 @@ def __init__(

self.sln_n = SeperableLayerNorm(
irreps=self.irreps_out,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)

self.sln_e = SeperableLayerNorm(
irreps=self.irreps_out,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)
Expand Down Expand Up @@ -741,6 +747,7 @@ def __init__(
res_update_ratios_learnable: bool = False,
avg_num_neighbors: Optional[float] = None,
so2_gate: bool = False,
spectral_balance: bool = False,
dtype: Union[str, torch.dtype] = torch.float32,
device: Union[str, torch.device] = torch.device("cpu"),
):
Expand All @@ -767,20 +774,22 @@ def __init__(

self.sln = SeperableLayerNorm(
irreps=self.irreps_in,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)

self.sln_e = SeperableLayerNorm(
irreps=self.edge_irreps_in,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)
Expand All @@ -807,6 +816,10 @@ def __init__(
irreps_gated # gated tensors
)

# P1: learnable per-l gain to counteract the Gate's l>0 attenuation. Init identity (1.0) so it
# does NOT disturb the data-based statistics head calibration at init (see PerLGain); off if None.
self.gain = PerLGain(self.activation.irreps_out, dtype=dtype, device=device) if spectral_balance else None

self.tp = SO2_Linear(
irreps_in=self.irreps_in+self.edge_irreps_in,
irreps_out=self.activation.irreps_in,
Expand Down Expand Up @@ -873,6 +886,8 @@ def forward(self, latents, node_features, edge_features, atom_type, node_onehot,
, dim=-1), edge_vector[active_edges], latents[active_edges], wigner=wigner) # full_out_irreps

message = self.activation(message)
if self.gain is not None:
message = self.gain(message)
message = self.lin_post(message)
scalars = message[:, :self.irreps_out[0].dim]

Expand Down Expand Up @@ -925,6 +940,7 @@ def __init__(
res_update_ratios: Optional[List[float]] = None,
res_update_ratios_learnable: bool = False,
so2_gate: bool = False,
spectral_balance: bool = False,
dtype: Union[str, torch.dtype] = torch.float32,
device: Union[str, torch.device] = torch.device("cpu"),
):
Expand Down Expand Up @@ -962,6 +978,10 @@ def __init__(
irreps_gated # gated tensors
)

# P1: learnable per-l gain to counteract the Gate's l>0 attenuation. Init identity (1.0) so it
# does NOT disturb the data-based statistics head calibration at init (see PerLGain); off if None.
self.gain = PerLGain(self.activation.irreps_out, dtype=dtype, device=device) if spectral_balance else None

self.tp = SO2_Linear(
irreps_in=self.node_irreps_in+self.irreps_in+self.node_irreps_in,
irreps_out=self.activation.irreps_in,
Expand All @@ -981,20 +1001,22 @@ def __init__(

self.sln_e = SeperableLayerNorm(
irreps=self.irreps_in,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)

self.sln_n = SeperableLayerNorm(
irreps=self.irreps_in,
eps=5e-3,
affine=True,
normalization='component',
eps=5e-3,
affine=True,
normalization='component',
std_balance_degrees=True,
per_l=spectral_balance,
dtype=self.dtype,
device=self.device
)
Expand Down Expand Up @@ -1062,6 +1084,8 @@ def forward(self, latents, node_features, node_onehot, edge_features, edge_index
scalars = new_edge_features[:, :self.tp.irreps_out[0].dim]
assert len(scalars.shape) == 2
new_edge_features = self.activation(new_edge_features)
if self.gain is not None:
new_edge_features = self.gain(new_edge_features)
new_edge_features = self.lin_post(new_edge_features)

scalars = new_edge_features[:, :self.irreps_out[0].dim]
Expand Down Expand Up @@ -1118,6 +1142,7 @@ def __init__(
res_update_ratios: Optional[List[float]] = None,
res_update_ratios_learnable: bool = False,
so2_gate: bool = False,
spectral_balance: bool = False,
dtype: Union[str, torch.dtype] = torch.float32,
device: Union[str, torch.device] = torch.device("cpu"),
):
Expand Down Expand Up @@ -1148,6 +1173,7 @@ def __init__(
res_update_ratios=res_update_ratios,
res_update_ratios_learnable=res_update_ratios_learnable,
so2_gate=so2_gate,
spectral_balance=spectral_balance,
dtype=dtype,
device=device,
)
Expand All @@ -1164,6 +1190,7 @@ def __init__(
res_update_ratios_learnable=res_update_ratios_learnable,
avg_num_neighbors=avg_num_neighbors,
so2_gate=so2_gate,
spectral_balance=spectral_balance,
dtype=dtype,
device=device,
)
Expand Down
Loading