matcha.torch.predictors.snn

Self-Normalizing Neural Network (SNN) predictor head with parallel MultiLn layers.

Classes

SNN	Self-Normalizing Neural Network (SNN) utility class.

Module Contents

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

Bases: matcha.torch.predictors.base_predictor.BasePredictor

Self-Normalizing Neural Network (SNN) utility class.

SNNs use SELU activations and AlphaDropout to maintain self-normalizing properties, as described in the paper “Self-Normalizing Neural Networks” (https://arxiv.org/abs/1706.02515).

This implementation uses MultiLn layers for parallel computation and averages the outputs across the num_parallel dimension for ensemble-like behavior.

It inherits from BasePredictor for common routines (e.g. forward pass).

It 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, goes directly from input to output.

• dropout (float) – dropout rate between all layers (uses AlphaDropout)

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

• num_parallel (int) – number of parallel heads in MultiLn layers (default: 8)

num_parallel = 8

encode(mol_features: torch.Tensor) → torch.Tensor

Extract the latent representation from all layers except the last.

Averages across the num_parallel dimension.

Parameters:

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

Returns:

averaged latent representation of shape (batch, hidden_dims[-1]).

Return type:

torch.Tensor

forward(mol_features: torch.Tensor) → torch.Tensor

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