Neural Additive Models: Interpretable Machine Learning with Neural Nets

In plain words

Doctors, lenders and judges often need to know why a model gave a particular answer. The most accurate tabular models are usually black boxes, and the most readable ones (short trees, straight-line scores) are usually less accurate. This paper tries to get both.

The idea is to build a model out of separate parts, one per input column. Each part is a small neural network that only sees its own column and outputs a single number. The model adds these numbers together, adds one overall bias, and passes the total through a link function such as the logistic one. Because each part sees only one input, you can plot each part as a curve against its input and read off exactly how that input moves the prediction. The parts are trained together with ordinary gradient descent, so the whole thing stays differentiable and can be combined with other neural-network pieces.

The paper also changes how the hidden units inside each part are built, so that a part can bend sharply where the data jumps, which plain neural networks struggle to do. The authors report that these models match a strong existing family of readable additive models on several tabular datasets, and they show two uses that tree-based versions cannot easily do: learning several related tasks at once and plugging the readable part into a larger model.

TL;DR

• NAMs are generalized additive models whose shape functions are small neural networks, one per input feature, trained jointly with backpropagation §2 eq. 1.
• Standard ReLU networks fail to fit sharp jumps in one-dimensional functions under mini-batch training; the ExU hidden unit (an exponentially scaled, shifted input clipped to [0, n]) fixes this §2 fig. 2.
• Author-reported 5-fold results on MIMIC-II, Credit, California Housing and FICO are close to EBMs and below XGBoost and DNNs tab. 1.
• The released code is Apache-2.0 TensorFlow 1.x; several paper hyperparameters differ from the code defaults (see Paper vs code).

Problem

Generalized additive models (GAMs) keep each feature's effect separate and therefore readable §1 eq. 1. The strongest GAMs at the time were tree-based (boosted trees, EBMs), and the authors argue these are hard to extend to multitask or multiclass settings and cannot be trained with GPUs or differentiated through §1. Linear models and shallow trees are readable but less accurate §3 tab. 1. Plain neural networks are accurate but opaque §1.

Two failure modes are stated as assumptions. First, real tabular data often has sharp jumps in a feature's effect, which smooth splines miss §2. Second, a ReLU network trained with mini-batch gradient descent does not fit a jagged one-dimensional target even when it has enough capacity §2 fig. 2. The authors note that full-batch training does not show this problem §2, footnote.

Method

Model. The prediction is a link function applied to a bias plus the sum of one shape function per input column §2 eq. 1. Each shape function is a neural net. For a model with K input columns:

• f_k is a feature net that maps column k alone to a scalar §2.
• The output is the sum of all f_k(x_k) plus a single bias term. The bias and the feature nets are trained jointly §2 models.py.
• The shape functions are not forced to zero mean in the training code. The paper says each graph is centred, with the bias absorbing the average prediction §2; see Paper vs code.

Feature net. Two variants are used §3 app. A.5 models.py:

1. Shallow: one hidden layer of n units, then a linear read-out with no bias.
2. Deep: one hidden layer of n units, then ReLU layers of 64 and 32 units, then a linear read-out.

The hidden layer is either standard ReLU units or ExU units. Dropout is applied after each hidden layer of a feature net models.py. The paper describes the deep variant as "3 hidden layers with 64, 64 and 32 units" §3; the code's first layer width is the one set by n (see Paper vs code).

ExU hidden unit. For a scalar input x, a unit with weight w, bias b and activation f computes §2 eq. 2:

• h(x) = f(exp(w) * (x - b)), with f = ReLU clipped at 1 (ReLU-1) in the experiments §2 app. A.4.
• The log-space weight means the slope exp(w) is always positive and can be very steep with a small w §2.
• ExU is used only for the datasets where the table says so; other datasets use standard ReLU units app. A.6.

Training objective app. A.3 eq. 3:

• Loss = mean task loss + λ1 × output penalty + λ2 × weight decay.
• Task loss: cross-entropy on the logit for binary classification, mean squared error for regression app. A.3.
• Output penalty: mean squared output of each feature net, averaged over feature nets, which keeps each net's contribution near zero unless the data supports it app. A.3.
• Weight decay: L2 penalty on the weights of the feature nets, averaged over feature nets app. A.3.
• Feature dropout: during training, whole feature nets are dropped for individual examples app. A.3.

Optimization app. A.5:

• Adam, batch size 1024, at most 1000 epochs, early stopping on the validation split.
• Learning rate multiplied by 0.995 after every epoch.
• Learning rate, λ1, λ2, dropout and feature dropout tuned by Bayesian optimization on each cross-validation fold.

# NAM forward pass  [§2] [eq. 1]
# Input:  x (B, K), one row per example
# Output: logit (B,)

def nam_forward(x, feature_nets, bias, training):
    # Each feature net sees only its own column  [§2]
    per_feature = []
    for k in range(K):
        per_feature.append(feature_net(feature_nets[k], x[:, k], training))
    # Stack the K scalar outputs per example  [code: models.py#L242]
    stacked = stack(per_feature, axis=-1)            # (B, K)
    # Feature dropout drops whole nets per example  [app. A.3]
    stacked = feature_dropout(stacked, training)     # (B, K)
    # Sum the feature outputs and add one shared bias  [eq. 1]
    return sum(stacked, axis=-1) + bias              # (B,)

# One feature net  [§2] [eq. 2] [app. A.3]
# Input:  value (B,) a single column
# Output: scalar contribution (B,)

def feature_net(net, value, training):
    # ExU unit: h = min(max(exp(w) * (value - b), 0), 1)  [eq. 2]
    hidden = exu_units(value, net.weight, net.center)  # (B, n)
    hidden = dropout(hidden, training)                  # (B, n)
    # Deep variant only: extra ReLU layers of 64 and 32  [app. A.5]
    hidden = relu_layer(hidden, 64)                     # (B, 64)
    hidden = relu_layer(hidden, 32)                     # (B, 32)
    # Linear read-out with no bias term  [code: models.py#L160-L165]
    return linear(hidden)                               # (B,)

# Penalised loss  [app. A.3] [eq. 3]
# Input:  batch_x (B, K), batch_y (B,)
# Output: scalar loss

def penalised_loss(model, batch_x, batch_y):
    # Task loss: cross-entropy or MSE on the summed output  [app. A.3]
    task = task_loss(model(batch_x, training=True), batch_y)
    # Output penalty: mean of squared feature-net outputs  [eq. 3]
    outputs = model.feature_outputs(batch_x, training=False)
    output_penalty = mean(square(outputs))
    # Weight decay: L2 over the feature-net weights  [eq. 3] [app. A.3]
    decay = l2_norm_squared(model.weights) / K
    return task + lambda_1 * output_penalty + lambda_2 * decay

# Training loop  [app. A.5]
# Input:  train (x, y), validation (x, y); hyperparameters from A.6
# Output: model with the best validation score

def train(model, train_xy, val_xy):
    # Adam with batch size 1024, at most 1000 epochs  [app. A.5]
    optimizer = adam(learning_rate)
    best_score = none
    for epoch in range(max_epochs):
        for batch in minibatches(train_xy, batch_size=1024):
            # One gradient step on the penalised loss  [app. A.3]
            optimizer.step(penalised_loss(model, *batch))
        # Learning rate is multiplied by 0.995 after every epoch  [app. A.5]
        learning_rate = learning_rate * 0.995
        # Early stopping on validation score  [app. A.5]
        score = validation_score(model, val_xy)
        if improved(score, best_score):
            best_score = score
            keep_checkpoint(model)
    return best_checkpoint

Interpretation. Each shape function is plotted against its input with the data density shown alongside. An ensemble of models gives several curves per feature, and their spread is read as confidence §2.

Multitask NAM §4.2 app. A.8. Each feature gets several subnets. For each task, a learned weighted sum over those subnets gives that task's shape function. The task outputs are summed and a bias is added §4.2. The synthetic experiment uses 6 subnets per feature for 6 tasks §4.2.1.

Evaluation protocol §3 app. A.5:

• Stratified (classification) or plain (regression) 5-fold cross-validation.
• In each fold, 20 random train/validation splits (train 70%, validation 10%, test 20% of the whole dataset). The NAM ensemble for a fold is the 20 trained models app. A.5.
• Classification metric in the text: area under the precision–recall curve; regression metric: RMSE §3 app. A.5.

Hyperparameters

Values are as stated in the paper. Where the paper is silent the cell says "not stated".

Tuned values for the NAM feature nets app. A.6 (tables in the appendix):

The table's "units" column is written as three layer widths for the deep nets and one width for the shallow ExU nets app. A.6.

Fixed settings not in the table: the DNN baseline and EBM settings follow prior work app. A.6. EBMs use the open-source InterpretML implementation with the parameters from the earlier EBM paper app. A.6. Logistic and linear regression and CART use sklearn with grid search §3 app. A.6. XGBoost uses the XGBoost implementation §3.

Results

Author-reported. Setup: five public tabular datasets, 5-fold cross-validation, 20 train/validation splits per fold, single-task NAM unless noted §3 app. A.5. Baselines: logistic or linear regression, CART, EBMs, a 10-layer DNN, and XGBoost §3 app. A.6.

Single-task results, mean ± standard deviation over folds tab. 1:

Read with care: the paper calls NAMs "comparable to EBMs" §3. On MIMIC-II and Credit the NAM and EBM intervals overlap; on CA Housing and FICO the NAM is slightly behind EBMs and behind XGBoost and the DNN on the regression tasks tab. 1.

Multitask synthetic data §4.2.1 tab. 2: mean MSE over 6 tasks, average of 20 runs.

The paper states multitask MSE is 34% lower on average and at least 25% lower on each task, and that multitask beat single-task in 119 of 120 trials §4.2.1.

COMPAS recidivism, ROC AUC by gender §4.2.2 tab. 3:

Other results are qualitative: a Covid-19 treatment-benefit model trained on over 3000 patients §4.1, and MIMIC-II shape functions reviewed with clinicians app. A.1. No numbers are given for these two.

Paper vs code

Code: google-research/google-research, subdirectory neural_additive_models/, commit 8bedbbd982e7dcf5c5280e6ae2f0769295c3369e, Apache-2.0 (repo root LICENSE). The paper links this repo from its project page. The README links a separate multitask repo at lemeln/nam [README]; that repo was not cloned or checked. The Covid model and the synthetic data generator are not in this repo (grep found neither).

1. Classification metric. Paper: area under the precision–recall curve §3 app. A.5. Code: ROC AUC via roc_auc_score graph_builder.py graph_builder.py. Table 3 also says ROC AUC tab. 3. Use ROC AUC when reproducing the classification numbers.
2. Hidden units per feature. Paper: one hidden layer of 1024 ExU units §3 app. A.6. Code: units per feature = min(1000, 2 × number of unique values in that column) graph_builder.py nam_train.py. The cap is 1000 and the width depends on the data. Setting the 1024 width requires changing the flags.
3. Deep feature net. Paper: three hidden layers of 64, 64 and 32 units §3 app. A.5. Code: a first layer of the width above, then layers of 64 and 32, then the linear read-out models.py. The paper's first layer is 64; the code's is the width from item 2.
4. Train/validation splits. Paper: 20 random splits per fold app. A.5. Code default: 3 nam_train.py. n_models defaults to 1 nam_train.py, and the code does not average predictions across models. Pass 20 for both, and build the ensemble yourself.
5. Epoch budget and early stopping. Paper: at most 1000 epochs, early stopping on validation app. A.5. Code: training_epochs has no default and must be given nam_train.py. Validation is checked only every 10 epochs nam_train.py nam_train.py. Training stops after 60 epochs without improvement nam_train.py nam_train.py. The paper does not state a patience.
6. Class balance in training. Not in the paper. For classification, batches are drawn from positive and negative examples with equal probability, so the training prior is 50/50 graph_builder.py graph_builder.py. Test data keeps its natural prior.
7. Feature scaling. Not in the paper. Code uses min–max scaling to [-1, 1] fit on the whole dataset before splitting into folds data_utils.py. The fold split happens afterwards nam_train.py. Test rows therefore affect the scaler. Fit the scaler on the training part of each fold when you reimplement.
8. Categorical features. Not in the paper. Each one-hot column becomes its own feature net, so K is the number of encoded columns, not original fields data_utils.py graph_builder.py.
9. Weight decay scope. Paper: penalty on the feature-net weights app. A.3. Code: tf.nn.l2_loss over all trainable variables, including the ExU centers, the ExU weights and the output bias, divided by the number of feature nets graph_builder.py. l2_loss is half the sum of squares.
10. Output penalty. Computed on feature-net outputs with dropout switched off graph_builder.py graph_builder.py. It is a mean over examples, then over nets. The paper's equation writes a sum over examples eq. 3.
11. Feature dropout scaling. The code uses TensorFlow dropout, which rescales surviving outputs by 1/(1 − p) during training models.py. The paper does not state this.
12. ExU initialization. Paper: weights ~ N(x, 0.5) with x in [3, 4], centers not stated §2. Code: weight ~ truncated normal, mean fixed at 4.0, s.d. 0.5; center ~ truncated normal, mean 0, s.d. 0.5 models.py models.py.
13. Learning-rate decay. Applied once per epoch after the epoch's batches nam_train.py. Matches the paper app. A.5.
14. Shape-function centring. Paper: each shape function is set to zero mean over the training data and the bias absorbs the average §2. Code: no centring step in the training path (models.py, graph_builder.py, nam_train.py read in full). The output penalty only pushes outputs towards zero. Centring may happen in the visualization notebook linked from the README; that was not checked.
15. Matching items. The 10-layer DNN with 100 ReLU units per layer app. A.6 matches models.py. The 0.995 decay and batch size 1024 match app. A.5 nam_train.py nam_train.py.

Pitfalls

• Unstated ExU details. The paper gives the initialization only loosely ("x in [3, 4]") §2. The code uses a fixed mean of 4.0 for the weight models.py. Since the unit's slope is exp(weight), that gives an initial slope of about exp(4) ≈ 55. This value is derived from the code, not stated in the paper.
• ReLU-1 saturation. Each ExU unit's output is clipped to [0, 1] models.py. Each unit's output is bounded before the read-out, so a steep change is built from several units with different centers, plus a large read-out weight. The paper's justification is that each unit is active only in a small input range §2.
• Stability claim. The paper reports no training instabilities with ExU across its datasets and attributes this to Adam and regularization app. A.4. This is the authors' observation, not a guarantee.
• Regularization interacts with the feature-net count. The output penalty and weight decay are averaged over nets graph_builder.py graph_builder.py, so changing the number of features or one-hot columns changes the effective penalty strength.
• Cost scales with features and ensembles. Each column has its own network of up to 1000 units (code) or 1024 units (paper). Each fold trains 20 models in the paper protocol app. A.5. Budget for K × units parameters per model and 20 models per fold.
• Interpretation limits. Shape plots are exact descriptions of what the model computes, but they are not causal explanations app. A.1. Data density matters: the authors caution that gaps in data give high variance in the curves §2.
• Fairness claim. On COMPAS the learned race graph is read as a sign of bias in the data. The paper suggests zeroing that graph may make the model fairer, at a cost to AUC §4.2.2. This is the authors' suggestion, not a validated fairness method.
• Small MTL differences. On COMPAS, the multitask and single-task intervals overlap tab. 3. The clearer multitask gain is on the synthetic data, which was built to favour multitask learning §4.2.1.
• Metric wording. The text says precision–recall AUC §3; the tables and the code use ROC AUC tab. 1 graph_builder.py. Confirm the metric before comparing numbers.
• Missing code. The Covid-19 model and the multitask synthetic generator are not in the official repo §4.1 §4.2.1. The paper says the generator code will be released app. A.8.1; the README points to a separate multitask repo that was not checked.

Lineage

• Builds on: generalized additive models with spline or tree shape functions §1 §5 (Caruana et al. 2015 and Lou et al. 2012/2013 are cited for boosted-tree GAMs §5). Not linked to a wiki card.
• Compared against: logistic and linear regression, CART, EBMs (boosted-tree GAMs), XGBoost and a 10-layer DNN §3.
• Related but distinct: GANNs, which used one hidden unit per feature and human-in-the-loop architecture changes §5 app. A.7.
• Superseded by: not stated in this card's source.