Writing

A neural network methylation classifier in TensorFlow, and when it beats a random forest

November 12, 2025 · Python, TensorFlow, machine learning, DNA methylation

In an earlier post I built a random forest classifier for brain tumor methylation profiles and evaluated it with nested cross-validation. The question I get most often after that one is some version of “why not a neural network?”

It’s a fair question. Most of my day-to-day work is methylation-based tumor classification (brain tumors, and more recently kidney tumors, as in “DNA methylation-based classification of kidney neoplasms”, Modern Pathology, 2025), and I use TensorFlow regularly. So this post builds a Keras multilayer perceptron (MLP) for the same task, trains it carefully, and compares it with the forest on identical folds. The goal is not to crown a winner. It’s to give you a setup where the comparison means something, so you can find out which model is better on your data.

Why a neural net is not automatically better here

Methylation classification is wide, small-n tabular data. The Capper et al. (2018) reference cohort, which I’ll use as the motivating dataset, has about 2,800 samples across 91 classes (82 tumor classes and 9 control tissue classes), and each 450k array measures over 450,000 CpG sites. Even after filtering to the few thousand most variable CpGs, you have more features than samples, many classes with only a handful of examples, and no spatial or sequential structure of the kind that convolutional or recurrent layers exploit.

That is the setting where tree ensembles tend to do well. Grinsztajn, Oyallon and Varoquaux (2022) benchmarked tree-based models against deep learning on medium-sized tabular datasets and found the trees still ahead on typical data. Among their explanations: neural networks are more hurt by uninformative features, and they are biased toward smooth decision functions when the true function is often irregular. Thousands of CpGs, most of which carry little information about any given class, is a fairly direct description of the first problem.

A random forest also asks very little of you. It doesn’t care whether you feed it beta values or M-values (trees are invariant to monotone transforms), it needs no scaling, it has few hyperparameters that matter, and it gives you out-of-bag estimates for free. An MLP needs scaling, regularization, a learning rate, an early stopping rule and a validation set. Each of those is a place to leak information or overfit.

None of that means the MLP will lose. It means the burden of proof is on it.

The data

The real reference cohort

The Capper reference cohort is public as GEO series GSE90496. How you get from GEO to a matrix (processed betas from the series supplementary files, or raw IDATs through minfi or SeSAMe with your own normalization) is a separate post. For what follows I assume you have:

• a samples-by-CpGs matrix of beta values, with CpG probe IDs as column names
• a vector of methylation class labels in the same sample order

import numpy as np
import pandas as pd

# Placeholder paths: point these at your own processed files.
betas = pd.read_parquet("gse90496_betas.parquet")   # rows: samples, columns: CpG IDs
labels = pd.read_csv("gse90496_labels.csv", index_col=0)["methylation_class"]
labels = labels.loc[betas.index]

X_beta = betas.to_numpy(dtype=np.float32)
cpg_ids = betas.columns.to_numpy()
class_names = np.array(sorted(labels.unique()))
y = np.searchsorted(class_names, labels.to_numpy())
n_classes = len(class_names)

y holds integer class codes from 0 to n_classes - 1, which is what the Keras loss below expects.

A synthetic stand-in

So that every block below runs, I generate a small synthetic dataset with scikit-learn and squash it into beta-like values between 0 and 1. It has 900 samples, 2,000 “CpGs” and 6 imbalanced classes. It is not methylation data and nothing about its results transfers to real tumors; it just has the right shape.

import numpy as np
from sklearn.datasets import make_classification

n_classes = 6
X_latent, y = make_classification(
    n_samples=900, n_features=2000, n_informative=40, n_redundant=0,
    n_classes=n_classes, n_clusters_per_class=1,
    weights=[0.35, 0.25, 0.15, 0.10, 0.10, 0.05],
    class_sep=1.5, random_state=0,
)
X_beta = (1.0 / (1.0 + 2.0 ** (-X_latent))).astype(np.float32)
cpg_ids = np.array([f"cg{i:08d}" for i in range(X_beta.shape[1])])
class_names = np.array([f"class_{k}" for k in range(n_classes)])

The transform is the inverse of the M-value formula, so the M-values of this fake data are just the original Gaussian features.

Outer folds first

The single most important thing for a fair comparison is that both models are evaluated on exactly the same held-out samples. I define the outer folds once, up front, and store them:

from sklearn.model_selection import StratifiedKFold

outer = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
folds = list(outer.split(X_beta, y))

If you saved fold assignments from the random forest work, load those instead. Everything learned from data (feature selection, scaling, class weights, early stopping, calibration) happens inside the training part of each outer fold. The outer test samples are touched once, for prediction.

Preprocessing inside the fold

Two steps: pick the most variable CpGs and convert beta values to M-values. Both are learned from training rows only. Selecting variable CpGs on the full dataset before splitting is the most common leak I see in methylation classifiers, and it makes every downstream number optimistic.

def beta_to_m(beta, eps=1e-3):
    beta = np.clip(beta, eps, 1 - eps)
    return np.log2(beta / (1 - beta))

def fit_preprocessing(beta_fit, n_cpgs):
    """Learn CpG selection and imputation values from the fitting rows only."""
    medians = np.nanmedian(beta_fit, axis=0)
    filled = np.where(np.isnan(beta_fit), medians, beta_fit)
    sd = filled.std(axis=0)
    keep = np.sort(np.argsort(sd)[::-1][:n_cpgs])
    return {"keep": keep, "medians": medians[keep]}

def apply_preprocessing(beta, prep):
    x = beta[:, prep["keep"]]
    x = np.where(np.isnan(x), prep["medians"], x)
    return beta_to_m(x).astype(np.float32)

I select on the standard deviation of beta values and then train on M-values. Beta values are bounded and compress differences near 0 and 1; M-values are roughly homoscedastic, which suits a model trained by gradient descent. The forest won’t care either way. The median imputation is there because real arrays have failed probes; the synthetic data has none.

Per-feature scaling (zero mean, unit variance) goes inside the model as a Normalization layer, adapted on the training rows. That way the scaling statistics are saved with the model and can’t drift out of sync with it.

The model

A plain MLP: normalization, a little input dropout, two hidden layers with dropout and L2 weight penalties, and a final layer that outputs logits. I leave the softmax out of the model on purpose. The loss can work from logits directly (from_logits=True, which is also numerically more stable), and temperature scaling needs the logits anyway.

import tensorflow as tf

def build_mlp(x_fit, n_classes, hidden=(256, 64), dropout=0.5,
              input_dropout=0.1, l2=1e-4, learning_rate=1e-3):
    norm = tf.keras.layers.Normalization(axis=-1, name="scale")
    norm.adapt(x_fit)
    reg = tf.keras.regularizers.L2(l2)

    inputs = tf.keras.Input(shape=(x_fit.shape[1],), name="m_values")
    x = norm(inputs)
    x = tf.keras.layers.Dropout(input_dropout)(x)
    for units in hidden:
        x = tf.keras.layers.Dense(units, activation="relu", kernel_regularizer=reg)(x)
        x = tf.keras.layers.Dropout(dropout)(x)
    logits = tf.keras.layers.Dense(n_classes, kernel_regularizer=reg, name="logits")(x)

    model = tf.keras.Model(inputs, logits)
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
    )
    return model

With a few thousand inputs, the first Dense layer holds most of the parameters (2,000 inputs times 256 units is already half a million weights, for 900 samples). That’s why the regularization is not optional. Dropout (Srivastava et al., 2014) on the hidden layers and a small amount on the inputs, plus L2, is the minimum. If you have very few samples per class, shrink the first layer before you add anything else.

Training: class weights, early stopping and the validation split

Early stopping needs a validation set, and that validation set has to come out of the outer training fold. Using the outer test fold to decide when to stop is the neural-net version of tuning on the test set. So each outer training fold is split again: about 85% for fitting, 15% for validation.

from scipy.optimize import minimize_scalar
from scipy.special import log_softmax
from sklearn.model_selection import train_test_split
from sklearn.utils.class_weight import compute_class_weight

def fit_temperature(logits, y_true):
    """Find T > 0 minimizing the negative log-likelihood of softmax(logits / T)."""
    def nll(log_t):
        logp = log_softmax(logits / np.exp(log_t), axis=1)
        return -logp[np.arange(len(y_true)), y_true].mean()
    result = minimize_scalar(nll, bounds=(np.log(0.05), np.log(20.0)), method="bounded")
    return float(np.exp(result.x))

def train_mlp_fold(X_beta, y, train_idx, n_classes, n_cpgs=1000, seed=0):
    fit_idx, val_idx = train_test_split(
        train_idx, test_size=0.15, stratify=y[train_idx], random_state=seed)

    prep = fit_preprocessing(X_beta[fit_idx], n_cpgs)
    x_fit = apply_preprocessing(X_beta[fit_idx], prep)
    x_val = apply_preprocessing(X_beta[val_idx], prep)

    weights = compute_class_weight("balanced", classes=np.arange(n_classes), y=y[fit_idx])
    class_weight = {k: float(w) for k, w in enumerate(weights)}

    tf.keras.utils.set_random_seed(seed)
    model = build_mlp(x_fit, n_classes)
    early_stop = tf.keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=20, restore_best_weights=True)
    model.fit(
        x_fit, y[fit_idx],
        validation_data=(x_val, y[val_idx]),
        epochs=300, batch_size=64,
        class_weight=class_weight,
        callbacks=[early_stop],
        verbose=0,
    )

    val_logits = model.predict(x_val, verbose=0)
    temperature = fit_temperature(val_logits, y[val_idx])
    return model, prep, temperature

A few notes on the choices.

Class weights. compute_class_weight("balanced", ...) weights each class inversely to its frequency, so a rare class contributes as much to the loss as a common one. Without it, an MLP on imbalanced data happily learns to predict the big classes. Two side effects to know about: the validation loss that early stopping watches is unweighted (you can pass sample weights as a third element of validation_data if you want it weighted), and the resulting probabilities reflect something closer to a balanced prior than the training class frequencies. That second point matters if you later interpret scores as probabilities in a population with very different class frequencies.

Early stopping. restore_best_weights=True is the part people forget. Without it, the model you keep is the one from the last epoch, which is patience epochs past the best one. With small validation sets the validation loss is noisy, so I use a generous patience.

Stratified inner split. stratify= needs at least two samples of every class in the outer training fold. With the real Capper classes, some of which are small, check this before running, and consider fewer outer folds or merging the rarest classes for a first pass.

tf.data. For data that fits in memory, NumPy arrays are fine. If you want a tf.data pipeline (for larger cohorts, or to stream from disk later), it’s a drop-in replacement, and class_weight still works:

def make_dataset(x, y, batch_size=64, shuffle=False, seed=0):
    ds = tf.data.Dataset.from_tensor_slices((x, y))
    if shuffle:
        ds = ds.shuffle(len(y), seed=seed, reshuffle_each_iteration=True)
    return ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)

# Inside train_mlp_fold, instead of passing arrays:
# model.fit(make_dataset(x_fit, y[fit_idx], shuffle=True, seed=seed),
#           validation_data=make_dataset(x_val, y[val_idx]),
#           epochs=300, class_weight=class_weight, callbacks=[early_stop], verbose=0)

When fit gets arrays it shuffles each epoch itself; with a dataset you have to ask for it, as above. Recent Keras versions print a warning that fit’s own shuffle argument is ignored for datasets, which is expected.

Temperature scaling

Neural networks trained with cross-entropy tend to be overconfident: the top softmax probability is systematically higher than the accuracy at that confidence level. Guo et al. (2017) showed this for modern networks and found that a one-parameter fix, temperature scaling, works surprisingly well.

The idea is simple. Divide the logits by a single scalar T before the softmax:

p = softmax(logits / T)

T > 1 flattens the probabilities (less confident), T < 1 sharpens them. T is fitted by minimizing negative log-likelihood on held-out data, which is what fit_temperature above does on the inner validation set. Because dividing every logit by the same positive number doesn’t change which one is largest, temperature scaling never changes the predicted class. Balanced accuracy is identical before and after; only the probabilities move.

One compromise in my code: the same inner validation set is used for early stopping and for fitting T. A separate calibration split would be cleaner, but with small classes I’d rather keep the samples for training. With a single parameter, the overfitting risk is small.

Calibration matters more than usual in this field because methylation classifier scores get thresholded. The Capper classifier itself reports calibrated scores (the paper calibrates the random forest output with a penalized multinomial logistic regression), and decisions like “report this class” or “call this result inconclusive” depend on those numbers meaning roughly what they say.

The random forest on the same folds

For the comparison, the forest gets the same outer folds and the same fold-internal CpG selection. Use the settings you ended up with in the earlier post; the ones below are a reasonable stand-in. I also fit a calibrated version with scikit-learn’s CalibratedClassifierCV, so both models get a calibration step.

from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import RandomForestClassifier

def train_rf_fold(X_beta, y, train_idx, n_cpgs=1000, seed=0, calibrate=False):
    prep = fit_preprocessing(X_beta[train_idx], n_cpgs)
    x_train = apply_preprocessing(X_beta[train_idx], prep)
    rf = RandomForestClassifier(
        n_estimators=500, class_weight="balanced", n_jobs=-1, random_state=seed)
    if calibrate:
        rf = CalibratedClassifierCV(rf, method="sigmoid", cv=3)
    rf.fit(x_train, y[train_idx])
    return rf, prep

The forest trains on the full outer training fold, while the MLP gives up 15% of it to validation. I think that’s the fair way to compare: needing a validation set is a real cost of the MLP, and the comparison should include it.

Evaluating

Plain accuracy is misleading with imbalanced classes, so the headline metric is balanced accuracy (the mean of per-class recall). For probabilities I look at log loss and a simple expected calibration error (ECE): bin predictions by their top probability and compare the average confidence with the observed accuracy in each bin.

from scipy.special import softmax
from sklearn.metrics import balanced_accuracy_score, confusion_matrix, log_loss, recall_score

def expected_calibration_error(y_true, proba, n_bins=10):
    confidence = proba.max(axis=1)
    correct = proba.argmax(axis=1) == y_true
    bins = np.minimum((confidence * n_bins).astype(int), n_bins - 1)
    ece = 0.0
    for b in range(n_bins):
        in_bin = bins == b
        if in_bin.any():
            ece += in_bin.mean() * abs(correct[in_bin].mean() - confidence[in_bin].mean())
    return ece

def evaluate(y_true, proba, n_classes):
    pred = proba.argmax(axis=1)
    return {
        "balanced_accuracy": balanced_accuracy_score(y_true, pred),
        "log_loss": log_loss(y_true, proba, labels=np.arange(n_classes)),
        "ece": expected_calibration_error(y_true, proba),
    }

Now the outer loop. Each fold trains both models, predicts the held-out samples once, and stores out-of-fold probabilities for every sample:

import pandas as pd

model_names = ["mlp_raw", "mlp_temp", "rf", "rf_calibrated"]
oof = {name: np.zeros((len(y), n_classes)) for name in model_names}
records = []

for k, (train_idx, test_idx) in enumerate(folds):
    mlp, mlp_prep, temperature = train_mlp_fold(X_beta, y, train_idx, n_classes, seed=k)
    logits = mlp.predict(apply_preprocessing(X_beta[test_idx], mlp_prep), verbose=0)
    logits = logits.astype(np.float64)   # float32 softmax rows can miss summing to 1
    oof["mlp_raw"][test_idx] = softmax(logits, axis=1)
    oof["mlp_temp"][test_idx] = softmax(logits / temperature, axis=1)

    for name, calibrate in [("rf", False), ("rf_calibrated", True)]:
        rf, rf_prep = train_rf_fold(X_beta, y, train_idx, seed=k, calibrate=calibrate)
        oof[name][test_idx] = rf.predict_proba(apply_preprocessing(X_beta[test_idx], rf_prep))

    for name in model_names:
        records.append({"fold": k, "model": name,
                        **evaluate(y[test_idx], oof[name][test_idx], n_classes)})
    print(f"fold {k}: temperature = {temperature:.2f}")

scores = pd.DataFrame(records)
scores.groupby("model")[["balanced_accuracy", "log_loss", "ece"]].agg(["mean", "std"])

predict_proba returns columns in the order of the forest’s classes_, which matches 0 ... n_classes - 1 as long as every class appears in the training fold. Stratified folds make sure of that, but it’s worth an assert if your smallest classes are tiny.

Because the folds are shared, compare models fold by fold rather than by their averages alone:

paired = scores.pivot(index="fold", columns="model", values="balanced_accuracy")
(paired["mlp_temp"] - paired["rf"]).describe()

If the per-fold differences flip sign from fold to fold, you don’t have evidence that either model is better, whatever the means say. Five folds is not much to go on; repeating the whole outer loop with a few different fold seeds gives a better sense of the spread.

Then look at where each model makes its mistakes, using the pooled out-of-fold predictions:

def per_class_report(y_true, proba, class_names):
    pred = proba.argmax(axis=1)
    labels = np.arange(len(class_names))
    recall = pd.Series(recall_score(y_true, pred, labels=labels, average=None, zero_division=0),
                       index=class_names, name="recall")
    cm = pd.DataFrame(confusion_matrix(y_true, pred, labels=labels),
                      index=pd.Index(class_names, name="true"),
                      columns=pd.Index(class_names, name="predicted"))
    return recall, cm

recall_mlp, cm_mlp = per_class_report(y, oof["mlp_temp"], class_names)
recall_rf, cm_rf = per_class_report(y, oof["rf"], class_names)
pd.concat({"mlp": recall_mlp, "rf": recall_rf}, axis=1)

Per-class recall is usually more informative than any single summary number. Two models with similar balanced accuracy can fail on completely different classes, and in methylation classification the confusions are often biologically sensible (neighboring subtypes of the same tumor family), which is a different kind of error from confusing unrelated classes.

Reproducibility

tf.keras.utils.set_random_seed(seed) seeds Python’s random, NumPy and TensorFlow in one call, and it’s what train_mlp_fold uses. That makes weight initialization, dropout masks and shuffling repeatable, but it does not make results bit-for-bit identical on every machine. Some GPU kernels use non-deterministic reductions, so two runs on a GPU with the same seed can differ slightly. TensorFlow has a switch for this:

tf.config.experimental.enable_op_determinism()

Call it at the start of the script, before building models. It makes ops deterministic where TensorFlow can, raises an error for ops that have no deterministic implementation, and can slow training down. For a model this small, CPU training is fast enough that I often skip the GPU for this kind of comparison.

The more useful habit is to treat the seed as a source of variance. Train with three to five seeds per fold and report the spread. If the MLP’s advantage over the forest is smaller than its seed-to-seed variation, there isn’t an advantage.

Pin your package versions too (pip freeze > requirements.txt, or a lock file). TensorFlow and Keras changed a lot between 2.15 and Keras 3, and a model that trained fine on one may behave differently on the other.

Saving the model

Use the native .keras format. It stores the architecture, weights and the Normalization statistics in a single file. The CpG selection, imputation values, temperature and class names live outside the model, so I save them next to it in JSON:

import json

model, prep, temperature = train_mlp_fold(X_beta, y, folds[0][0], n_classes, seed=0)
model.save("mlp_fold0.keras")

with open("mlp_fold0_preprocessing.json", "w") as f:
    json.dump({
        "cpg_ids": cpg_ids[prep["keep"]].tolist(),
        "impute_medians": prep["medians"].tolist(),
        "temperature": temperature,
        "class_names": class_names.tolist(),
    }, f)

restored = tf.keras.models.load_model("mlp_fold0.keras")

Storing CpG IDs rather than column positions matters: a new sample’s beta matrix may have its probes in a different order, or come from a different array version with some probes missing. Select by ID, and fail loudly if any are absent.

The cross-validated models are for estimating performance. For an actual deployable model, retrain on all the data with the same procedure (holding out a validation slice for early stopping and temperature), and report the cross-validated estimate as the expected performance of that procedure.

Where a neural net earns its place

I’d keep the random forest as the default for single-platform methylation classification at reference-cohort scale. It is quick to train, hard to break, insensitive to preprocessing choices, and easy to explain to a pathologist. The MLP is worth the extra work when your problem has one of these features.

More samples. The trade-off shifts as cohorts grow. With tens of thousands of profiles, the MLP’s capacity starts to be useful instead of a liability, and mini-batch training scales better than growing ever larger forests.

Multi-task setups. Methylation classes are hierarchical (class families and subclasses), and you may want to predict other things from the same profile. A shared trunk with several heads is natural in Keras and awkward with a forest:

def build_two_head_mlp(x_fit, n_families, n_classes, l2=1e-4):
    norm = tf.keras.layers.Normalization(axis=-1)
    norm.adapt(x_fit)
    reg = tf.keras.regularizers.L2(l2)
    inputs = tf.keras.Input(shape=(x_fit.shape[1],))
    x = norm(inputs)
    x = tf.keras.layers.Dense(256, activation="relu", kernel_regularizer=reg)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    family = tf.keras.layers.Dense(n_families, name="family")(x)
    subclass = tf.keras.layers.Dense(n_classes, name="subclass")(x)
    model = tf.keras.Model(inputs, {"family": family, "subclass": subclass})
    model.compile(
        optimizer="adam",
        loss={"family": tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              "subclass": tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)},
        loss_weights={"family": 0.5, "subclass": 1.0},
    )
    return model

You’d fit it with y={"family": y_family, "subclass": y_class}. Keras doesn’t support class_weight for multi-output models, so rebalancing has to go through sample_weight instead.

Combining inputs. A methylation array also gives you a copy number profile, and many studies have sequencing calls and clinical variables for the same samples (our “Multiplatform molecular analyses refine classification of gliomas arising in patients with neurofibromatosis type 1” paper in Acta Neuropathologica, 2022, is that kind of study). A network with one input branch per data type, joined before the output layer, is a clean way to combine them. With a forest you’d concatenate everything into one wide table and hope.

Representation learning. Unlabeled methylation profiles are much more plentiful than labeled reference samples. Pretraining an encoder on them, or training with augmentations that mimic real-world problems such as low tumor purity, is something neural nets do and forests don’t. Sample quality is a real issue in clinical use (how much tumor DNA a classifier needs is the question behind our paper “Tumor DNA requirements for accurate epigenetic-based classification of CNS neoplasia”, Neuro-Oncology, 2021), and a model trained to cope with it is one place I’d expect a network to justify itself.

If none of those apply, the forest is the sensible default, and the main thing this post buys you is a fair way to check.

Checklist before claiming the MLP wins

1. Both models use the same outer folds, defined once and saved.
2. CpG selection, imputation and scaling are fitted on training rows only, inside every fold.
3. Early stopping and temperature scaling use a validation split carved from the outer training fold, never the test fold.
4. restore_best_weights=True is set.
5. Class imbalance is handled for both models (class_weight for the MLP, class_weight="balanced" or equivalent for the forest).
6. You report balanced accuracy, per-class recall, a confusion matrix, and a calibration measure, not accuracy alone.
7. Differences are compared fold by fold and across several seeds, and they’re larger than the seed-to-seed spread.
8. Both models received comparable tuning effort.
9. The saved model comes with its CpG IDs, imputation values, temperature and class names.

References

• Capper D, et al. DNA methylation-based classification of central nervous system tumours. Nature 555, 469–474 (2018).
• Grinsztajn L, Oyallon E, Varoquaux G. Why do tree-based models still outperform deep learning on typical tabular data? NeurIPS Datasets and Benchmarks Track (2022).
• Guo C, Pleiss G, Sun Y, Weinberger KQ. On Calibration of Modern Neural Networks. ICML (2017).
• Srivastava N, et al. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. JMLR 15, 1929–1958 (2014).
• Abadi M, et al. TensorFlow: A system for large-scale machine learning. OSDI (2016).

← All writing