2 Prediction models

Three models use the extracted measurements: one predicts shielding tensors, one predicts shifts from static structures, and one predicts shifts from trajectory conformations. When training a shift model, the fitting code can initialise its encoders with weights learned by the shielding model.

This chapter describes the calculations and the transfer of weights. Appendix A follows the implementation from its data records through a complete prediction call.

2.1 The three predictions

Shielding. An equivariant encoder learns a representation of each atom’s environment from chemical categories, geometry and selected measurements. Two readouts convert these features into contributions to the atom’s symmetric shielding tensor. Their sum gives six predicted components, which the prediction code converts from normalised values back to shielding units.

Static shifts. The shift model starts with copies of the shielding encoder and is fitted to experimental shifts. Its readouts use features from the encoders, local descriptors and OpenFold residue embeddings to predict shifts for particular atom roles. The model adds rotamer corrections to those predictions.

Trajectory shifts. A separately fitted shift model uses the static model’s architecture without the OpenFold embedding. The evaluation code calls this model once for each of 100 conformations, then takes the equally weighted mean of each atom’s 100 predicted shifts.

The published MD fitting configuration uses one selected conformation per protein. The 100-frame average is calculated during evaluation, not during training. Each prediction uses one frame’s coordinates; the model does not receive averaged coordinates. See md/config.toml and md/evaluate.py::predict_trajectory.

Solid arrows in the overview indicate calculation flow. Dashed arrows indicate parameter copying before fitting, not a shielding-model call during prediction. “Decode to ppm” labels the conversion from normalised predictions to shifts (Section 2.2).

2.2 Preparation and units

Before passing per-atom measurements to a network, the preparation code subtracts a fitted centre from each scalar channel and divides by a fitted scale. For a vector or rank-2 channel, it divides all components by one scale and does not subtract a centre. Scaling the components together preserves their rotation rules.

A saved vocabulary assigns an integer identifier to each chemical category. The encoder uses that identifier to select a row of its learned embedding table.

The shift models keep the shielding model’s input normalisation and extend its vocabularies with categories found in the shift-training inputs. The fitting code calculates separate normalisation values for the additional rotamer descriptors and the shift targets. The fitted model includes these vocabularies and normalisation values for use during prediction.

The networks are trained to predict normalised targets. To recover shielding values, ShieldingTargetNormalizer selects the fitted centre and scales for the atom’s category. It rescales the isotropic part and restores its centre, then rescales the five rank-2 components together without adding a centre. For shifts, absolute_prediction multiplies the prediction by the scale fitted for the atom’s output role and adds the corresponding centre to obtain a shift in ppm. These conversions are separate from the neural readouts.

Source: shielding/preparation.py::_normalize_input, ShieldingTargetNormalizer; shift_preparation.py::prepare_shift_input; shift_targets.py::absolute_prediction.

2.3 Encoder architecture

ModelBatch contains the input arrays for one complete protein. Row i refers to the same atom in the position, label and measurement arrays. Bond endpoints use these row numbers; residue indices identify the atoms’ residues.

EquivariantEncoder maps each categorical label to a learned list of numbers, called an embedding. Its 17 embedding tables cover element, atom name, residue type, atom role, hybridisation and prepared chemical state. The encoder combines these embeddings with the input measurements and applies three message-passing layers over a neighbour graph.

The graph includes up to 64 nearest neighbours within 8 angstroms of each atom, plus edges for covalent bonds and selected electronic pairs. An atom can therefore receive messages from more than 64 neighbours. Edge directions are represented by spherical harmonics of ranks 0, 1 and 2, and distances by 10 radial basis functions.

At each layer, MessageConvolution calculates a message from each neighbouring atom to the receiving atom. The tensor product combines the neighbour’s features with spherical harmonics of the direction from neighbour to receiver. A small radial network calculates the mixing weights from distance and connectivity descriptors. The layer averages the incoming messages and adds that average to a learned transformation of the receiving atom’s own features. An e3nn.nn.Gate then applies nonlinear functions to scalar features and uses scalar values to scale each vector or rank-2 feature as a whole.

After each layer the hidden representation has 12 scalar channels, 12 vector channels and 12 rank-2 channels: 12 + 12*3 + 12*5 = 108 numeric components per atom. This implementation is equivariant under the rotations in SO(3). The e labels used in the code do not establish equivariance under physical reflections of all the inputs.

Source: equivariant.py::ModelBatch, EquivariantEncoder, MessageConvolution; atom_labels.py::LABEL_EMBEDDING_DIMS.

2.4 Shielding readout

ShieldingModel passes the encoder’s features to two readouts. The first, e3nn.o3.Linear, uses the same learned weights for every atom. The second, ResidueAtomCorrection, uses weights specific to the residue:atom category. The model sums their outputs: one scalar component for isotropic shielding and five components for the symmetric-traceless part. The model does not predict the three antisymmetric components of a full tensor.

The diagram shows ShieldingTargetNormalizer separately because it is not a child module of the neural model. Its conversion to shielding units is described in Section 2.2.

When the selected inputs include the MOPAC peptide rank-2 contribution, the model sets that input channel to zero before calling the encoder. It multiplies the original peptide contribution by a learned coefficient for the atom name and adds the result to the five rank-2 output components. This term leaves the scalar output unchanged. The structure-only model has no such term, and the transfer code does not copy these coefficients to the shift models.

Source: shielding/model.py::ShieldingModel, ResidueAtomCorrection; shielding/preparation.py::ShieldingTargetNormalizer.

2.5 Shift readout and chemical roles

ShiftPredictor combines two calculations for each of H, C and N: a prediction from IndependentEncoderBranches and an added dihedral-based correction from NucleusRotamerResidual.

IndependentEncoderBranches contains a separate encoder and readouts for each of H, C and N. Each encoder calculates features for the whole molecular graph; its readouts predict shifts only for its own element. The three branches have separate learned parameters but use the same EdgeContext, which stores edge geometry and descriptors rather than learned values.

Within each branch, chemical labels determine additional scale factors for features after message layers two and three. These factors apply to atoms of the branch’s element. To prepare the readout’s inputs, the branch keeps the encoder’s twelve final scalar features and takes the norm of each of its twelve vector and twelve rank-2 features. The resulting 36 values per atom do not change when the molecule and its directional measurements rotate together.

A small network supplies a further 24 values from the atom’s original invariant inputs. These are concatenated with the 36 encoder-derived values for the readout. This “local bypass” lets the readout use local information that has not passed through the message layers; each output role has its own bypass. In the static model, each H/C/N branch also reduces the 384-value OpenFold residue embedding to 32 values and appends them to the readout’s input, using the same values for all atoms in a residue.

The nine output roles are:

Each role has a network that predicts a shift from the concatenated features. A second readout uses chemical labels to weight the same network’s hidden values, and its output is added to the first prediction.

The rotamer modules calculate three further corrections from dihedral descriptors. One uses the descriptors of individual sidechain angles; another uses products of descriptors for selected pairs of angles. The third also uses the encoder’s invariant features as context. Support flags determine which corrections are added to an atom’s encoder-branch prediction.

The prediction code then converts the summed, normalised output to ppm as described in Section 2.2.

Source: shift_model.py::IndependentEncoderBranches, PreparedChemistryGate, NucleusRotamerResidual, ShiftPredictor; shift_roles.py::NUCLEUS_ROLES; shift_targets.py::absolute_prediction.

2.6 Transferring weights from shielding to shifts

shift_transfer.py::transfer_encoder_parameters_ copies the learned parameters from the shielding encoder’s three message convolutions into each H/C/N encoder. The function also copies rows from the shielding encoder’s 17 embedding tables for categories encountered during shielding training. Each shift encoder gets a separate copy. Rows for new categories keep their initial random values. The ordinary e3nn gates have no fitted parameters to copy.

The function does not copy parameters from the shielding readouts, residue:atom correction or optional peptide output term. Training on shift data fits the parameters of the shift model’s chemistry gates, local bypasses, OpenFold projections, role readouts and rotamer modules.

The published comparison also trains an otherwise matched shift model without copying encoder weights. The diagrams show the version that uses transferred weights.

Source: shift_transfer.py::transfer_encoder_parameters_; shift_fit.py::_fit_initialization_arms; shift_parent.py.