PR: Consolidate NODE as the Single Neural Tabular Path (CheMeleon + NODE)

[thomasATbayer]

Consolidate NODE as the Single Neural Tabular Path

Summary

This branch adds two complementary capabilities to Mother:

1. CheMeleon molecular representations: pretrained chemistry-GNN fingerprints that turn SMILES into learned molecular features.
2. NODE estimators: differentiable oblivious decision ensembles for regression and classification, with multiple heads, flow-based prediction, uncertainty estimation, embeddings, and MotherTuner integration.

SMILES -> CheMeleon fingerprint --------------------+
                                                     +-> Mother estimator
Tabular data / descriptors / fingerprints -----------+   (NODE, CatBoost, ...)
                                                         |
                                                         +-> point prediction
                                                         +-> flow distribution
                                                         +-> uncertainty

CheMeleon is a feature generator and is not coupled to NODE. The current CheMeleon example uses Mother's CatboostRegressorMother; the NODE example is maintained separately in examples/notebooks/05_advanced/05_NODE.ipynb.

NODE architecture

NODE uses differentiable oblivious decision trees (ODSTs). An oblivious tree uses one split feature and threshold at each depth. For a depth-$d$ tree:

$$ h(x) = \sum_{\ell=1}^{2^d} P(\ell \mid x) w_\ell $$

Here, $P(\ell \mid x)$ is the differentiable probability that sample $x$ reaches leaf $\ell$, and $w_\ell$ is that leaf's learned response. Soft routing allows the tree to be trained with gradient descent.

NODE layers use dense connections:

Layer 1: [x]           -> h1
Layer 2: [x, h1]       -> h2
Layer 3: [x, h1, h2]   -> h3

Sparse activations such as entmax15, sparsemax, entmoid15, and sparsemoid make feature selection and split gates differentiable while encouraging sparse routing.

NODE heads

Head	Purpose
subset	Compact NODE readout.
linear	Linear readout over tree features.
mlp	Nonlinear multilayer readout.
flow	Conditional predictive distribution $p(y \mid x)$ for regression.

from mother.ml.models.m_node import NODEClassifier, NODERegressor

reg = NODERegressor(head_type="mlp")
clf = NODEClassifier(head_type="subset")

Dropout controls

NODE has three independent dropout probabilities:

Parameter	Acts on	Granularity
input_dropout	Input received by an ODST layer	Individual feature channels.
tree_dropout	ODST outputs	Complete tree channels.
mlp_dropout	Hidden MLP layers	Individual hidden units.

For dropout probability $p$, inverted dropout uses a mask $m$:

$$ \tilde{h} = \frac{m}{1-p}h, \qquad m \sim \text{Bernoulli}(1-p) $$

so the expected value of the retained signal is preserved.

Two boolean controls determine where input and tree dropout are applied:

• input_dropout_only_input=True applies feature dropout only to original input features. False also applies it to dense between-layer inputs.
• tree_dropout_only_head=True applies tree dropout to the final representation before the head. False applies tree dropout inside every ODST layer.

The default behavior keeps feature dropout over dense inputs and applies tree dropout at the head boundary. Setting all dropout rates to zero is valid and produces a deterministic model.

Uncertainty and flow heads

For non-flow heads, MC dropout estimates knowledge uncertainty by repeating prediction with different dropout masks. For $T$ stochastic predictions:

$$ \bar{y}(x) = \frac{1}{T}\sum_{t=1}^{T} f_t(x) $$

$$ \sigma(x) = \sqrt{\frac{1}{T-1}\sum_{t=1}^{T}\left(f_t(x)-\bar{y}(x)\right)^2} $$

A flow head models a complete conditional distribution $p(y \mid x)$. Samples from one fixed flow represent data uncertainty. Different dropout passes produce different flow experts and can additionally represent knowledge uncertainty.

Plain-language uncertainty glossary

• Entropy summarizes how spread out the possible outcomes are. A narrow prediction distribution has less spread; a wide one has more.
• Nats are simply the unit used to report entropy, like metres for distance.
• Differential entropy is entropy for continuous numeric outcomes such as molecular properties. It describes a continuous density rather than a short list of class probabilities.
• A continuous density can be greater than 1, so differential entropy can be negative. This does not mean negative uncertainty; it usually indicates a very concentrated density on the current measurement scale.
• Numerical precision means that computers store only a finite number of digits. A mathematical zero may appear as -0.0000001 after floating-point rounding. Tiny negative values close to zero should be treated as zero within tolerance.

Let the expert distributions be $p_t(y \mid x)$ and their pooled distribution be:

$$ \bar{p}(y \mid x) = \frac{1}{T}\sum_{t=1}^{T}p_t(y \mid x) $$

For the BALD entropy decomposition:

$$ U_{data} = \frac{1}{T}\sum_{t=1}^{T} H[p_t] $$

$$ U_{total} = H[\bar{p}] $$

$$ U_{knowledge} = U_{total} - U_{data} $$

The knowledge term is the mutual-information disagreement signal and is non-negative up to floating-point precision. Differential entropy for a continuous flow can itself be negative; that is valid for a sharply concentrated density.

The combined API supports two knowledge methods:

stats = reg.predict_with_combined_uncertainty(
    X,
    knowledge_method="bald",
)

stats_emd = reg.predict_with_combined_uncertainty(
    X,
    knowledge_method="balsa_emd",
)

• bald uses the additive entropy decomposition above.
• balsa_emd uses a sampled Earth Mover's Distance disagreement score. It is not an entropy term and should not be added to data_uncertainty to reconstruct total_uncertainty.

Flow types and tuning

Flow type	Plain-language description	Main tuning parameters
NICE	Fast additive shifts between groups of values.	flow_transforms
NSF	Flexible learned spline curves; a strong general-purpose choice.	flow_transforms, flow_bins
RealNVP	Reversible stretch-and-shift coupling transformations.	flow_transforms
NAF	Expressive autoregressive transformations.	flow_transforms, flow_signal
UNAF	More flexible unconstrained autoregressive transformations.	flow_transforms, flow_signal
BPF	Smooth monotonic Bernstein polynomial transformations.	flow_degree
GMM	Several Gaussian components, useful for multimodal targets.	flow_components

Parameter meanings:

• flow_transforms: number of reversible transformation blocks.
• flow_bins: number of spline pieces for NSF.
• flow_degree: polynomial degree for BPF.
• flow_signal: hidden width for NAF and UNAF.
• flow_components: number of Gaussian peaks for GMM.

Start with NICE for speed or NSF for a strong general-purpose baseline. Increase flow complexity only when validation results support it.

Batching for wide one-layer NODE models

A model with num_layers=1 and many trees is wide and shallow:

X -> [many trees in one ODST layer] -> head -> prediction
          T1  T2  T3 ... T2048

batch_size batches rows, not trees. Every row in each minibatch passes through the full tree ensemble:

Full data -> [rows 1..B]      -> all trees -> loss/update
         -> [rows B+1..2B]    -> all trees -> loss/update
         -> [... ]             -> all trees -> loss/update

Reducing batch_size is usually the first memory adjustment for a wide model. Reducing num_trees also saves memory but changes model capacity. max_layers_retained affects dense skip connections only when num_layers > 1.

CheMeleon workflow

CheMeleon converts molecular graphs into fixed-width learned embeddings:

SMILES
  -> atoms and bonds
  -> pretrained CheMeleon GNN
  -> molecular fingerprint
  -> Mother-compatible estimator

from mother.feature_generation import CheMeleonFingerprintFactory
from mother.ml import CatboostRegressorMother

factory = CheMeleonFingerprintFactory(
    output_dim=2048,
    batch_size=128,
    device="cpu",
)
fingerprinter = factory.get_fingerprint_generator()

model = CatboostRegressorMother(
    iterations=500,
    max_depth=6,
    learning_rate=0.03,
    loss_function="RMSE",
    random_seed=42,
    verbose=False,
)

Use the repository's locked Chemprop extra before running the example from a fresh environment:

uv sync --extra chemprop

Notebook documentation

Updated examples/notebooks/05_advanced/05_NODE.ipynb with accessible, plain-language explanations and high-contrast callouts covering:

• NODE architecture, differentiable trees, heads, and molecular features.
• Major hyperparameters, flow types, flow tuning parameters, and batching for wide one-layer models.
• Input, tree, and MLP dropout, including placement options and max_layers_retained.
• Cross-layer dropout comparisons for one- and three-layer models, with deterministic zero-dropout checks.
• BALD and BALSA-EMD, including uncertainty intuition, differential entropy, nats, numerical precision, and ASCII diagrams.
• CheMeleon, Morgan fingerprints, and physicochemical descriptors.

Dependencies

The NODE optional dependency group contains all NODE-specific dependencies:

node = [
  "skorch>=1.4.0,<2",
  "zuko>=1.6.0,<2",
  "torch>=2.3.0,<3",
]

skorch is not a base dependency. The environment was verified with:

uv sync --extra node

Breaking changes

The neural modeling path is consolidated under NODE. Downstream imports should use:

mother.ml.models.node_utils
mother.ml.models.node_head_utils
mother.ml.models.node_balsa_acquisition

Standalone flow and MLP wrappers are no longer the supported public path. Use NODERegressor(head_type="flow") or NODERegressor(head_type="mlp") instead.

Validation focus

The branch includes focused coverage for:

• NODE regression and classification.
• Dropout placement and complete-tree masking.
• Zero-dropout deterministic behavior and uncertainty warnings.
• Cross-layer dropout comparisons.
• Flow uncertainty and BALD/BALSA-EMD methods.
• Estimator cloning and MotherTuner integration.
• CheMeleon fingerprint generation and CatBoost integration.

References

• Popov, Morozov, and Babenko. Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data.
• Wielopolski, Furman, and Zięba. NodeFlow: End-to-end Flexible Probabilistic Regression on Tabular Data.
• Gal, Islam, and Ghahramani. Deep Bayesian Active Learning with Image Data.
• Werner and Schmidt-Thieme. Bayesian Active Learning by Distributional Disagreement.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 24 changed files in this pull request and generated no new comments.

Suppressed comments (2)

scripts/check_docs_python_fences.py:26

• FENCE_RE ends with \s*$. Because \s matches newlines, the regex can accidentally consume content after the closing fence (including subsequent blocks) and make matching dependent on file layout. It should only allow trailing spaces on the closing fence line.

FENCE_RE = re.compile(
    r"(?ms)^[ \t]*```(?P<lang>python|py)(?P<attrs>[^\n]*)\n(?P<code>.*?)\n[ \t]*```\s*$"
)

src/mother/pipeline_utils.py:881

• The else branch clears fold_estimators to “free memory”, but fold_estimators is only appended to when return_estimators is True (see line 770), so it is already empty here. The branch is redundant and may mislead future readers about actual memory behavior.

    else:
        fold_estimators.clear()  # Clear the list to free memory if not returning
        module_logger.info("Returning performance_data only")

[Copilot]

Pull request overview

Copilot reviewed 18 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

src/mother/feature_generation/fp_gnn_gen.py:39

• The default CheMeleon checkpoint is downloaded over the network and written to a user cache without any integrity verification (e.g., SHA256). This is a supply-chain risk and can also lead to silent corruption (partial/cached files) being treated as valid. Consider verifying the downloaded file against a known hash (from Zenodo metadata) before moving it into the cache.

                with urllib.request.urlopen(_CHEMELEON_ZENODO_URL, timeout=60) as response:
                    shutil.copyfileobj(response, tmp_file)

            tmp_path.replace(_CHEMELEON_CACHE_PATH)  # atomic on POSIX; avoids partial files

[Copilot]

Pull request overview

Copilot reviewed 18 out of 24 changed files in this pull request and generated no new comments.

Suppressed comments (1)

src/mother/feature_generation/fp_gnn_gen.py:39

• The CheMeleon checkpoint is downloaded from a remote URL and cached locally, but the download path does not verify integrity (e.g., sha256) before trusting/executing the weights. This is a supply-chain risk if the remote asset changes or is tampered with.

Consider pinning the expected checksum (or fetching it from a trusted manifest) and validating the downloaded file before moving it into the cache path; fail fast with a clear error if verification fails.

                with urllib.request.urlopen(_CHEMELEON_ZENODO_URL, timeout=60) as response:
                    shutil.copyfileobj(response, tmp_file)

            tmp_path.replace(_CHEMELEON_CACHE_PATH)  # atomic on POSIX; avoids partial files

[Copilot]

Pull request overview

Copilot reviewed 18 out of 24 changed files in this pull request and generated no new comments.

Suppressed comments (1)

src/mother/feature_generation/fp_gnn_gen.py:147

• CheMeleonFingerprintTransformer.transform() currently materializes values via np.array(list(X), ...). When X is a numpy object array (e.g., shape (n, 1)), iterating yields row sub-arrays, so values becomes an array of arrays and isinstance(..., Chem.Mol) fails for valid molecules. This contradicts the new unit test that expects single-column object arrays to be flattened into molecules.

        values = np.array(list(X), dtype=object).reshape(-1)

[Copilot]

Pull request overview

Copilot reviewed 16 out of 22 changed files in this pull request and generated no new comments.

Suppressed comments (2)

src/mother/feature_generation/fp_gnn_gen.py:148

• CheMeleonFingerprintTransformer.transform() builds values via np.array(list(X), ...), which turns a 2D object array like np.array([[mol]], dtype=object) into an array of row-arrays (not Chem.Mol). That makes valid_mask false and returns NaNs, contradicting the intended single-column flattening behavior (and the added unit test).

    def transform(self, X: Iterable) -> np.ndarray:
        check_is_fitted(self, "is_fitted_")

        values = np.array(list(X), dtype=object).reshape(-1)
        out = np.full((len(values), self.output_dim), np.nan, dtype=np.float32)
        if len(values) == 0:
            return out

        valid_mask = np.array([isinstance(compound, Chem.Mol) for compound in values], dtype=bool)

src/mother/ml/models/node_head_utils.py:106

• nn.Parameter is a subclass of torch.Tensor, so the isinstance(obj, torch.Tensor) branch will run first and the later isinstance(obj, nn.Parameter) branch is unreachable. If the intent is to avoid moving Parameters via .to(), the checks should be reordered.

        if isinstance(obj, torch.Tensor):
            if obj.device != device:
                return obj.to(device)
            return obj

[Copilot]

Pull request overview

Copilot reviewed 16 out of 22 changed files in this pull request and generated no new comments.

Suppressed comments (3)

src/mother/feature_generation/fp_gnn_gen.py:24

• get_default_chemeleon_checkpoint() downloads a model checkpoint from a remote URL and caches it locally, but it doesn’t verify integrity (e.g., SHA256) before trusting the file. This is a supply-chain and reliability risk (corrupt/truncated downloads, transparent proxy tampering, etc.).

Consider publishing/embedding an expected hash (or fetching a Zenodo-provided checksum) and verifying it after download; alternatively allow users to opt out of network downloads by default and require an explicit checkpoint_path.

def get_default_chemeleon_checkpoint() -> Path:
    """Return path to chemeleon_mp.pt, downloading from Zenodo on first use."""
    if not _CHEMELEON_CACHE_PATH.exists():
        _CHEMELEON_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)

src/mother/ml/models/node_utils.py:649

• These warnings are emitted from inside the library, but they don’t set stacklevel, so users will see the warning originating from node_utils.py rather than their own call site. Adding stacklevel=2 makes the warning location more actionable for callers.

        if input.shape[0] < 256:
            warn(
                "Data-aware initialization is performed on less than 256 data points. "
                "This may reduce threshold initialization quality on some datasets. "
                "Prefer at least 256 samples for stable initialization; 512+ can be more robust "

src/mother/ml/models/node_utils.py:785

• This warning also lacks stacklevel, which makes it point at the library code instead of the caller location that passed an invalid max_layers_retained. Using stacklevel=2 improves debuggability.

        if effective_max_layers_retained is not None and effective_max_layers_retained < 1:
            warn(
                f"max_layers_retained={effective_max_layers_retained} is smaller than 1; "
                "using max_layers_retained=1 to keep dimensions consistent."
            )

[Copilot]

Pull request overview

Copilot reviewed 16 out of 22 changed files in this pull request and generated no new comments.