Auxiliary outputs

These outputs, which are generally idenfified by the mtt::aux:: prefix, represent additional information that the model may provide. They are not conventional trainable outputs, in the sense that they do not correspond to training targets. For example, such quantities might be the internal representation of the model, uncertainty estimates, or non-trainable quantities.

The following auxiliary outputs are currently supported by one or more architectures in metatrain:

• mtt::aux::{target}_last_layer_features: The representation of the model at the last layer, before the final linear transformation to produce target target.

• mtt::aux::{target}_uncertainty: An uncertainty estimate for target target, as computed by the LLPR module.

• mtt::aux::{target}_ensemble: An ensemble prediction for target target, as computed by the LLPR module.

mtt::aux::{target}_last_layer_features

This output contains the representation of the model at the last layer, before the final linear transformation to produce target target. If the model produces multiple targets, the corresponding representations might be different.

Metadata for last-layer features

Metadata	Names	Description
keys	"_" if invariant or non-symmetrized, ["o3_lambda", "o3_sigma"] if equivariant	if invariant or non-symmetrized, a single dimension named "_", with a single entry set to 0, or, if equivariant, two dimensions named "o3_lambda" and "o3_sigma" with the corresponding O3 equivariant indices.
samples	["system", "atom"] or ["system"]	if a per_atom output is requested, the sample names will be ["system", "atom"], otherwise they will be ["system"]. "system" ranges from 0 to the number of systems given as input to the model. "atom" ranges between 0 and the number of atoms/particles in the corresponding system.
components	No components if invariant or non-symmetrized, ["o3_mu"] if equivariant	Nothing if invariant or non-symmetrized, the O3 equivariant mu number if equivariant.
properties	"feature"	the last-layer features have a single property dimension named "feature", with entries ranging from 0 to the number of features in the last layer.

Last-layer features are supported by the following architectures:

PET

NanoPET (deprecated)

SOAP-BPNN

MACE (Experimental)

FlashMD (Experimental)

mtt::aux::{target}_uncertainty

This output contains an uncertainty estimate for target target, as computed by the LLPR module. The metadata for this output is the same as for the corresponding target, except that there is a single property named "uncertainty".

When target is energy, this output is instead named energy_uncertainty, creating the corresponding standard metatomic output.

This output is only supported by the LLPR architecture.

mtt::aux::{target}_ensemble

This output contains an ensemble prediction for target target, as computed by the LLPR module. The metadata for this output is the same as for the corresponding target, except that there is a single property named "ensemble_prediction", with entries ranging from 0 to the number of ensemble members.

When target is energy, this output is instead named energy_ensemble, creating the corresponding standard metatomic output.

This output is only supported by the LLPR architecture.