dnamite: A Python Package for Neural Additive Models

TL;DR

dnamite addresses the need for interpretable yet flexible additive modeling on tabular data by implementing Neural Additive Models with a DNAMite architecture. It enables regression, binary classification, and survival analysis, integrating feature selection through learnable gates and offering a scikit-learn–style interface built on PyTorch. The package supports advanced capabilities such as survival IPCW loss, monotonic constraints, and calibration plots, and demonstrates practical gains on the MIMIC-III clinical dataset. By combining interpretability, automatic feature selection, and survival analysis in a single Python toolkit, dnamite has potential to streamline deployment of transparent NAMs in healthcare and other domains with complex tabular data.

Abstract

Additive models offer accurate and interpretable predictions for tabular data, a critical tool for statistical modeling. Recent advances in Neural Additive Models (NAMs) allow these models to handle complex machine learning tasks, including feature selection and survival analysis, on large-scale data. This paper introduces dnamite, a Python package that implements NAMs for these advanced applications. dnamite provides a scikit-learn style interface to train regression, classification, and survival analysis NAMs, with built-in support for feature selection. We describe the methodology underlying dnamite, its design principles, and its implementation. Through an application to the MIMIC III clinical dataset, we demonstrate the utility of dnamite in a real-world setting where feature selection and survival analysis are both important. The package is publicly available via pip and documented at dnamite.readthedocs.io.

Figures (4)

• Figure 1: The DNAMite architecture. For simplicity, only individual features (and not interactions) are shown. The choice of link and loss function depends on the prediction task. Feature gates are only used for feature selection. The embedding module outputs a weighted sum of neighborhood embeddings for each feature.
• Figure 2: Visualization of the smooth-step function (Equation \ref{['eq:smooth_step']}) for different $\gamma$'s.
• Figure 3: Class structure for dnamite models.
• Figure 4: Comparison of dnamite to other additive model packages. For each package, the shape function for age is given as well as the test AUC.