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="ple" | "pbld" | "plr" | "plr-lite" | "pl" | "periodic" — PLE-vB ("ple", the TabM† embedding) plus the Fourier/periodic embedding zoo (arXiv:2203.05556 + PBLD); token models (ft_transformer, tab_transformer) use the PLR family as feature tokenizers, but not PLE
• num_embedding_cols=[...] / linear_skip_cols=[...] (new in 0.9.1) — input routing: embed only these numeric columns and let the rest enter the first layer linearly, and/or send these numeric columns straight to the output through a zero-initialized linear skip (raw = trunk(x) + x_skip @ W_skip + b_skip, own param group at linear_skip_lr_factor, no weight decay; realmlp). For numeric blocks that mix raw measurements with target-encoded log-odds columns
• 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
• inner ensembling — tabm (always), realm (always, new in 0.9.2) and ft_transformer (k>1, new in 0.8.0) run k weight-shared members inside one model, composing with the outer n_ens; the n_ens·k members are exposed by predict_members / predict_proba_members. tabm and ft_transformer use the TabM-mini adapter; tabm additionally offers variant="full", while realm uses full BatchEnsemble over the RealMLP trunk
• 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)
• masamlp.realm_td_params(task, k=8) — the same recipe on the realm architecture: identical training knobs, k BatchEnsemble members instead of one model

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 0.5.0 models verified on Kaggle v5e-8 and Colab v5e-1 (tabm, new in 0.6.0, has XLA CI coverage but no TPU measurement yet). 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.