masaMLP

Extensible tabular deep learning — TabularResNet, DANet, and TabularLNN behind sklearn-compatible estimators with first-class sample_weight, custom objectives, custom metrics, and early stopping on any metric. The sibling library of repleafgbm (same author, same API philosophy), for the neural side of tabular ML.

Status: alpha (0.2.x). Built with heavy use of Claude Code (coding and architecture design).

Why masaMLP

Excellent tabular DL libraries exist — pytabkit ships state-of-the-art models like RealMLP and TabM, and rtdl provides reference modules. What they don't make easy is extension: sample_weight in fit, custom training losses, custom evaluation metrics, and early stopping driven by them. masaMLP is built around exactly those hooks:

• fit(X, y, sample_weight=..., eval_set=...) — LightGBM-style, sklearn compatible. Weights flow through a single reduction (loss * w).sum() / w.sum() that every objective shares.
• Custom objectives are per-sample torch losses — a plain function (or nn.Module with trainable parameters). Because the trainer owns the weighted reduction, your loss gets correct sample_weight and class_weight handling for free.
• Custom metrics are plain NumPy callables via make_metric, and any of them (minimize or maximize) can drive early stopping with best-epoch weight restoration.
• Multiclass, multioutput regression, class_weight, label smoothing supported natively; built-in preprocessing (quantile scaling, missing values, categorical embeddings) so DataFrames go straight into fit.
• CPU / CUDA / MPS / multi-GPU behind device="auto": device-resident tensors with no DataLoader overhead, automatic full-batch mode for small data, per-model bf16 AMP on CUDA, opt-in torch.compile with eager fallback — and when several GPUs are detected, n_ens members train concurrently, one worker per GPU.

masaMLP deliberately does not try to re-benchmark the field — see docs/attribution.md for the research and libraries it builds on.

Models

Third-party architectures plug in with register_model and get the whole estimator surface (weights, objectives, metrics, early stopping) for free.

Every architecture is fully configurable through model_params — depth, width, dropout, and the model-specific knobs (e.g. TabR's context_size). The complete list per model, with defaults and sizing notes, is in docs/parameters.md:

clf = MasaClassifier(model="ft_transformer",
                     model_params={"n_blocks": 4, "d_block": 256})

# realmlp is the free-form MLP: hidden_sizes IS the architecture
# (length = depth, entries = per-layer width)
reg = MasaRegressor(model="realmlp",
                    model_params={"hidden_sizes": (512, 256, 128)})

Layer-by-layer width control (512 → 256 → 128) is a realmlp feature — its hidden_sizes is the whole stack. The other architectures keep a constant width per their papers (residual streams require it); size those with d/d_token × n_blocks/n_layers.

Key parameters

Constructor parameters shared by both estimators (the full reference, including everything below plus preprocessing, ensembling, and hardware options, is docs/parameters.md):

RealMLP insights are composable options

The tricks from the RealMLP paper are estimator-level options usable with any model (lnn included), not baked into one architecture:

• numeric_scaler="rssc" — robust scale + smooth clip preprocessing
• cat_encoding="onehot" — RealMLP-style one-hot (binary → ±1, missing → 0)
• num_embedding="pbld" | "plr" | "plr-lite" | "pl" | "periodic" — the numeric embedding zoo (arXiv:2203.05556 + PBLD); token models (ft_transformer, tab_transformer) use the same options as feature tokenizers
• model_params={"num_scaling": True} — learnable per-feature input scale
• lr_scheduler="coslog4", optimizer_betas=(0.9, 0.95) — the training schedule
• clip_predictions=True (regressor) — clip to the observed target range
• n_ens=k — seed ensembling as in pytabkit's RealMLP: k members trained with seeds random_state + i, predictions averaged on the probability / value scale; works with every model including the retrieval ones. ens_mode="vectorized" trains all members in one vmapped forward/backward (torch.func) for BatchNorm-free models — pytabkit's speed trick
• weight_decay_schedule="flat_cos" — RealMLP-TD's scheduled weight decay (param groups can opt out, e.g. biases)
• ema_decay=0.999 — exponential moving average (Polyak averaging) of the weights; evaluation, early stopping, and the final model all use the averaged parameters
• candidate_budget=N — bound the retrieval corpus of tabr/modernnca with a seeded, class-stratified subsample (keeps memory/compute in check on large data; no-op for other models)
• masamlp.realmlp_td_params(task) — the full RealMLP-TD recipe: parametric activations, flat_cos-scheduled dropout and weight decay, PBLD embeddings with their own lr factor, and hybrid categorical encoding (one-hot ≤ 9 categories, embeddings above)

from masamlp import MasaClassifier, realmlp_params

clf = MasaClassifier(**realmlp_params("classification"))    # the TD-S recipe
clf = MasaClassifier(**{**realmlp_params("classification"),
                        "num_embedding": "pbld"})           # toward RealMLP-TD

Install

pip install masamlp        # torch, numpy, pandas, scikit-learn

Quickstart

import numpy as np
from masamlp import MasaClassifier, make_metric

def f1(y_true, y_proba):
    pred = y_proba >= 0.5
    tp = np.sum(pred & (y_true == 1))
    return 2 * tp / max(pred.sum() + (y_true == 1).sum(), 1)

clf = MasaClassifier(
    model="resnet",
    eval_metric=make_metric(f1, name="f1", minimize=False),
    early_stopping_rounds=15,
    class_weight="balanced",
)
clf.fit(X_train, y_train, sample_weight=w_train, eval_set=[(X_val, y_val)])
proba = clf.predict_proba(X_test)
print(clf.best_iteration_, clf.best_score_, clf.evals_result_["valid_0"]["f1"][:3])

Custom objective (regression, asymmetric loss):

import torch
from masamlp import MasaRegressor

def asymmetric_mse(y_true, raw_pred):          # -> per-sample (n,) tensor
    err = raw_pred - y_true                    # raw_pred: (n, out_dim)
    return torch.where(err < 0, 4.0 * err**2, err**2).mean(dim=1)

reg = MasaRegressor(model="danet", objective=asymmetric_mse)
reg.fit(X, y, sample_weight=w)                 # weights just work

Save/load is a plain directory (manifest.json + tensors, loaded with weights_only=True — no pickle execution):

reg.save_model("model_dir")
reg2 = MasaRegressor.load_model("model_dir")

Devices

device="auto" resolves tpu > cuda > mps > cpu. CUDA gets bf16 AMP by default (per-model policies apply) and optional compile=True; MPS and CPU train in float32. With multiple GPUs and n_ens > 1, ensemble members are sharded across all GPUs and trained concurrently; opt out with device="cuda:0".

TPU (experimental, 0.4.0/0.5.0): device="tpu" (or "xla") trains on Cloud/Kaggle/Colab TPUs via torch_xla — bf16 by default, all ten models, verified on Kaggle v5e-8 and Colab v5e-1. Kaggle grants TPU quota separately from GPU quota, so this is extra free accelerator time for competition workloads. 0.5.0 adds opt-in bf16 prediction (amp_predict) and large-corpus TabR eval fusion (predict −44% at 345k rows on TPU). Details, install pairing, tuning guidance, and measured numbers: docs/devices.md.

Development

pip install -e ".[dev]"
bash scripts/check.sh      # ruff + pytest + examples/quickstart.py

Development rules live in CLAUDE.md; roadmap in docs/roadmap.md.

License

MIT. Architecture attributions: docs/attribution.md.