matcha.torch.predictors.snn

Self-Normalizing Neural Network (SNN) predictor with TabM BatchEnsemble layers.

Classes

BatchEnsembleLinear

Parameter-efficient TabM-style BatchEnsemble linear layer.

SNN

Self-Normalizing Neural Network (SNN) predictor with TabM BatchEnsemble.

Module Contents

class matcha.torch.predictors.snn.BatchEnsembleLinear(in_features: int, out_features: int, num_parallel: int, *, first_layer: bool = False, bias: bool = True)[source]

Bases: torch.nn.Module

Parameter-efficient TabM-style BatchEnsemble linear layer.

Factorizes k parallel linear layers as a single shared weight W of shape (out, in) plus per-member rank-1 adapters R of shape (k, in) and S of shape (k, out), and an optional per-member bias of shape (k, 1, out). For a member i and input row x, the layer computes:

\[l_i(x) = \bigl((x \odot R_i)\, W^{\top}\bigr) \odot S_i + B_i,\]

which reduces the parameter cost of a naive parallel ensemble from k * in * out to in * out + k*(in + out) + k*out (bias).

The layer accepts either a 2D input (batch, in) — internally broadcast to (k, batch, in) — or a 3D input (k, batch, in). Output is always (k, batch, out).

Initialization follows the TabM recipe:

  • W — LeCun-normal (Kaiming-normal with mode='fan_in', nonlinearity='linear').

  • R — random Rademacher \pm 1 when first_layer=True to diversify the k submodels at initialization; deterministic 1 otherwise. Without the first-layer Rademacher init, all k branches collapse to identical outputs.

  • S — deterministic 1.

  • bias — 0.

Paper (TabM): Gorishniy et al. 2024, arXiv:2410.24210.

Parameters:

  • in_features (int) – Input feature dimensionality.

  • out_features (int) – Output feature dimensionality.

  • num_parallel (int) – Number of parallel ensemble members k.

  • first_layer (bool) – If True, initialize R with Rademacher \pm 1 values to diversify the submodels; else initialize to 1.

  • bias (bool) – Whether to include a per-member bias term.

in_features
out_features
num_parallel
first_layer = False
weight
R
S
reset_parameters() → None[source]

Apply the TabM init recipe to W, R, S, and bias.

num_extra_parameters() → int[source]

Return the number of per-member (non-shared) parameters.

extra_repr() → str[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x: torch.Tensor) → torch.Tensor[source]

Run the BatchEnsemble forward pass.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch, in_features) or (num_parallel, batch, in_features).

Returns:

Output tensor of shape (num_parallel, batch, out_features).

Return type:

torch.Tensor

Raises:

ValueError – If x is neither 2D nor 3D.

class matcha.torch.predictors.snn.SNN(input_dim: int, hidden_dims: list[int] | None, num_endpoints: int, dropout: float, num_parallel: int = 8)[source]

Bases: matcha.torch.predictors.base_predictor.BasePredictor

Self-Normalizing Neural Network (SNN) predictor with TabM BatchEnsemble.

Combines the self-normalizing recipe (SELU + AlphaDropout) from Klambauer et al. 2017 with the parameter-efficient BatchEnsemble factorization from TabM (Gorishniy et al. 2024): each linear layer is a BatchEnsembleLinear with a shared weight W and per-member rank-1 adapters R, S plus per-member bias. Outputs from the num_parallel ensemble members are averaged along the leading dimension at the end of forward() and (for a body with hidden layers) at the end of encode().

Papers:

It inherits from BasePredictor for common routines and is intended to be used inside a BaseClassicModel instance.

Parameters:

  • input_dim (int) – input feature dimensionality.

  • hidden_dims (list[int]) – shape of hidden layers in the predictor. If None or empty, goes directly from input to output.

  • num_endpoints (int) – number of endpoints (if multitasking) or classes (if classification) to predict.

  • dropout (float) – dropout rate applied between layers (uses AlphaDropout).

  • num_parallel (int) – number of parallel ensemble members k (default: 8). Must be >= 2; use the MLP predictor for the single-model case.

Raises:

ValueError – If num_parallel <= 1.

num_parallel = 8
layers
post
encode(mol_features: torch.Tensor) → torch.Tensor[source]

Extract the latent representation from all layers except the last.

Averages across the num_parallel dimension when a hidden body is present; otherwise returns the raw input.

Parameters:

mol_features – input tensor of shape (batch, input_dim).

Returns:

averaged latent representation of shape (batch, latent_dim).

Return type:

torch.Tensor

forward(mol_features: torch.Tensor) → torch.Tensor[source]

Run the full forward pass and average across the num_parallel dimension.

Parameters:

mol_features – input tensor of shape (batch, input_dim).

Returns:

predictions of shape (batch, num_endpoints).

Return type:

torch.Tensor