Add **Neural Oblivious Decision Ensembles (NODE)** and MLP as a first-class estimator in the Mother framework, alongside CatBoost and TabPFN.

[thomasATbayer]

feat: Add NODE (Neural Oblivious Decision Ensembles) to Mother

Summary

Add Neural Oblivious Decision Ensembles (NODE) and MLP as first-class estimators in the
Mother framework, alongside CatBoost and TabPFN.

NODE brings differentiable decision trees to tabular data — combining the
interpretability of tree ensembles with the flexibility of deep learning, and
uniquely offering calibrated probabilistic predictions via normalising flows.

---

Motivation

Mother currently supports gradient-boosted trees (CatBoost) and in-context
learners (TabPFN). CatBoost does provide uncertainty decomposition:

• Regression — RMSEWithUncertainty (aleatoric) + virtual ensembles (epistemic)
• Binary classification — virtual ensembles give both data and knowledge uncertainty

However, CatBoost's uncertainty is unavailable for multiclass classification and
does not model the full predictive distribution.

NODE complements this by offering:

• Full predictive distributions — sample from $p(y \mid x)$ at any quantile
  via normalising flow heads (NICE, NSF, …)
• Decomposed uncertainty for multiclass classification — MC Dropout epistemic
  uncertainty works across all classification tasks, not just binary
• Learned tree-layer embeddings — transferable representations for downstream
  tasks (UMAP, clustering, standalone heads)
• Architecture flexibility — swap head types (subset / linear / MLP / flow)
  without retraining the tree backbone

---

Proposed Additions

Core Models

Class	Task
NODERegressor	Regression (single- and multi-target, NaN-safe)
NODEClassifier	Binary / multiclass / multi-label classification

Head Architectures

Head	When to use
subset	Default for classification; zero overhead
linear	Lightweight baseline
mlp	Default for regression; configurable depth
flow	Probabilistic regression with full density estimation

Standalone Heads (no NODE backbone needed)

• MLPHeadRegressor / MLPHeadClassifier
• FlowHeadRegressor

All heads are sklearn-compatible, support MotherTuner, and include
auto dimension detection (no need to specify input_dim/output_dim).

Uncertainty Quantification

# Decomposed uncertainty — unique to NODE with flow head
pred, knowledge_unc, data_unc = model.predict_with_combined_uncertainty(X_test)
# knowledge_unc → reducible with more data (epistemic)
# data_unc      → irreducible noise (aleatoric)

GNN Fingerprints (CheMeleon)

Add CheMeleonFingerprintFactory / CheMeleonFingerprinter — sklearn transformers
that produce 2048-d molecular embeddings from SMILES using a pretrained
bond-message-passing network (chemprop >= 2.0).

---

Dependencies (optional extras)

[tool.poetry.extras]
node = ["torch", "skorch", "zuko"]
chemeleon = ["chemprop"]

---

References

• Popov et al. (2020). Neural Oblivious Decision Ensembles for Deep Learning on
  Tabular Data. ICLR 2020.
• Wielopolski et al. (2024). NodeFlow: Towards End-to-end Flexible Probabilistic
  Regression on Tabular Data. Entropy, 26(7).

🌲 NODE – What's New in This PR

PR #440 — Introduce NODE to the Mother package

---

🎯 At a Glance

This PR adds Neural Oblivious Decision Ensembles (NODE) to Mother — a deep-learning model for tabular data that rivals gradient-boosted trees while offering probabilistic predictions, uncertainty decomposition, and learned embeddings.

from mother.ml.models.m_node import NODERegressor

reg = NODERegressor(head_type="flow", flow_type="NSF", max_epochs=100)
reg.fit(X_train, y_train)
preds = reg.predict(X_test)

---

✨ Highlights

1. Four Head Architectures

Choose the right prediction layer for your task:

Head	Best for	Special powers
subset	Classification (default)	Lightweight, fast
linear	Simple baselines	Minimal overhead
mlp	Regression (default)	Configurable depth, BatchNorm, GELU
flow	Probabilistic regression	Full density estimation via normalising flows

2. NodeFlow — Probabilistic Predictions

The flow head implements the NodeFlow architecture, combining NODE with conditional normalising flows (powered by zuko):

• 4 flow types: CNF, MAF, NSF (recommended), NICE
• Sample from $p(y \mid x)$ — credible intervals, quantiles, full density
• Massive speed improvement — optimised vectorised mode estimation via shared compute_flow_mode_and_uncertainty()

3. Uncertainty Decomposition

Three levels of uncertainty quantification:

Method	What it captures	Requirements
MC Dropout	Epistemic (knowledge)	Any head + input_dropout > 0
Flow sampling	Aleatoric (data)	Flow head
Combined	Both, decomposed	Flow head + dropout

preds, knowledge_unc, data_unc = model.predict_with_combined_uncertainty(X)
# knowledge_unc → reducible with more data
# data_unc      → irreducible noise

4. Standalone Heads — Use Anywhere

The MLP and Flow heads are extracted into independent sklearn-compatible estimators:

from mother.ml.models.m_heads import MLPHeadRegressor, FlowHeadRegressor, MLPHeadClassifier

• Auto dimension detection — DimensionSetter callback infers input_dim/output_dim from data
• Batteries included — AdamW, EarlyStopping, ReduceLROnPlateau built-in
• MotherTuner compatible — suggest_hyperparameters() for Optuna integration
• Perfect for pre-computed embeddings or lightweight deployment

5. Adaptive Sparse Activations

NODE uses entmax15 and entmoid15 by default — sparse activation functions that produce exactly-zero outputs for irrelevant features, improving interpretability without sacrificing performance.

6. Multi-Task with NaN Masking

Train on incomplete label matrices — NODE's custom get_loss() automatically masks NaN targets during backpropagation:

reg = NODERegressor(target_type="multi_target")
reg.fit(X, y_with_nans)  # NaN values are masked, not propagated

7. Learned Embeddings

Extract tree-layer representations for downstream tasks:

embeddings = model.get_embeddings(X)  # (n_samples, num_layers * num_trees * tree_dim)

Use for UMAP visualisation, clustering, transfer learning, or as input to standalone heads.

---

🏗️ Architecture

┌─────────────────────────────────────────────────────┐
│                  NODERegressor                       │
│                  NODEClassifier                      │
├─────────────────────────────────────────────────────┤
│  Embedding Layer  →  Dense ODST Block  →  Head      │
│  (BatchNorm,         (differentiable      (subset,  │
│   categorical         oblivious trees,     linear,  │
│   embeddings)         sparse activations)  mlp,     │
│                                            flow)    │
└─────────────────────────────────────────────────────┘
         ↓ get_embeddings()          ↓ predict()
    [n, tree_dim]            point predictions / distributions

---

📦 New Files

File	What	Lines
mother/ml/models/m_node.py	NODERegressor, NODEClassifier, CompletePyTorchTabularNODE	~2800
mother/ml/models/m_node_utils.py	ODST trees, sparse activations, embeddings	~760
mother/ml/models/m_heads.py	MLPHeadRegressor, MLPHeadClassifier, FlowHeadRegressor	~1400
mother/ml/models/m_head_utils.py	Shared flow mode/uncertainty computation	~50
test/unit/test_node_unit.py	Comprehensive unit tests	—
mkdocs/docs/mother/node.md	Full documentation page	~480
examples/notebooks/example_NODE.ipynb	Interactive tutorial (11 sections)	—

---

🔧 Dependencies

Added as optional [node] extra in pyproject.toml:

[tool.poetry.extras]
node = ["torch", "skorch", "zuko"]

torch   >= 2.3.0
skorch  >= 1.2.0
zuko   >= 1.4.1

---

🔑 Key Design Decisions

1. Skorch wrapping — NODE uses skorch's NeuralNetRegressor/NeuralNetClassifier for full sklearn compatibility (pipelines, cross_val_score, cloning)
2. Auto-detection callbacks — InputOutputShapeSetter and DimensionSetter infer dimensions from data, so users never need to manually specify them
3. Shared utilities — compute_flow_mode_and_uncertainty() is used by both NODE and standalone heads, eliminating code duplication
4. MC Dropout in eval mode — Dropout layers are selectively set to training mode while keeping BatchNorm in eval mode, ensuring stable uncertainty estimates

---

📊 What Comes Next

• GPU acceleration testing and benchmarks
• Integration with Mother's cross_validate framework
• Categorical feature embedding examples with real datasets
• Performance comparison against CatBoost, XGBoost, TabPFN on standard benchmarks

NODE & Standalone Heads — In-Depth Guide

A comprehensive reference for NODE (Neural Oblivious Decision Ensembles) and the standalone head estimators in the Mother package.

---

Table of Contents

1. NODE Overview
2. Architecture Deep Dive
3. NODERegressor
4. NODEClassifier
5. Head Types
6. Flow Head & Normalising Flows
7. Uncertainty Estimation
8. Standalone Heads
   • MLPHeadRegressor
   • MLPHeadClassifier
   • FlowHeadRegressor
9. Auto Dimension Detection
10. Hyperparameter Tuning
11. Advanced Topics

---

NODE Overview

NODE implements differentiable oblivious decision trees that can be trained end-to-end with gradient descent. Unlike traditional tree ensembles (XGBoost, CatBoost), NODE's trees are smooth and differentiable, enabling:

• Gradient-based optimisation — no need for greedy splitting heuristics
• Representation learning — tree layers produce learned embeddings
• Flexible heads — swap the prediction layer without retraining the tree backbone
• Probabilistic outputs — flow heads model full conditional distributions

When to use NODE

Scenario	Recommendation
Small tabular dataset (<1000 rows)	CatBoost or TabPFN may be better
Medium-large tabular data	✅ NODE is competitive
Need uncertainty estimates	✅ NODE with flow head
Need learned representations	✅ NODE embeddings
Need probabilistic predictions	✅ NODE flow head
Categorical-heavy data	✅ NODE with categorical embeddings
Real-time inference needed	CatBoost is faster at inference

---

Architecture Deep Dive

┌───────────────────────────────────────────────────────────────┐
│                   CompletePyTorchTabularNODE                  │
├───────────────────────────────────────────────────────────────┤
│                                                               │
│  ┌─────────────────────┐                                     │
│  │  Embedding Layer     │  BatchNorm (optional)               │
│  │  (Embedding1dLayer)  │  Categorical embeddings (optional)  │
│  │                      │  Input dropout                      │
│  └──────────┬──────────┘                                     │
│             ▼                                                 │
│  ┌─────────────────────┐                                     │
│  │  Dense ODST Block    │  num_trees differentiable trees     │
│  │  (DenseODSTBlock)    │  Each tree: depth splits            │
│  │                      │  Sparse feature selection            │
│  │                      │  Dense connections between layers    │
│  │                      │  Tree dropout (optional)            │
│  └──────────┬──────────┘                                     │
│             ▼                                                 │
│  ┌─────────────────────┐                                     │
│  │  Head Layer          │  One of: subset, linear, mlp, flow  │
│  │                      │  Converts tree output → predictions │
│  └──────────┬──────────┘                                     │
│             ▼                                                 │
│         Predictions                                           │
└───────────────────────────────────────────────────────────────┘

Sparse Activations

NODE uses sparse activation functions that produce exact zeros:

Function	Used for	Description
entmax15	Feature selection (choice_function)	Projects onto the 1.5-simplex; exactly sparse
sparsemax	Feature selection (alternative)	Projects onto the probability simplex
entmoid15	Bin boundaries (bin_function)	Sparse sigmoid-like function
sparsemoid	Bin boundaries (alternative)	Hard sigmoid variant

The defaults (entmax15 + entmoid15) produce trees that use only a subset of features at each split, improving interpretability.

ODST (Oblivious Decision Stump Tree)

Each tree in the ensemble is an oblivious decision tree:

• All nodes at the same depth share the same split feature and threshold
• This makes the tree equivalent to a lookup table of size $2^{\text{depth}}$
• Each leaf stores a learned response vector

The tree output is computed as:

$$\text{output} = \sum_{l=1}^{2^d} w_l \cdot \mathbb{1}[\text{sample reaches leaf } l]$$

where $w_l$ is the response vector for leaf $l$ and the indicator is soft (differentiable) via the bin function.

Dense Connections

The DenseODSTBlock stacks multiple ODST layers with dense (DenseNet-style) connections:

Layer 1 input:  [X]
Layer 1 output: [h1]
Layer 2 input:  [X, h1]
Layer 2 output: [h2]
Layer 3 input:  [X, h1, h2]
...

This enables later layers to refine predictions from earlier layers.

---

NODERegressor

from mother.ml.models.m_node import NODERegressor

Constructor Parameters

Core Architecture

Parameter	Type	Default	Description
num_trees	int	2048	Number of trees in the ensemble
depth	int	6	Tree depth (each tree has $2^{\text{depth}}$ leaves)
num_layers	int	1	Number of stacked ODST layers

Head Configuration

Parameter	Type	Default	Description
head_type	str	"mlp"	"subset", "linear", "mlp", or "flow"
mlp_hidden_dims	list	[128, 64, 32]	Hidden layer sizes for MLP head
mlp_activation	str	"ReLU"	"ReLU", "GELU", or "LeakyReLU"
flow_type	str	"CNF"	"CNF", "MAF", "NSF", "NICE"
flow_transforms	int	3	Number of flow transformation layers
flow_bins	int	8	Number of spline bins (NSF only)

Dropout & Regularisation

Parameter	Type	Default	Description
input_dropout	float	0.1	Dropout on input features (needed for MC Dropout UQ)
tree_dropout	float	0.0	Dropout on tree outputs
mlp_dropout	float	0.1	Dropout in MLP head layers
embedding_dropout	float	0.0	Dropout on categorical embeddings

Training

Parameter	Type	Default	Description
max_epochs	int	100	Training epochs
lr	float	0.01	Learning rate
batch_size	int	128	Batch size
optimizer	type	Adam	Optimizer class
criterion	type	MSELoss	Loss function
device	str	auto	"cpu" or "cuda"

Multi-Target

Parameter	Type	Default	Description
target_type	str	"single_target"	"single_target" or "multi_target"
task_weights	list	None	Per-target loss weights

Methods

Method	Returns	Description
fit(X, y)	self	Train the model
predict(X)	ndarray	Point predictions (mode for flow heads)
predict_uncertainty(X, num_samples)	DataFrame	Predictions with uncertainty estimates
predict_with_combined_uncertainty(X)	tuple	Decomposed uncertainty (flow only)
predict_flow_head(X, num_samples)	ndarray	Raw flow mode predictions
get_embeddings(X)	ndarray	Tree-layer representations
suggest_hyperparameters(X, y, trial)	dict	Optuna hyperparameter space

Multi-Target with NaN Masking

NODE's get_loss() method automatically detects and masks NaN values in multi-target regression:

reg = NODERegressor(
    head_type="mlp",
    target_type="multi_target",
    task_weights=[1.0, 2.0, 0.5],  # optional per-target weighting
)

# y can contain NaN values — they are masked during loss computation
y_train[some_mask] = np.nan
reg.fit(X_train, y_train)

The loss is computed only over non-NaN targets:

$$\mathcal{L} = \frac{1}{|\text{valid}|} \sum_{i \in \text{valid}} w_i \cdot \text{MSE}(\hat{y}_i, y_i)$$

---

NODEClassifier

from mother.ml.models.m_node import NODEClassifier

Supports:

• Binary classification (CrossEntropyLoss)
• Multiclass classification (CrossEntropyLoss)
• Multi-label classification (BCEWithLogitsLoss)

Key Differences from Regressor

Feature	Regressor	Classifier
Default head	mlp	subset
Default input_dropout	0.1	0.0
Default criterion	MSELoss	CrossEntropyLoss
Flow head	✅	❌ Not supported
predict_proba()	❌	✅
predict_with_combined_uncertainty()	✅ (flow)	❌

Class Weights

import torch.nn as nn

# Inverse frequency weighting
weights = torch.tensor([1.0, 5.0])  # e.g., 5:1 minority class upweight

clf = NODEClassifier(
    criterion=nn.CrossEntropyLoss,
    criterion__weight=weights,
)

Multi-Label Classification

clf = NODEClassifier(
    model_type="classification_multilabel",
    criterion=nn.BCEWithLogitsLoss,
)

# y should be (n_samples, n_labels) with 0/1 values
clf.fit(X, y_multilabel)
preds = clf.predict(X_test)     # (n_test, n_labels), thresholded at 0.5
probas = clf.predict_proba(X_test)  # (n_test, n_labels), sigmoid outputs

---

Head Types

Subset Head

The default for classification. Selects a weighted subset of tree outputs and averages them:

$$\hat{y} = \frac{1}{|\text{subset}|} \sum_{t \in \text{subset}} w_t \cdot \text{tree}_t(x)$$

• No additional parameters
• Fastest inference
• Good for classification where tree outputs are already class-aligned

Linear Head

Single linear transformation of the flattened tree output:

$$\hat{y} = W \cdot \text{flatten}(\text{trees}(x)) + b$$

• Minimal added complexity
• Good baseline to isolate tree-layer contribution

MLP Head

Multi-layer perceptron with configurable architecture:

tree_output → Linear → BatchNorm → Activation → Dropout → ... → Linear → output

Parameters:

• mlp_hidden_dims=[128, 64, 32] — layer sizes
• mlp_activation="ReLU" — nonlinearity
• mlp_dropout=0.1 — regularisation

Kaiming initialisation is used for all linear layers.

Flow Head

Conditional normalising flow that models $p(y \mid x)$:

tree_output → (Optional Tanh MLP) → Conditional Flow → Distribution

The flow transforms a simple base distribution (standard normal) into the target distribution through a series of invertible transformations.

---

Flow Head & Normalising Flows

How It Works

1. NODE tree layers produce embeddings $h = f_\theta(x)$
2. The flow head defines a conditional distribution $p(y \mid h)$
3. Training minimises negative log-likelihood: $\mathcal{L} = -\mathbb{E}[\log p(y \mid h)]$
4. At inference, the mode is found by sampling and selecting the highest log-probability point

Flow Types

Type	Full Name	Architecture	Recommendation
CNF	Continuous Normalising Flow	ODE-based, flexible	Default, good general choice
MAF	Masked Autoregressive Flow	Autoregressive transforms	Good for complex densities
NSF	Neural Spline Flow	Monotonic rational-quadratic splines	Best quality/speed tradeoff
NICE	Non-linear Independent Components Estimation	Volume-preserving transforms	Fast but less expressive

Configuration

reg = NODERegressor(
    head_type="flow",
    flow_type="NSF",         # Recommended
    flow_transforms=3,       # Number of transformation layers (more = more expressive)
    flow_bins=8,             # Spline bins (NSF only; more = finer density resolution)
)

Target Standardisation

Critical: Flow heads require standardised targets for numerical stability of log-probability calculations:

from sklearn.preprocessing import StandardScaler

y_scaler = StandardScaler()
y_train_scaled = y_scaler.fit_transform(y_train.reshape(-1, 1)).ravel()

reg = NODERegressor(head_type="flow", flow_type="NSF")
reg.fit(X_train, y_train_scaled)

preds_scaled = reg.predict(X_test)
preds = y_scaler.inverse_transform(preds_scaled.reshape(-1, 1)).ravel()

---

Uncertainty Estimation

MC Dropout (All Heads)

Multiple forward passes with dropout enabled produce a distribution of predictions. The spread measures model confidence.

Requirements: input_dropout > 0 (or tree_dropout > 0)

model = NODERegressor(input_dropout=0.1, head_type="mlp")
model.fit(X_train, y_train)

df = model.predict_uncertainty(X_test, num_samples=100)
# DataFrame columns:
#   mean_predictions     — average across MC passes
#   knowledge_uncertainty — std (or IQR) across MC passes
#   data_uncertainty      — None for non-flow heads
#   total_uncertainty     — same as knowledge_uncertainty

Flow Uncertainty

For flow heads, predict_uncertainty() provides both data and knowledge uncertainty:

model = NODERegressor(head_type="flow", input_dropout=0.1)
model.fit(X_train, y_train)

df = model.predict_uncertainty(X_test, num_samples=100)
# data_uncertainty is now populated (from flow distribution width)
# knowledge_uncertainty from MC Dropout

Combined Decomposition

The most powerful uncertainty method — decomposes into aleatoric and epistemic:

preds, knowledge_unc, data_unc = model.predict_with_combined_uncertainty(
    X_test,
    num_mc_samples=50,      # MC Dropout passes
    num_flow_samples=100,   # Samples per flow evaluation
)

Algorithm:

1. Single forward pass WITHOUT dropout → find mode, compute $-\log p(\text{mode})$ = data uncertainty
2. num_mc_samples forward passes WITH dropout → evaluate $\log p(\text{mode})$ under each mask
3. Knowledge uncertainty = IQR of $-\log p(\text{mode})$ across MC passes
4. Total = data + knowledge

Return All Details

results = model.predict_with_combined_uncertainty(X_test, return_all=True)
# results is a dict:
#   'predictions'           — mean predictions
#   'knowledge_uncertainty' — epistemic
#   'data_uncertainty'      — aleatoric
#   'total_uncertainty'     — sum of both
#   'mc_means'              — per-MC-pass means
#   'mc_stds'               — per-MC-pass stds

Quantile Predictions

df, quantiles = model.predict_uncertainty(
    X_test,
    return_quantiles=True,
    quantiles=[0.025, 0.5, 0.975],
)
# quantiles.shape = (n_test, 3)

---

Standalone Heads

The MLP and Flow heads are available as independent estimators that don't require NODE's tree layers. These are ideal when:

• You have pre-computed feature embeddings (e.g. from a foundation model, autoencoder, or another model's get_embeddings())
• You want a lightweight neural net without the ODST overhead
• You want to use the flow architecture on arbitrary features

All standalone heads:

• Implement the sklearn API (fit, predict, predict_proba)
• Implement Mother's AbstractMotherPipeline interface
• Support auto dimension detection via DimensionSetter
• Support hyperparameter tuning via suggest_hyperparameters()
• Include sensible defaults (AdamW, EarlyStopping, LR scheduling)

---

MLPHeadRegressor

from mother.ml.models.m_heads import MLPHeadRegressor

A multi-layer perceptron regressor with built-in training best practices.

Constructor

MLPHeadRegressor(
    input_dim=1,              # Auto-detected from data
    output_dim=1,             # Auto-detected from data
    hidden_dims=[256, 128, 64],  # Layer sizes
    dropout=0.05,             # Dropout rate
    batch_norm=True,          # BatchNorm between layers
    activation="ReLU",        # "ReLU", "GELU", or "LeakyReLU"
    max_epochs=500,           # Training epochs
    lr=0.005,                 # Learning rate
    # Built-in defaults:
    # optimizer=AdamW (weight_decay=0.01)
    # train_split=ValidSplit(cv=0.1)
    # callbacks=[DimensionSetter, EarlyStopping(patience=20), ReduceLROnPlateau]
)

Example

from mother.ml.models.m_heads import MLPHeadRegressor
from sklearn.datasets import make_regression

X, y = make_regression(n_samples=500, n_features=20, noise=5)

reg = MLPHeadRegressor(
    hidden_dims=[128, 64],
    dropout=0.1,
    activation="GELU",
    max_epochs=200,
)
reg.fit(X.astype("float32"), y.astype("float32"))

preds = reg.predict(X_test)

Note: No need to specify input_dim or output_dim — the DimensionSetter callback detects them automatically from the training data.

---

MLPHeadClassifier

from mother.ml.models.m_heads import MLPHeadClassifier

Identical architecture to MLPHeadRegressor but configured for classification:

• Uses CrossEntropyLoss (expects integer class labels)
• Provides predict_proba() for probability estimates

Constructor

MLPHeadClassifier(
    input_dim=1,              # Auto-detected
    output_dim=1,             # Auto-detected (number of classes)
    hidden_dims=[256, 128, 64],
    dropout=0.05,
    batch_norm=True,
    activation="ReLU",
    max_epochs=500,
    lr=0.005,
    # criterion=CrossEntropyLoss (default)
)

Example

from mother.ml.models.m_heads import MLPHeadClassifier

clf = MLPHeadClassifier(
    hidden_dims=[128, 64],
    max_epochs=100,
)
clf.fit(X_train, y_train)  # y must be integer labels

predictions = clf.predict(X_test)       # class labels
probabilities = clf.predict_proba(X_test)  # (n, n_classes)

---

FlowHeadRegressor

from mother.ml.models.m_heads import FlowHeadRegressor

A normalising flow regressor that models the full conditional distribution $p(y \mid x)$.

Constructor

FlowHeadRegressor(
    input_dim=1,              # Auto-detected
    output_dim=1,             # Auto-detected
    flow_type="CNF",          # "CNF", "MAF", "NSF", "NICE"
    flow_transforms=3,        # Number of transformation layers
    flow_bins=8,              # Spline bins (NSF only)
    max_epochs=100,
    lr=0.001,
    # Loss: negative log-likelihood (automatic via get_loss())
)

Key Methods

Method	Returns	Description
fit(X, y)	self	Train with NLL loss
predict(X, num_samples)	ndarray	Point predictions (mode via MAP estimation)
predict_flow(X, num_samples)	ndarray	Full distribution samples (n, num_samples, d)
predict_uncertainty(X)	DataFrame	Uncertainty estimates (MC Dropout)
get_loss(y_pred, y_true)	Tensor	Negative log-likelihood
suggest_hyperparameters(X, y, trial)	dict	Optuna search space

Example: Full Probabilistic Workflow

from mother.ml.models.m_heads import FlowHeadRegressor
import numpy as np

reg = FlowHeadRegressor(
    flow_type="NSF",
    flow_transforms=3,
    max_epochs=200,
    lr=0.001,
)
reg.fit(X_train, y_train)

# Point predictions (MAP estimate)
preds = reg.predict(X_test, num_samples=200)

# Full distribution sampling
samples = reg.predict_flow(X_test, num_samples=1000)
# samples.shape = (n_test, 1000, output_dim)

# Compute statistics
mean = samples.mean(axis=1)
std = samples.std(axis=1)
q05 = np.quantile(samples, 0.05, axis=1)
q95 = np.quantile(samples, 0.95, axis=1)

# Credible intervals
print(f"90% CI width: {(q95 - q05).mean():.3f}")

---

Auto Dimension Detection

Both NODE and standalone heads use callback-based auto-detection of input/output dimensions.

For NODE: InputOutputShapeSetter

Runs at on_train_begin and handles:

• Input dimension from X.shape[1]
• Output dimension from y (regression: shape, classification: unique values)
• Categorical feature detection from DataFrame column dtypes
• Label encoding and embedding dimension calculation for categoricals

# Categorical features are auto-detected from category dtype
df = pd.DataFrame({"num1": [1.0, 2.0], "cat1": pd.Categorical(["a", "b"])})
reg = NODERegressor()
reg.fit(df, y)  # Automatic: num1 → continuous, cat1 → embedded

For Standalone Heads: DimensionSetter

Simpler version that detects:

• input_dim from X.shape[1]
• output_dim from y (shape for regression, unique values for classification)

# No need to specify dimensions!
reg = MLPHeadRegressor()  # input_dim=1, output_dim=1 are placeholders
reg.fit(X, y)  # DimensionSetter detects actual dimensions

---

Hyperparameter Tuning

All NODE and standalone head estimators expose a suggest_hyperparameters() (or get_hyperparameter_space()) method compatible with Mother's MotherTuner.

NODE Tuning Space

Parameter	Range	Scale
num_trees	256 – 4096	log
depth	3 – 8	linear
num_layers	1 – 3	linear
input_dropout	0.0 – 0.3	linear
tree_dropout	0.0 – 0.5	linear
lr	1e-4 – 1e-1	log
head_type	mlp / subset / linear / flow	categorical
mlp_hidden_dims	adaptive	derived
flow_type	CNF / MAF / NSF / NICE	categorical

When tune_head=True (default), head-specific parameters are included in the search space.

Standalone Head Tuning Space

Parameter	Range	Scale
num_hidden_layers	1 – 4	linear
hidden_dim_1	adaptive to input_dim	linear
dropout	0.0 – 0.5	linear
batch_norm	True / False	categorical
activation	ReLU / GELU / LeakyReLU	categorical
lr	1e-5 – 1e-2	log

MotherTuner Integration

from mother.optimization import MotherTuner

# NODE
tuner = MotherTuner(
    estimator=NODERegressor(max_epochs=50, device="cpu"),
    X=X_train, y=y_train,
    n_trials=50,
    scoring="neg_mean_squared_error",
)

# Standalone head
tuner = MotherTuner(
    estimator=MLPHeadRegressor(max_epochs=100),
    X=X_train, y=y_train,
    n_trials=30,
    scoring="neg_mean_squared_error",
)

tuner.run()
best = tuner.best_estimator

---

Advanced Topics

Embeddings as Features

Use NODE's tree layers as a feature extractor, then feed embeddings into a standalone head or other model:

# Train NODE
node = NODEClassifier(num_trees=2048, max_epochs=100)
node.fit(X_train, y_train)

# Extract embeddings
emb_train = node.get_embeddings(X_train)
emb_test = node.get_embeddings(X_test)

# Use embeddings with a standalone flow head
flow = FlowHeadRegressor(flow_type="NSF", max_epochs=200)
flow.fit(emb_train, y_train_reg)

MC Dropout Implementation Details

NODE implements MC Dropout carefully to ensure correct uncertainty estimates:

1. Eval mode for BatchNorm — BatchNorm uses running statistics (not batch statistics)
2. Train mode for Dropout — Only nn.Dropout modules are switched to training mode
3. Temporary dropout override — The temporary_dropout_override context manager allows using a different dropout rate for UQ than was used during training

DataFrame Support

NODE natively supports pandas DataFrames:

df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(X.shape[1])])
reg = NODERegressor()
reg.fit(df, y)         # Works!
preds = reg.predict(df)  # Works!

Categorical columns with pd.Categorical dtype are automatically detected and embedded.

Device Management

# Auto-detect GPU
reg = NODERegressor()  # Uses CUDA if available

# Force CPU
reg = NODERegressor(device="cpu")

# Specify GPU
reg = NODERegressor(device="cuda:0")

Skorch Integration

NODE estimators are built on skorch, so all skorch features work:

from skorch.callbacks import EarlyStopping, LRScheduler

reg = NODERegressor(
    callbacks=[
        EarlyStopping(patience=10, monitor="valid_loss"),
        LRScheduler(policy="CosineAnnealingLR", T_max=50),
    ],
    train_split=skorch.dataset.ValidSplit(cv=0.15),
)

---

Quick Reference

Imports

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

# Standalone heads
from mother.ml.models.m_heads import MLPHeadRegressor, MLPHeadClassifier, FlowHeadRegressor

# Tuning
from mother.optimization import MotherTuner

Common Patterns

# Regression with uncertainty
reg = NODERegressor(head_type="flow", flow_type="NSF", input_dropout=0.1)
reg.fit(X, y)
df = reg.predict_uncertainty(X_test)

# Classification with probabilities
clf = NODEClassifier(input_dropout=0.1)
clf.fit(X, y)
probas = clf.predict_proba(X_test)
df = clf.predict_uncertainty(X_test)

# Standalone MLP on embeddings
mlp = MLPHeadRegressor(hidden_dims=[256, 128])
mlp.fit(embeddings, targets)

# Standalone flow for probabilistic predictions
flow = FlowHeadRegressor(flow_type="NSF")
flow.fit(X, y)
samples = flow.predict_flow(X_test, num_samples=1000)

[thomasATbayer]

Uncertainty for NODE and MLP models

This page details how uncertainty is estimated for the neural tabular models in
Mother — NODE (Neural Oblivious Decision Ensembles) and the standalone
MLP / Flow heads — and how their outputs map onto the shared
predict_uncertainty interface.

Each model exposes fundamentally different kinds of uncertainty. Rather than
forcing every backend into one definition, Mother keeps each method's natural
uncertainty semantics and only standardises the interface (column schema,
optional quantiles, and the optimisation-friendly scalar).

Shared interface

All estimators implement:

predict_uncertainty(
    X,
    return_quantiles=False,
    quantiles=DEFAULT_QUANTILES,   # [0.25, 0.5, 0.75]
    uncertainty_for_opt=False,
    num_samples=...,
    use_std=True,                  # regression MC-dropout heads only
) -> pd.DataFrame | pd.Series | tuple[pd.DataFrame, np.ndarray]

The default return is a DataFrame with a fixed 5-column schema:

Column	Meaning
pred	Point prediction (class for classifiers, target value for regressors).
mean_predictions	Mean prediction (MC-dropout mean, flow mode, or class probability).
knowledge_uncertainty	Epistemic (model / knowledge) uncertainty, or None if unavailable.
data_uncertainty	Aleatoric (data / noise) uncertainty, or None if unavailable.
total_uncertainty	The primary uncertainty measure used downstream.

Interface conventions:

• return_quantiles=True → returns (DataFrame, quantile_array).
  Only flow heads support this (see Quantiles); every other
  head raises ValueError.
• uncertainty_for_opt=True → returns a pd.Series (regression) or
  single-column DataFrame (classification) of total_uncertainty /
  knowledge_uncertainty, for active learning and acquisition ranking.

!!! note "Different uncertainty types by design"
total_uncertainty is not in the same units across all backends. For
MC-dropout heads it is a spread of predictions (target units); for flow
heads it is a negative log-likelihood (nats). Compare values within a
model, not across different model families.

Uncertainty estimation methods

Two mechanisms are used under the hood.

Monte-Carlo (MC) dropout

Multiple stochastic forward passes are run with dropout layers kept active at
inference time. The variability across passes estimates epistemic
(knowledge) uncertainty.

• Regression: knowledge_uncertainty = std (or IQR if use_std=False) of
  the predictions across passes; mean_predictions is the mean across passes.
• Classification: probabilities are averaged across passes and the
  entropy of the mean probability vector is the uncertainty.
• Requires at least one active dropout (input_dropout, tree_dropout, or
  mlp_dropout > 0). With all dropout disabled there is no stochasticity, so
  MC dropout cannot produce meaningful uncertainty.

Normalizing-flow likelihood

Flow heads (using zuko flows: GMM, NICE,
RealNVP, NAF, UNAF, NSF, BPF) model the full conditional distribution
$p(y \mid x)$. This gives:

• A mode prediction (highest-density sample).
• Data uncertainty = $-\log p(y_{\text{mode}} \mid x)$ — the negative
  log-likelihood (NLL) of the mode. High NLL indicates a wide and/or
  out-of-distribution prediction (aleatoric + OOD signal).
• Genuine predictive quantiles by sampling $p(y \mid x)$ directly.

NODE

NODERegressor supports four head types: subset, linear, mlp, and
flow. NODEClassifier supports subset, linear, and mlp (no flow).
The uncertainty behaviour depends on the head and whether dropout is active.

NODE regression

Head	Dropout	mean_predictions	knowledge_uncertainty	data_uncertainty	total_uncertainty
flow	on	flow mode	std of NLL across MC-dropout passes	NLL of mode (no-dropout pass)	data + knowledge
flow	off	flow mode	None	NLL of mode	= data
subset / linear / mlp	on	MC-dropout mean	std/IQR across passes	None	= knowledge
subset / linear / mlp	off	deterministic prediction	(falls back to zeros — dropout required)	None	—

Notes:

• For the flow + dropout case, epistemic uncertainty is the standard
  deviation of the mode's negative log-likelihood across MC-dropout passes.
  The mode itself is fixed from a single no-dropout pass; each dropout pass then
  re-evaluates log_prob of that fixed mode. This is computed in
  predict_with_combined_uncertainty.
• The data-uncertainty term always comes from the no-dropout pass, so it
  reflects the flow's aleatoric estimate rather than dropout noise.

Classification (NODE and MLP)

NODEClassifier and MLPHeadClassifier share the same MC-dropout entropy
decomposition. Running num_samples stochastic forward passes yields per-pass
class probabilities $p_t$ with mean $\bar p = \tfrac{1}{T}\sum_t p_t$:

• pred = argmax of the mean probability $\bar p$.
• mean_predictions = mean-over-dropout probability of the reported class
  (class-1 probability for binary, predicted-class probability for multiclass).
• total_uncertainty = entropy of the mean probability $H(\bar p)$.
• data_uncertainty (aleatoric) = mean per-pass entropy
  $\tfrac{1}{T}\sum_t H(p_t)$ (expected entropy).
• knowledge_uncertainty (epistemic) = mutual information
  $H(\bar p) - \tfrac{1}{T}\sum_t H(p_t)$ =
  total_uncertainty - data_uncertainty, matching CatBoost's decomposition
  (0 when dropout is inactive).

Standalone MLP / Flow heads

These are Skorch-wrapped estimators usable independently of the NODE backbone.

Estimator	Mechanism	knowledge_uncertainty	data_uncertainty	total_uncertainty
MLPHeadRegressor	MC dropout	std/IQR across passes	None	= knowledge
MLPHeadClassifier	MC dropout	mutual information (total - data)	mean per-pass entropy	entropy of mean probs
FlowHeadRegressor	flow likelihood	None (no dropout in head)	NLL of mode	= data

• MLPHeadRegressor reports a deterministic point prediction in pred and the
  MC-dropout mean in mean_predictions.
• MLPHeadClassifier follows the classification decomposition
  shared with NODEClassifier.
• FlowHeadRegressor has no dropout, so it exposes aleatoric uncertainty only,
  plus predictive quantiles from the flow.

Quantiles

Predictive quantiles are only meaningful when the model represents a genuine
predictive distribution. Therefore quantiles are flow-only:

Model	return_quantiles=True	predict_quantiles()
NODE flow regressor	✅ sampled from $p(y \mid x)$	✅
NODE non-flow regressor	❌ ValueError	❌ ValueError
NODE classifier	❌ ValueError	❌ ValueError
FlowHeadRegressor	✅ sampled from flow	—
MLPHeadRegressor / MLPHeadClassifier	❌ ValueError	—

Key points:

• Quantiles are computed by sampling the flow distribution in eval mode
  (dropout off) and taking torch.quantile of the samples. MC dropout is
  never used to build the quantile array — even in the flow + dropout case,
  dropout only feeds the scalar uncertainty columns.
• The returned array has shape (n_samples, n_quantiles) for single-target
  regression (an extra trailing axis for multi-target), matching the TabPFN and
  RandomForest quantile conventions.
• As with TabPFN / RandomForest, the requested quantiles are merged with
  DEFAULT_QUANTILES and sorted ascending, so column i corresponds to sorted
  quantile i. predict_quantiles additionally filters the result back to the
  exact quantiles requested.

Examples

from mother.ml.models.m_node import NODERegressor

# Flow head with dropout: both aleatoric (NLL) and epistemic (std of NLL).
reg = NODERegressor(head_type="flow", flow_type="NICE", input_dropout=0.1)
reg.fit(X_train, y_train)

df = reg.predict_uncertainty(X_test)
# columns: pred, mean_predictions, knowledge_uncertainty, data_uncertainty, total_uncertainty

# Predictive quantiles sampled from the flow (dropout off):
df, q = reg.predict_uncertainty(
    X_test, return_quantiles=True, quantiles=[0.025, 0.5, 0.975]
)
lower, median, upper = q[:, 0], q[:, 1], q[:, 2]

# Acquisition-friendly scalar for active learning:
total = reg.predict_uncertainty(X_test, uncertainty_for_opt=True)  # pd.Series

from mother.ml.models.m_node import NODERegressor

# Non-flow head: MC-dropout epistemic uncertainty only, no quantiles.
reg = NODERegressor(head_type="mlp", input_dropout=0.1)
reg.fit(X_train, y_train)

df = reg.predict_uncertainty(X_test)          # 5-column schema, data_uncertainty is None
reg.predict_uncertainty(X_test, return_quantiles=True)   # raises ValueError

[thomasATbayer]

Classification Uncertainty: Modelled After CatBoost

This page explains how the classification uncertainty for the neural tabular
classifiers — NODEClassifier and MLPHeadClassifier — is derived to match
CatBoost's information-theoretic decomposition.

1. What CatBoost actually does

CatBoost estimates classification uncertainty with virtual ensembles: a
single trained model is split into $T$ independent sub-models (virtual ensemble
members). Predicting with prediction_type="TotalUncertainty" returns, per
sample, two entropy-based quantities. In the Mother wrapper
(get_virtual_prediction) these are read out
as:

data_uncertainty      = virtual_prediction[:, 0]
total_uncertainty     = virtual_prediction[:, 1]
knowledge_uncertainty = total_uncertainty - data_uncertainty

This is the standard Malinin et al. information-theoretic decomposition of
predictive uncertainty for an ensemble of probabilistic classifiers. Given
per-member class-probability vectors $p_t$ (one per ensemble member
$t = 1 \dots T$) with mean $\bar p = \tfrac{1}{T}\sum_t p_t$:

$$ \underbrace{H(\bar p)}_{\text{total}} ;=; \underbrace{\tfrac{1}{T}\sum_t H(p_t)}_{\text{data (aleatoric)}} ;+; \underbrace{\Big[H(\bar p) - \tfrac{1}{T}\sum_t H(p_t)\Big]}_{\text{knowledge (epistemic)}} $$

where $H(p) = -\sum_c p_c \log p_c$ is the Shannon entropy.

Term	CatBoost name	Formula	Meaning
Total	total uncertainty	$H(\bar p)$ — entropy of the mean prediction	Overall predictive uncertainty (predictive entropy)
Data	data uncertainty	$\tfrac{1}{T}\sum_t H(p_t)$ — mean of the per-member entropies	Irreducible aleatoric noise / class overlap
Knowledge	knowledge uncertainty	$H(\bar p) - \tfrac{1}{T}\sum_t H(p_t)$	Epistemic disagreement between members = mutual information

The key subtlety: data uncertainty is the (mean) entropy, not the entropy
of the mean. The entropy of the mean is the total. Their difference (which
is always $\ge 0$ by Jensen's inequality, since entropy is concave) is the
mutual information — the epistemic term.

2. Mapping CatBoost's virtual ensemble onto NODE / MLP

Neural heads don't have virtual ensembles, but MC dropout gives us the same
structure: each stochastic forward pass with dropout active is treated as one
ensemble member. Running num_samples passes yields a stack of per-pass
probability vectors $p_t$, exactly analogous to CatBoost's $T$ virtual members.

CatBoost	NODE / MLP equivalent
virtual ensemble member $t$	one MC-dropout forward pass
$T$ = virtual_ensembles_count	$T$ = num_samples
per-member probs $p_t$	per-pass softmax probs $p_t$
$\bar p$ = ensemble mean	mean over dropout passes

3. The implementation

Both NODEClassifier.predict_uncertainty
(m_node.py) and
MLPHeadClassifier.predict_uncertainty
(m_heads.py) use identical logic:

from scipy.stats import entropy

# probabilities: (num_samples, n_datapoints, n_classes) — one slice per MC-dropout pass
probabilities = self._mc_dropout_proba_samples(X, num_samples)

# Predictive (mean) distribution across MC passes.
mean_probs = probabilities.mean(axis=0)          # (n_datapoints, n_classes)
pred = mean_probs.argmax(axis=1)

# Uncertainty decomposition matching CatBoost (Malinin et al.):
#   total     = entropy of the mean predictive distribution  H(mean_p)
#   data      = mean per-pass entropy (expected entropy)      E_t[H(p_t)]
#   knowledge = total - data (mutual information; 0 when dropout inactive)
total_uncertainty = entropy(mean_probs, axis=1)              # H(mean_p)
per_pass_entropy  = entropy(probabilities, axis=2)           # H(p_t) per pass -> (T, n)
data_uncertainty  = per_pass_entropy.mean(axis=0)            # E_t[H(p_t)]
knowledge_uncertainty = total_uncertainty - data_uncertainty  # mutual information

scipy.stats.entropy normalises each probability vector internally and uses
natural-log units (nats), consistent with CatBoost's entropy-based outputs.

Column semantics

• pred = argmax of $\bar p$.
• mean_predictions = mean-over-dropout probability of the reported class
  (class-1 probability for binary, max class probability for multiclass).
• total_uncertainty = $H(\bar p)$.
• data_uncertainty = $\tfrac{1}{T}\sum_t H(p_t)$.
• knowledge_uncertainty = total − data (mutual information).

Interface conventions preserved

• uncertainty_for_opt=True → single-column DataFrame of
  knowledge_uncertainty (the epistemic term, which is what active-learning
  acquisition wants — it drives sampling toward regions of model disagreement
  rather than irreducible noise).
• return_quantiles=True → ValueError (quantiles are flow-only; a classifier
  has no continuous predictive distribution to sample).

4. Why this differs from the previous version

The earlier implementation had two mismatches with CatBoost:

	Old (incorrect)	New (CatBoost-aligned)
data_uncertainty	$H(\bar p)$ — entropy of the mean	$\tfrac1T\sum_t H(p_t)$ — mean entropy
knowledge_uncertainty	$\operatorname{std}_t H(p_t)$ — std of per-pass entropies (ad-hoc)	$H(\bar p) - \tfrac1T\sum_t H(p_t)$ — mutual information
total_uncertainty	data + knowledge (of the wrong terms)	$H(\bar p)$ (predictive entropy)

The old data term was actually the total (entropy of the mean), and
knowledge was a heuristic spread with no information-theoretic guarantee (it
could be arbitrarily scaled and wasn't the epistemic complement of the aleatoric
term). The new version restores the exact additive decomposition
$\text{total} = \text{data} + \text{knowledge}$ that CatBoost uses.

5. Guarantees and numerical notes

• Additivity holds exactly: total == data + knowledge.
• Mutual information is non-negative by Jensen's inequality ($H$ concave
  $\Rightarrow H(\bar p) \ge \tfrac1T\sum_t H(p_t)$). In practice, NODE's float32
  probabilities can produce a tiny negative knowledge value
  ($\sim -1.8\times10^{-7}$) when passes are near-identical (confident /
  low-dropout regime); this is pure floating-point noise, not a modelling error.
  The MLP (float64 path) stays $\ge 0$.
• Requires active dropout. With all dropout disabled every pass is
  identical, so $p_t = \bar p$, giving knowledge = 0 and total = data. This
  mirrors CatBoost degenerating to a single deterministic member.
• Units. Entropy is in nats, comparable within the classifier family.
  Don't compare these magnitudes against the flow-head NLL numbers.

References

• A. Malinin, L. Prokhorenkova, A. Ustimenko, Uncertainty in Gradient
  Boosting via Ensembles, ICLR 2021.
• CatBoost uncertainty documentation:
  https://catboost.ai/docs/en/references/uncertainty.
• See also the companion page
  Uncertainty for NODE and MLP models.