+1 (415) 366-1364

Tabular ML on TensorFlow: Decision Forests vs. Keras Neural Networks

Most of the models we are asked to build for clients are not language models or vision models. They are tabular: churn, fraud, credit risk, propensity, lead scoring, demand forecasting on a wide feature table. Every one of our other tutorials - Gemma fine-tuning, two-tower retrieval, LiteRT - is about deep networks. This one is about the boring, profitable case, and about the honest answer to a question clients ask constantly: do we even need a neural network here?

Usually the answer is no. Gradient-boosted trees still win on small and mid-size tabular data, they train in minutes on a CPU, they need almost no tuning, and they are far easier to explain to a risk committee. The good news is you do not have to leave the TensorFlow stack to use them.

The two tools

TensorFlow Decision Forests (TF-DF) wraps the Yggdrasil Decision Forests (YDF) library in a Keras-style API: Random Forest, Gradient Boosted Trees, CART. It trains on CPU, consumes a tf.data.Dataset or a pandas DataFrame, and exports a SavedModel that TensorFlow Serving loads like any other model.

Keras 3 preprocessing layers plus a dense network is the path when you need embeddings for high-cardinality IDs, when the tabular features feed a larger multimodal model, or when you want one differentiable graph end to end.

A note on versions before you install anything: TF-DF is packaged separately from TensorFlow and tracks specific TF releases, and the YDF team now also ships a standalone ydf package with the same algorithms and a faster API. Pin both tensorflow and tensorflow_decision_forests in the same lockfile and check the compatibility table in the TF-DF documentation before you upgrade - this is the single most common way these projects break, and the same discipline we describe in our TensorFlow 2.21 upgrade guide.

Baseline first: gradient boosted trees in ten lines

pip install tensorflow tensorflow_decision_forests
import pandas as pd
import tensorflow_decision_forests as tfdf

train_df = pd.read_csv("train.csv")
val_df = pd.read_csv("val.csv")

train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(train_df, label="churned")
val_ds = tfdf.keras.pd_dataframe_to_tf_dataset(val_df, label="churned")

model = tfdf.keras.GradientBoostedTreesModel(
    num_trees=500,
    early_stopping="LOSS_INCREASE",
    validation_ratio=0.1,
)
model.fit(train_ds)

Things that are deliberately missing from that snippet, and that trip up people arriving from the deep-learning side:

  • No normalization, no scaling. Trees split on raw values. Standardizing your features does nothing except cost you time.
  • No one-hot encoding. TF-DF handles categorical strings natively, including high-cardinality ones. One-hot encoding them usually makes the model worse.
  • No imputation. Missing values are a first-class branch condition. Filling them with the mean destroys information the tree could have used.
  • No learning rate schedule, no epochs. fit is called once. Boosting is not SGD.

Evaluate and inspect:

model.compile(metrics=["AUC"])
print(model.evaluate(val_ds, return_dict=True))
print(model.summary())          # includes variable importances

model.summary() gives you out-of-bag or validation loss per tree, the number of trees actually kept after early stopping, and variable importances (NUM_AS_ROOT, SUM_SCORE, and others). That output alone answers most of the "which features matter?" questions a stakeholder will ask in the first review meeting, and it is why we always ship a tree baseline even on projects that end up neural.

Tuning, briefly

The default hyperparameters are strong. When you do tune, TF-DF has a built-in tuner so you do not need an external sweep framework:

tuner = tfdf.tuner.RandomSearch(num_trials=40, use_predefined_hps=True)
model = tfdf.keras.GradientBoostedTreesModel(tuner=tuner)
model.fit(train_ds)

use_predefined_hps=True searches a curated space that the YDF team maintains, which beats a hand-rolled grid in our experience. Budget forty trials on a CPU box; it is minutes, not hours. If forty trials buy you less than a point of AUC over the defaults, stop tuning - your ceiling is in the features, not the model.

The neural alternative, done properly

Reach for a network when you have high-cardinality identifiers that deserve learned embeddings, text or image columns alongside the tabular ones, or a need to fine-tune the whole thing jointly. Keras 3 preprocessing layers keep the transformation inside the model, so serving cannot drift from training:

import keras
from keras import layers

numeric = ["tenure_months", "monthly_spend", "support_tickets"]
categorical = ["plan", "region", "acquisition_channel"]

inputs, encoded = {}, []

for col in numeric:
    inp = keras.Input(shape=(1,), name=col, dtype="float32")
    norm = layers.Normalization()
    norm.adapt(train_df[col].values)     # adapt on TRAIN only
    inputs[col] = inp
    encoded.append(norm(inp))

for col in categorical:
    inp = keras.Input(shape=(1,), name=col, dtype="string")
    lookup = layers.StringLookup(output_mode="int")
    lookup.adapt(train_df[col].astype(str).values)
    emb = layers.Embedding(lookup.vocabulary_size(), 16)
    inputs[col] = inp
    encoded.append(layers.Flatten()(emb(lookup(inp))))

x = layers.Concatenate()(encoded)
x = layers.Dense(128, activation="relu")(x)
x = layers.Dropout(0.3)(x)
x = layers.Dense(64, activation="relu")(x)
out = layers.Dense(1, activation="sigmoid")(x)

nn = keras.Model(inputs, out)
nn.compile(optimizer=keras.optimizers.AdamW(1e-3),
           loss="binary_crossentropy",
           metrics=[keras.metrics.AUC(name="auc")])

Two rules we enforce on every tabular network:

  1. Adapt preprocessing layers on the training split only. Adapting on the full dataset leaks the validation distribution into the model and inflates your score. It is the most common silent bug we find in inherited tabular code.
  2. Keep Normalization and StringLookup inside the exported graph. If your service normalizes in pandas before calling the model, you now maintain the same arithmetic in two languages.

Feed it with tf.data and the input-pipeline hygiene from our tf.data tutorial:

def to_ds(df, shuffle=False):
    labels = df.pop("churned")
    ds = tf.data.Dataset.from_tensor_slices((dict(df), labels))
    if shuffle:
        ds = ds.shuffle(len(df))
    return ds.batch(512).prefetch(tf.data.AUTOTUNE)

Class imbalance, which tabular problems always have

Fraud and churn labels are rarely balanced. Do not oversample first - it is usually unnecessary and it distorts your calibrated probabilities. Start with class weights and an appropriate metric:

nn.fit(train_ds, validation_data=val_ds, epochs=30,
       class_weight={0: 1.0, 1: 12.0},
       callbacks=[keras.callbacks.EarlyStopping(monitor="val_auc", mode="max",
                                                patience=5, restore_best_weights=True)])

For TF-DF, GradientBoostedTreesModel(..., weights=...) or a weight column serves the same purpose. And judge these models on PR-AUC, not accuracy or ROC-AUC: at one percent positives, a model that predicts "no" every time is 99% accurate and worth nothing.

Which one won?

Score both on a held-out test set you touched exactly once, and compare on the metric the business actually cares about - expected saved dollars at your operating threshold, not AUC in the abstract. In client work the pattern is consistent:

  • Under roughly 100k rows, purely tabular: trees win, often by a clear margin, and train in a fraction of the time.
  • Millions of rows with strong ID features: the network catches up and embeddings start earning their keep.
  • Mixed tabular plus text or images: the network wins because trees cannot consume those columns at all.
  • Anything where you must explain individual decisions to a regulator: trees, plus per-prediction attributions, unless you have a compelling reason otherwise.

Ensembling the two - averaging a boosted-tree score with a network score - reliably adds a small amount and doubles your operational surface. We recommend it only when that small amount is worth real money.

Shipping it

Both paths export the same way, which is the practical argument for staying on this stack:

model.save("serving/churn_gbt/1")      # TF-DF -> SavedModel
nn.export("serving/churn_nn/1")        # Keras 3 -> SavedModel

From there it is the standard TensorFlow Serving deployment: versioned directories, canary labels, Prometheus metrics. Tabular models are also the ones most exposed to drift, because upstream business logic changes without anyone telling the ML team - so pair the deployment with schema and distribution monitoring the way we describe in Catching Data Drift Early with TensorFlow Data Validation.

One last piece of advice we give every client on day one: build the tree baseline in the first week, get it into production behind a feature flag, and only then argue about architectures. A deployed model with 0.78 AUC beats an unfinished one with 0.81 every quarter of the year.

If you want a second pair of eyes on a tabular pipeline - or a team to build it - our TensorFlow consultants do this work continuously. Get in touch with a description of your data and what you are trying to predict.