Output envelope: `learned_pair_envelope`

Role

Matrix elements of localized orbitals decay roughly exponentially with distance and span several orders of magnitude within one cutoff sphere. An envelope multiplies the network output by a known radial profile so that the network only has to learn a slowly varying, order-one correction.

Built-in members of this kind: the fitted Slater envelopes (hamiltonian_envelope_mode, hamiltonian_envelope_path) of the flat network and gaussian_envelope with scale_features in recipes.

The new layer

Instead of fitting the envelope beforehand, learned_pair_envelope makes the Slater-type profile part of the model:

\[\bm x_{ij}\longmapsto A_p\,(1+r_{ij})^{\nu_p}\,e^{-\beta_p r_{ij}}\;\bm x_{ij}, \qquad p=(a_i,a_j),\]

with three parameters per ordered species pair, stored as \(\log A_p\), \(\nu_p\) and \(\beta_p=\operatorname{softplus}(\theta_p)\ge0\). At initialization \(A=1\), \(\nu=0\) and \(\beta\approx0\), so the layer starts as the identity and learns the decay jointly with the rest of the network. In log space the envelope is linear in its parameters,

\[\log E_p(r)=\log A_p+\nu_p\log(1+r)-\beta_p r,\]

which keeps its gradients well conditioned. The paper’s ablation found a 3.8 % lower median test Hamiltonian MAE from a fitted envelope; a learned envelope removes the separate fitting step and adapts the profile to the objective.

Symmetry. One invariant factor per edge multiplies every component (checklist 3c); equivariance is unaffected whatever the parameters.

Implementation

class LearnedPairEnvelope(nn.Module):
    r"""``x_ij <- A_p (1 + r_ij)^nu_p exp(-beta_p r_ij) x_ij`` with ``p = (a_i, a_j)``.

    One invariant factor scales every component of an edge, so equivariance is
    untouched. The network then predicts a slowly varying correction to a
    physically motivated exponential decay. ``A = 1, nu = beta = 0`` at start.
    """

    def __init__(self, num_pairs):
        super().__init__()
        self.log_amplitude = nn.Parameter(torch.zeros(num_pairs))
        self.power = nn.Parameter(torch.zeros(num_pairs))
        self.raw_decay = nn.Parameter(torch.full((num_pairs,), -6.0))

    def forward(self, values, lengths, pair_index):
        r = lengths.reshape(-1)
        decay = nn.functional.softplus(self.raw_decay)[pair_index]
        log_factor = (
            self.log_amplitude[pair_index]
            + self.power[pair_index] * torch.log1p(r)
            - decay * r
        )
        return values * log_factor.exp()[:, None]

@register_transform("learned_pair_envelope")
def _learned_pair_envelope(options, inputs, mapper):
    values, lengths, pair_index = inputs
    require_index(pair_index, "edge")
    if lengths.dim != 1:
        raise ValueError("learned_pair_envelope expects one length per edge")
    require_same_support(values, lengths, pair_index)
    global_irreps(values, "edge")
    return LearnedPairEnvelope(**options), [values]

Usage

- id: envelope
  role: readout
  transform: learned_pair_envelope
  inputs: [coupled, edge_lengths, edge_type_idx]
  outputs: [coupled_scaled]
  options: {num_pairs: '${species.pair_count}'}

Variations

Initialize \((A,\nu,\beta)\) from a fit to the training targets, as the fitted envelopes do, and fine-tune them.
Use separate envelopes per irrep degree \(L\): higher-\(L\) couplings involve more anisotropic orbitals and often decay differently.

Output envelope: learned_pair_envelope