Adding a new pretraining model

Note

This page is written for LLM agents contributing to MATCHA. It trades prose for a step-by-step recipe with concrete file paths and code snippets.

Pretraining in MATCHA is a self-supervised, CLI-only path. Unlike classic supervised models, pretraining models are not wrapped by an sklearn estimator — they are trained by the pretrain_encoder CLI command via the PretrainingModelRegistry and reused later by feeding their encoder weights into a matching classic *Model for finetuning.

Two reference implementations to copy from:

2D case — GINPretraining (src/matcha/torch/models/pretraining/gin_pretraining.py) + GraphPretrainingDataModule (src/matcha/datamodules/pretraining/graph_pretraining_datamodule.py).
3D case — E3GNNPretraining, GPS3DPretraining, GT3DPretraining (src/matcha/torch/models/pretraining/{e3gnn,gps3d,gt3d}_pretraining.py) + Graph3DPretrainingDataModule (src/matcha/datamodules/pretraining/graph_3d_pretraining_datamodule.py). Coordinates ride on graph.pos — no separate collate key.

A new pretraining model touches four layers, in this order:

torch/encoders/   →  datamodules/pretraining/   →  utils/schemas/   →  torch/models/pretraining/   →  tests/
(canonical enc)      (batch producer)              (pydantic)          (Lightning module)             (parity + fit)

The encoder itself is shared with the classic path — do not create a pretraining-specific encoder twin. Issue #24 deleted every one of those; the invariant “one canonical encoder per architecture” is load-bearing.

1. Encoder — reuse the canonical class

Pretraining models consume the same encoder as the classic model of the same family. Add a new encoder only when a genuinely new architecture is being introduced (in which case follow adding-a-model.md first, then come back here).

The pretraining base class calls encoder.forward_nodes_per_layer(graph) to obtain one node-feature tensor per layer, then feeds the list through the shared jumping-knowledge merge and the two prediction heads. Any encoder exposing that method Just Works — see BaseGraphPretrainingModel._get_per_layer_embeddings in src/matcha/torch/models/pretraining/base_graph_pretraining.py.

For 3D encoders the input contract is the PyG convention: coordinates are read from graph.pos. Missing pos must raise a ValueError with a clear message pointing at the datamodule — that error is the tripwire for misconfigured pipelines.

2. Datamodule — `src/matcha/datamodules/pretraining/<name>_datamodule.py`

Pretraining datamodules produce batches with a graph field, per-atom targets y_node, and per-molecule targets y_graph. Add positional/structural fields (like 3D coords on graph.pos) on top of the parent class rather than parallel to it.

Inherit from GraphPretrainingDataModule (never from a classic Graph3DDataModule — that would carry the ETKDG conformer path, which is out of scope for pretraining) and register on DataModuleRegistry:

# see src/matcha/datamodules/pretraining/graph_3d_pretraining_datamodule.py
from matcha.datamodules.base_datamodule import DataModuleRegistry
from matcha.datamodules.pretraining.graph_pretraining_datamodule import (
    GraphPretrainingDataModule,
)

@DataModuleRegistry.register("graph3d_pretraining")
class Graph3DPretrainingDataModule(GraphPretrainingDataModule):
    def featurize(self, mol_list, y_graph, y_node, coords, is_training=True, n_jobs=None):
        # validate + reorder coords, attach to Data.pos, parent collate handles the rest
        ...
    def export_to_classic(self) -> Graph3DDataModule:
        # mirror pretraining PE settings into the classic 3D datamodule
        ...

Key rules:

Wire coords through Data.pos, not a parallel collate key. Batch.from_data_list auto-concatenates pos, so the inherited collate function needs no changes. Any 3D encoder reads graph.pos inside its per-layer hook.
Validate against the canonical-SMILES atom count. The parent’s _validate_node_labels uses this convention for y_node; new node-aligned fields (coords, per-atom features) must mirror it so shape errors fail fast with a clear message.
Reorder user-supplied per-atom data to canonical order. GraphDataModule._calculate_graph reparses each molecule from its canonical SMILES; user rows come back misaligned otherwise. Use mol.GetSubstructMatch(canonical_mol) to remap.
Zero-pad virtual nodes for coordinate-like fields, never NaN — NaN would poison distance-based features (E3GNN’s Fourier distance) on real neighbours of a virtual node. y_node’s NaN padding is safe only because MultitaskLoss masks it.
Override export_to_classic(). It returns a classic datamodule that mirrors the pretraining PE settings so downstream finetuning inherits the same featurization.
Set a stable state_dict "ID" string so serialization survives round-trips.

3. Schema — `src/matcha/utils/schemas/`

Three touches:

Datamodule schema in datamodules.py — subclass the parent pretraining schema, override datamodule_type, and add to the union at the bottom:

# see src/matcha/utils/schemas/datamodules.py
class Graph3DPretrainingDataModuleInputModel(GraphPretrainingDataModuleInputModel):
    datamodule_type: Literal["graph3d_pretraining"] = "graph3d_pretraining"

DataModuleModel = (
    ...
    | Graph3DPretrainingDataModuleInputModel
)

Model schema in torch_api.py — compose PretrainingMatchaModel + GraphMixin + GraphPretrainingMixin + the architecture’s *Mixin. PretrainingMatchaModel and GraphPretrainingMixin live in generic_models.py and cover loss/optimizer/scheduler and the joint head fields (num_node_targets, num_graph_targets, node_head_dims, graph_head_dims, node_loss_weight, …) that every graph pretraining model shares.
```
# see src/matcha/utils/schemas/torch_api.py
class E3GNNPretrainingInputModel(
    PretrainingMatchaModel, GraphMixin, GraphPretrainingMixin, E3GNNMixin
):
    torch_type: Literal["e3gnn_pretraining"] = "e3gnn_pretraining"
    pred_hidden_dims: list[int] | None = None
    pred_task_head_dims: list[int] | None = None

TorchModel = (
    ...
    | E3GNNPretrainingInputModel
)
```
pred_hidden_dims / pred_task_head_dims come from GraphMixin but are unused in pretraining — default them to None so callers don’t need to pass them.
Export both new schemas from utils/schemas/__init__.py (import + __all__).

4. Lightning module — `src/matcha/torch/models/pretraining/<name>_pretraining.py`

Subclass BaseGraphPretrainingModel, register on PretrainingModelRegistry, and implement _build_encoder(). The base class handles the per-layer hook, JK merge, both heads, both losses, training_step, validation_step, and per-task logging.

# see src/matcha/torch/models/pretraining/e3gnn_pretraining.py
from matcha.torch.encoders.e3gnn import E3GNN
from matcha.torch.models.pretraining.base_graph_pretraining import (
    BaseGraphPretrainingModel,
)
from matcha.torch.models.pretraining.base_pretraining_model import (
    PretrainingModelRegistry,
)

@PretrainingModelRegistry.register()
class E3GNNPretraining(BaseGraphPretrainingModel):
    def __init__(self, num_node_targets=1, num_graph_targets=1,
                 enc_num_layers=3, enc_atom_hidden_dim=128, ...,
                 node_head_dims=None, graph_head_dims=None, ...):
        super().__init__(num_node_targets=..., num_graph_targets=..., ...)
        self.save_hyperparameters()
        self._build_encoder()
        head_input_dim = enc_atom_hidden_dim * enc_num_layers if enc_jk == "concat" else enc_atom_hidden_dim
        self.node_head  = self._build_prediction_head(head_input_dim, node_head_dims, num_node_targets, ...)
        self.graph_head = self._build_prediction_head(head_input_dim, graph_head_dims, num_graph_targets, ...)
        self._parse_train_config()

    def _build_encoder(self):
        self.encoder = E3GNN(
            num_layers=self.hparams["enc_num_layers"],
            atom_input_dim=self.hparams["enc_atom_input_dim"] + self.hparams["enc_laplacian_k"] + ...,
            ...
        )

Rules:

Never override forward / training_step / _get_per_layer_embeddings. For 3D encoders, coordinates flow through graph.pos — the base class reads them via self.encoder.forward_nodes_per_layer(batch["graph"]). Adding an override reintroduces the drift issues #24 and #26 were designed to eliminate.
enc_*-prefix every encoder-facing hyperparameter. The classic *Model uses the same prefixes, which is what makes encoder-weight transfer at finetuning time a strict load_state_dict.
Add the class to torch/models/pretraining/__init__.py (import + __all__).

5. Tests — parity, fit, and encoder transfer

The pretraining test suite lives under tests/pretraining/. Add a new file test_<name>_pretraining.py covering:

Registry lookup: "<name>pretraining" in PretrainingModelRegistry.
Canonical encoder wiring: isinstance(model.encoder, <CanonicalEncoder>).
Forward shapes: model(batch)["node"].shape == (N_atoms, num_node_targets) and ["graph"].shape == (batch_size, num_graph_targets).
Per-layer hook length: len(model._get_per_layer_embeddings(batch)[0]) == num_layers.
One-step fit: model.training_step(batch, 0) returns a finite scalar that requires_grad.
Optimizer step actually moves weights: compare parameters_to_vector(model.encoder.parameters()) before and after.
save_hyperparameters round-trip through state_dict/load_state_dict.
Encoder weight transfer to the classic *Model — the headline acceptance criterion. Strict load_state_dict from the pretraining encoder into the classic encoder, then parameters_to_vector equality:
```
missing, unexpected = classic.encoder.load_state_dict(
    pretrain.encoder.state_dict(), strict=True
)
assert missing == []
assert unexpected == []
```

Then extend the parametrized parity suites so the new model is exercised alongside the existing ones:

tests/pretraining/test_encoder_parity.py — add a pytest.param(<ClassicModel>, <PretrainingModel>, dict(enc_*=...), id="<name>") entry. The three checks (parameter-key equality, module-tree string equality, allclose on encoder output after a weight-sync load_state_dict) then run automatically.
tests/pretraining/test_graph_pretraining_encoders.py — add pytest.param(<PretrainingModel>, <CanonicalEncoder>, dict(enc_*=...), id="<name>") for the is_canonical / per-layer-length / forward-shape checks.

For 3D architectures also add a classic ↔ pretraining datamodule parity test: build a small mol list, featurize once through Graph3DDataModule (ETKDG coords), reuse the resulting graph.pos as user-supplied coords in the pretraining datamodule, and assert that E3GNN.forward_nodes_per_layer produces bit-identical per-layer outputs on both batches. See tests/pretraining/test_e3gnn_pretraining.py::test_classic_and_pretraining_datamodules_produce_identical_e3gnn_features for the shape.

6. Wire the pretraining path

torch/models/pretraining/__init__.py — import + __all__.
datamodules/pretraining/__init__.py — import + __all__.
datamodules/__init__.py — add a lazy-import entry (_LAZY_IMPORTS); a plain top-level import will reintroduce the circular chain.
utils/schemas/__init__.py — export both schemas.

No sklearn wiring is required — the pretraining path never goes through an sklearn estimator.

CLI wiring (2D and 3D). Both pretrain_encoder and pretrain_multitask accept precomputed per-molecule 3D conformers via a shared flat + offsets npz layout — the datamodules deliberately never run ETKDG, so coordinates must be supplied upfront.

pretrain_encoder — EncoderPretrainDataset.task_type is Literal["mlm", "graph", "graph3d"]. The graph3d branch (cli/pretrain_encoder.py::_run_graph3d_pretraining) dispatches to Graph3DPretrainingDataModule (registered as "graph3d_pretraining") and requires train_coords / val_coords alongside the four y_graph / y_node npz files. The schema rejects a graph3d config missing any of those six paths, and rejects coords supplied under mlm / graph. See cli/example_configs/pretrain_encoder_graph3d.yaml for a runnable example.
pretrain_multitask — coords are discovered by convention from the dataset directory: if {dataset_dir}/train_coords.npz and {dataset_dir}/val_coords.npz exist, they are loaded and passed to OnTheFlyDataModule.set_data(..., train_coords=..., val_coords=...). There is no schema field to opt in — dropping the two files in the prepared dataset directory is the switch. When the resolved multitask architecture uses a 3D-capable base datamodule (one whose generate_features accepts a coords kwarg), the wrapper forwards them via base.generate_features(..., coords=..., n_jobs=1). When the base is 2D-only, coords are silently dropped and a single logging.warning is emitted (stdlib logger, not warnings.warn — pytest’s filterwarnings=error would trip on the latter). Absence of the coord files is byte-identical to the pre-issue-29 behavior.
Coord npz layout. Both CLIs read the same on-disk shape as train_y_node / val_y_node:
```
np.savez_compressed(
    "train_coords.npz",
    flat=flat,        # (total_atoms, 3) float
    offsets=offsets,  # (N + 1,) int, monotonic non-decreasing, offsets[-1] == flat.shape[0]
)
```
cli.utils._load_coords_npz performs startup assertions (flat.ndim == 2, flat.shape[1] == 3, offsets monotonic, offsets[-1] == flat.shape[0]) so malformed inputs fail at CLI parse time rather than at featurize time.

Rows in train_coords / val_coords must align 1-to-1 with the SMILES parquet in the same split. Downstream reordering to the canonical-SMILES atom order is handled inside Graph3DPretrainingDataModule.featurize — user-supplied rows should follow the original SMILES atom order.

Checklist

Reuse the canonical encoder from torch/encoders/ (no pretraining twin).
Datamodule inherits from GraphPretrainingDataModule, registered with DataModuleRegistry, validates + reorders per-atom user data to canonical order, ships an export_to_classic() override.
3D specifics: coordinates on graph.pos, zero-padded for virtual nodes; encoder raises ValueError on missing pos.
Datamodule schema in utils/schemas/datamodules.py + model schema in torch_api.py (composed from PretrainingMatchaModel, GraphMixin, GraphPretrainingMixin, the architecture mixin), both re-exported.
Lightning module in torch/models/pretraining/, registered with PretrainingModelRegistry, enc_* / pred_* prefixes, _build_encoder only — no forward / _get_per_layer_embeddings override.
Tests: per-model test_<name>_pretraining.py (registry, forward shapes, fit, encoder transfer), plus entries in test_encoder_parity.py and test_graph_pretraining_encoders.py.
3D only: classic ↔ pretraining datamodule parity test.
uv run pytest -k 'not gpu' passes locally.

If any of these steps feels awkward — e.g. the new pretraining path requires an override on BaseGraphPretrainingModel — flag it in the PR rather than adding it silently. The one-canonical-encoder + coords-on-pos invariants exist because parallel pretraining twins drifted the last three times.