Diagnosing Overfitting in a PPI Transformer: Hyperparameter Search and What Actually Helped

Summary

A hyperparameter search was conducted to diagnose and mitigate overfitting in a transformer-based model for predicting Protein-Protein Interactions (PPIs). The search explored variations in dropout rates, latent dimensions, and weight decay. The results showed that increasing dropout and reducing model size mildly improved generalization, while weight decay had a less pronounced effect. The best configuration achieved an AUROC improvement of ~0.02 on the leakage-reduced test set compared to the baseline. These results suggest that while hyperparameter tuning can provide modest gains, the model may have reached a possible performance ceiling given the current architecture and mean-pooled ESM2 embeddings. Further improvements are likely to require architectural changes or data augmentation strategies.

The problem

In my previous post, I built a transformer-based model to predict PPIs trained from the HuRI dataset. The model achieved an AUROC of 0.75 and accuracy of 0.67 on leakage-reduced test data, which is comparable to the performance of existing models in the literature (Reim et al. 2025). However, looking at the training and validation AUROC curves, it was clear that the model was overfitting to the training data quite early on. The training AUROC quickly reached 0.9 within the first few epochs, while the validation AUROC plateaued around 0.74 and even started to decline after epoch 10. This indicated that the model was learning patterns specific to the training data that did not generalize well to unseen data.

Training and validation AUROC curves showing overfitting.

Hypothesis

I hypothesized that the overfitting was due to the model being too complex for the amount of training data available, and that tuning the hyperparameters could help mitigate this issue. Specifically, I wanted to investigate how different hyperparameters such as latent dimension (model size), dropout rate, and weight decay affected the model’s performance and generalization ability.

Experimental design

Let’s look at the model architecture again:

Model architecture diagram.

I defined 8 sets of hyperparameter configurations to test, including variations in dropout rates for the projection, attention, and MLP layers, as well as different latent dimensions and weight decay values. The configurations were designed to systematically explore the effects of increasing dropout, reducing model size, and applying stronger regularization through weight decay.

# Define hp space
# Label, Dropout (proj, attn, mlp), latent dim, weight_decay
CONFIGS = [
    ("baseline", (0.1, 0.1, 0.3), 512, 0.1),
    ("high_dropout", (0.3, 0.3, 0.5), 512, 0.1),
    ("small_model", (0.1, 0.1, 0.3), 256, 0.1),
    ("small+dropout", (0.2, 0.2, 0.4), 256, 0.1),
    ("small+wd", (0.1, 0.1, 0.3), 256, 0.3),
    ("small+all", (0.2, 0.2, 0.4), 256, 0.3),
    ("tiny_model", (0.1, 0.1, 0.3), 128, 0.1),
    ("tiny+dropout", (0.3, 0.3, 0.5), 128, 0.2),
]

Many of the training utilities were already defined in the previous post so will not be shown here, such as:

• PPIModel, train, evaluate, get_loss_fn, MetricHistory
  
• train_loader, val_loader, test_loader
  
• train_dataset, val_dataset, test_dataset
  

Here are some of the new functions that I defined for this hyperparameter search:

• build_model_and_optimizer: This function takes in the hyperparameters and builds the PPIModel, sets up the optimizer with appropriate weight decay, and defines a learning rate scheduler. It also counts the number of trainable parameters in the model for later analysis.
  
• run_hyperparameter_search: This function iterates through the defined CONFIGS, builds the model for each configuration, and runs the training loop. It records the training history, number of parameters, and checkpoint path for each configuration in a dictionary.
  
• evaluate_all_on_test: After training all models, this function loads each model from its checkpoint and evaluates it on the test set, recording various metrics such as AUROC, AUPRC, F1 score, MCC, and accuracy in a DataFrame for easy comparison.

# Function to instantiate model and utilities
def build_model_and_optimizer(
    proj_drop, attn_drop, mlp_drop, latent_dim, weight_decay,
    lr=1e-4,
):
    torch.manual_seed(42)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"  Using device: {device}")

    model = PPIModel(
        esm2_dim = 1280,
        latent_dim = latent_dim,
        num_heads = max(4, latent_dim // 64),  # scale heads with latent_dim
        num_layers = 3,
        ffn_dim = latent_dim * 2,
        hidden_dim = latent_dim // 2,
        proj_drop = proj_drop,
        attn_drop = attn_drop,
        mlp_drop = mlp_drop,
    ).to(device)

    decay_params = []
    no_decay_params = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if 'norm' in name or 'bias' in name or param.dim() == 1:
            no_decay_params.append(param)
        else:
            decay_params.append(param)

    optimizer = torch.optim.AdamW([
    {'params': decay_params, 'weight_decay': weight_decay},
    {'params': no_decay_params, 'weight_decay': 0.0},
    ], lr=1e-4)

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='max', patience=5, factor=0.5, min_lr=1e-7
    )

    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    return model, optimizer, scheduler, device, n_params

# Function to run hp search
def run_hyperparameter_search(
    configs = CONFIGS,
    train_loader = None,
    val_loader = None,
    loss_fn = None,
    save_path = None,
    n_epochs = 100,
    early_stop_limit = 15,
):
    """
    Iterate through CONFIGS to build model and perform training.
    Record metrics and hp information in a dictionary.
    """

    all_results = {} # record: history, n_params, config, checkpoint

    for i, (label, dropout, latent_dim, weight_decay) in enumerate(configs):

      proj_drop, attn_drop, mlp_drop = dropout

      print(f"\n{'═'*65}")
      print(f"  Config {i+1}/{len(configs)}: {label}")
      print(f"  dropout=({proj_drop},{attn_drop},{mlp_drop})  "
                f"latent_dim={latent_dim}  weight_decay={weight_decay}")
      print(f"{'═'*65}")

      model, optimizer, scheduler, device, n_params = build_model_and_optimizer(
          proj_drop, attn_drop, mlp_drop, latent_dim, weight_decay
      )
      print(f"  Parameters: {n_params:,}  ({n_params/len(train_loader.dataset):.1f} params/sample)")

      history = train(
          model = model,
          train_loader = train_loader,
          val_loader = val_loader,
          optimizer = optimizer,
          scheduler = scheduler,
          loss_fn = loss_fn,
          device = device,
          n_epochs = n_epochs,
          early_stop_limit = early_stop_limit,
          checkpoint_path = f"{save_path}/data/hparam_{label.replace('+','_')}.pt",
      )

      all_results[label] = {
          "history": history,
          "n_params": n_params,
          "config": (dropout, latent_dim, weight_decay),
          "ckpt": f"{save_path}/data/hparam_{label.replace('+','_')}.pt",
      }

    print("All hyperparameter searches completed.")
    return all_results

# Function to evaluate on test data
def evaluate_all_on_test(all_results, test_loader, loss_fn):
    """
    Evaluate all models on test data.
    Build model -> load checkpoint -> evaluate -> record metrics.
    Return a pd.DataFrame with all metrics.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    rows = []

    for label, res in all_results.items():
      dropout, latent_dim, weight_decay = res["config"]
      proj_drop, attn_drop, mlp_drop = dropout

      model, _, _, _, n_params = build_model_and_optimizer(
          proj_drop, attn_drop, mlp_drop, latent_dim, weight_decay
      )

      model.load_state_dict(torch.load(res["ckpt"]))
      model.to(device)
      model.eval()

      _, auroc, auprc, probs, labels = evaluate(
          model, test_loader, loss_fn, device
      )

      preds = (probs >= 0.5).astype(int)

      rows.append({
            "config": label,
            "n_params": n_params,
            "latent_dim": latent_dim,
            "dropout": f"({proj_drop},{attn_drop},{mlp_drop})",
            "weight_decay": weight_decay,
            "test_auroc": round(auroc, 4),
            "test_auprc": round(auprc, 4),
            "f1": round(f1_score(labels, preds), 4),
            "precision": round(precision_score(labels, preds), 4),
            "recall": round(recall_score(labels, preds), 4),
            "mcc": round(matthews_corrcoef(labels, preds), 4),
            "accuracy": round(accuracy_score(labels, preds), 4),
            # Best val AUROC from training (for reference)
            "best_val_auroc": round(max(res["history"].history["val_auroc"]), 4),
            # Train-val gap at best epoch
            "gap_at_best": round(
                res["history"].history["train_auroc"][
                    int(np.argmax(res["history"].history["val_auroc"]))
                ] - max(res["history"].history["val_auroc"]), 4
            ),
        })

      print(f"  {label:20s}  test_auroc={auroc:.4f}  test_auprc={auprc:.4f}  "
              f"mcc={rows[-1]['mcc']:.4f}")

    df = pd.DataFrame(rows).sort_values("test_auroc", ascending=False).reset_index(drop=True)
    # columns are metrics, rows are configs
    return df

Results & discussions

The different hyperparameter configurations had varying effects on the training dynamics and final performance. To visualize these differences, I created several plotting functions:

• plot_learning_curves_comparison: This function plots the validation AUROC, AUPRC, validation loss, and the gap between training and validation AUROC for all configurations on the same axes. This allows for easy comparison of how each configuration affected the learning curves and overfitting behavior.
  
• plot_train_val_curves_per_config: This function creates a grid of plots, where each plot shows the training and validation AUROC curves for a single configuration. This is a more traditional way to diagnose overfitting for each model individually.
  
• plot_test_metrics_comparison: This function creates bar charts comparing the test AUROC, AUPRC, F1 score, and MCC across all configurations. The bars are annotated with the metric values and the best-performing configuration is highlighted.
  

def plot_learning_curves_comparison(all_results, save_path=None):
    """
    Plots val AUROC, val AUPRC, val Loss, and train/val AUROC gap
    for all configs on the same axes for easy comparison.
    """
    labels = list(all_results.keys())
    colors = cm.tab10(np.linspace(0, 1, len(labels)))

    fig, axes = plt.subplots(2, 2, figsize=(16, 10))
    fig.suptitle("Hyperparameter Comparison — Learning Curves", fontsize=14, fontweight='bold')

    metrics = [
        ("val_auroc",  axes[0, 0], "Val AUROC",  "max", [0.5, 1.0]),
        ("val_auprc",  axes[0, 1], "Val AUPRC",  "max", [0.5, 1.0]),
        ("val_loss",   axes[1, 0], "Val Loss",   "min", None),
        (None,         axes[1, 1], "Train−Val AUROC Gap", "min", None),  # special
    ]

    for metric_key, ax, title, _, ylim in metrics:
        for label, color in zip(labels, colors):
            h = all_results[label]["history"].history
            epochs = range(1, len(h["train_auroc"]) + 1)

            if metric_key is None:
                # Gap plot
                gap = [tr - vl for tr, vl in zip(h["train_auroc"], h["val_auroc"])]
                ax.plot(epochs, gap, label=label, color=color, linewidth=1.8)
                ax.axhline(y=0, color='gray', linestyle='--', alpha=0.4)
            else:
                ax.plot(epochs, h[metric_key], label=label, color=color, linewidth=1.8)

            # Mark best val AUROC epoch with a dot
            best_ep = int(np.argmax(h["val_auroc"])) + 1
            best_val = h["val_auroc"][best_ep - 1]
            if metric_key == "val_auroc":
                ax.scatter(best_ep, best_val, color=color, s=60, zorder=5)

        ax.set_title(title, fontweight='bold')
        ax.set_xlabel("Epoch")
        ax.grid(True, alpha=0.3)
        if ylim:
            ax.set_ylim(ylim)
        ax.legend(fontsize=7, loc='best')

    plt.tight_layout()
    if save_path:
        out = f"{save_path}/hparam_learning_curves.png"
        plt.savefig(out, dpi=150, bbox_inches='tight')
        print(f"Saved → {out}")
    plt.show()

def plot_train_val_curves_per_config(all_results, save_path=None):
    """
    For each config, plots train vs val AUROC side by side —
    the standard overfitting diagnostic view.
    """
    n = len(all_results)
    ncols = 4
    nrows = int(np.ceil(n / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 4.5, nrows * 3.5))
    axes = axes.flatten()

    fig.suptitle("Train vs Val AUROC per Config", fontsize=13, fontweight='bold')

    for i, (label, res) in enumerate(all_results.items()):
        h   = res["history"].history
        ax  = axes[i]
        eps = range(1, len(h["train_auroc"]) + 1)
        best_ep = int(np.argmax(h["val_auroc"])) + 1

        ax.plot(eps, h["train_auroc"], label="Train", color="steelblue", lw=1.8)
        ax.plot(eps, h["val_auroc"],   label="Val",   color="coral",     lw=1.8)
        ax.axvline(x=best_ep, color='gray', linestyle='--', alpha=0.6,
                   label=f"Best ep {best_ep}")

        best_val  = max(h["val_auroc"])
        gap       = h["train_auroc"][best_ep-1] - best_val
        n_params  = res["n_params"]
        ax.set_title(f"{label}\nBest val={best_val:.4f}  gap={gap:.3f}\n{n_params:,} params",
                     fontsize=8)
        ax.set_ylim([0.5, 1.0])
        ax.set_xlabel("Epoch", fontsize=8)
        ax.legend(fontsize=7)
        ax.grid(True, alpha=0.3)

    # Hide unused subplots
    for j in range(i + 1, len(axes)):
        axes[j].set_visible(False)

    plt.tight_layout()
    if save_path:
        out = f"{save_path}/hparam_per_config_curves.png"
        plt.savefig(out, dpi=150, bbox_inches='tight')
        print(f"Saved → {out}")
    plt.show()

def plot_test_metrics_comparison(results_df, save_path=None):
    """
    Bar chart comparing test AUROC, AUPRC, F1, and MCC across all configs.
    Configs sorted by test AUROC descending.
    """
    metrics = ["test_auroc", "test_auprc", "f1", "mcc"]
    titles = ["Test AUROC", "Test AUPRC", "F1 Score", "MCC"]
    labels = results_df["config"].tolist()
    x = np.arange(len(labels))
    colors = cm.tab10(np.linspace(0, 1, len(labels)))

    fig, axes = plt.subplots(2, 2, figsize=(16, 9))
    fig.suptitle("Test Set Performance by Hyperparameter Config", fontsize=13, fontweight='bold')

    for ax, metric, title in zip(axes.flatten(), metrics, titles):
        vals = results_df[metric].tolist()
        bars = ax.bar(x, vals, color=colors, edgecolor='white', linewidth=0.5)

        # Annotate bars
        for bar, v in zip(bars, vals):
            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.005,
                    f"{v:.3f}", ha='center', va='bottom', fontsize=7)

        ax.set_title(title, fontweight='bold')
        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=35, ha='right', fontsize=8)
        ax.set_ylim([max(0, min(vals) - 0.05), min(1.0, max(vals) + 0.08)])
        ax.grid(True, axis='y', alpha=0.3)

        # Highlight best
        best_idx = int(np.argmax(vals))
        bars[best_idx].set_edgecolor('black')
        bars[best_idx].set_linewidth(2.5)

    plt.tight_layout()
    if save_path:
        out = f"{save_path}/hparam_test_metrics.png"
        plt.savefig(out, dpi=150, bbox_inches='tight')
        print(f"Saved → {out}")
    plt.show()

Now let’s run everything and discuss the results as we go along.

# Run hp search
loss_fn = get_loss_fn(label_smoothing=0.1, pos_weight=None)

save_path = "PPI_prediction"

all_results = run_hyperparameter_search(
    configs = CONFIGS,
    train_loader = train_loader,
    val_loader = val_loader,
    loss_fn = loss_fn,
    save_path = save_path,
    n_epochs = 100,
    early_stop_limit = 15,
)

# Plot learning curves: compare configs
plot_learning_curves_comparison(all_results, save_path=save_path)

• Looking at validation AUROC and AUPRC, the metrics appear to peak at around 0.74 no matter the model configuration, suggesting that this might be the performance plateau that the current data setting can reach - meaning that the bottleneck is likely not the model architecture, but the limitation of mean-pooled ESM2 embeddings and the size of the HuRI dataset to generalize to unseen PPIs.

• Looking at the train-val AUROC gap curve, it appears that reducing the model size or increasing regularization/dropouts did reduce the severity of overfitting, albeit not sufficient to significantly improve the best validation AUROC.

Let’s look at the actual train vs val AUROC curves for each configuration to see how the overfitting dynamics differ:

# Plot learning curves: compare train vs. val per config
plot_train_val_curves_per_config(all_results, save_path=save_path)

• Looking at the per-config train-vs-val curves, the reduction in the train-val gap is driven primarily by a lower training AUROC at later epochs — not by a meaningful increase in validation AUROC.

Now we will evaluate the best model checkpoint for each config on the leakage-reduced test set to see if there is any improvement in test performance:

results_df = evaluate_all_on_test(all_results, test_loader, loss_fn)
print(results_df.to_string(index=False))

Output:

       config  n_params  latent_dim       dropout  weight_decay  test_auroc  test_auprc     f1  precision  recall    mcc  accuracy  best_val_auroc  gap_at_best
small+dropout   2666241         256 (0.2,0.2,0.4)           0.1      0.7690      0.7683 0.5110     0.8101  0.3732 0.3393    0.6429          0.7375       0.2000
     small+wd   2666241         256 (0.1,0.1,0.3)           0.3      0.7666      0.7451 0.5851     0.7854  0.4662 0.3709    0.6694          0.7356       0.1759
  small_model   2666241         256 (0.1,0.1,0.3)           0.1      0.7663      0.7448 0.5867     0.7848  0.4684 0.3715    0.6700          0.7363       0.1751
    small+all   2666241         256 (0.2,0.2,0.4)           0.3      0.7653      0.7693 0.5401     0.8180  0.4031 0.3636    0.6567          0.7443       0.1988
   tiny_model    833409         128 (0.1,0.1,0.3)           0.1      0.7636      0.7623 0.5513     0.7837  0.4252 0.3462    0.6539          0.7358       0.2144
 high_dropout   9330177         512 (0.3,0.3,0.5)           0.1      0.7619      0.7538 0.5873     0.7548  0.4806 0.3483    0.6622          0.7439       0.1814
 tiny+dropout    833409         128 (0.3,0.3,0.5)           0.2      0.7599      0.7469 0.5802     0.7770  0.4629 0.3608    0.6650          0.7343       0.1945
     baseline   9330177         512 (0.1,0.1,0.3)           0.1      0.7466      0.7470 0.6314     0.7195  0.5626 0.3518    0.6717          0.7390       0.1688

# Plot test metrics across configs
plot_test_metrics_comparison(results_df, save_path=save_path)

Test AUROC ranked by AUROC

Test AUPRC ranked by AUROC

• Looking at the test metrics, which is the most important indication of how the model is performing, the smaller models with stronger dropouts/regularization generally improved the performance when compared with the baseline. Among these models, the “small+dropout” showed consistently higher AUROC and AUPRC, which both improved by ~0.02, suggesting that this model config is most adequate for the data.
• Why is there a disconnect between some AUROCs vs AUPRCs? AUROC measures the model’s ability to rank positives above negatives across all classification thresholds, treating both classes symmetrically. AUPRC, by contrast, focuses on the precision-recall trade-off and is more sensitive to how well the model identifies true positives with high confidence. Even on balanced datasets, the two metrics can diverge when a model ranks most positives correctly (good AUROC) but with poor confidence separation at high-score ranges (lower AUPRC). This is worth keeping in mind when calling out the best-performing model — a configuration that wins on AUROC may not win on AUPRC, as we see here with small+dropout vs small+all.

Future directions

Now that there is a better understanding of how hyperparameter tuning affects the performance of the PPI transformer model, perhaps there are other avenues to explore. In future posts, I’ll investigate whether per-token (non-mean-pooled) embeddings or isoform-based data augmentation — selecting multiple isoforms per protein within a defined similarity band — can push past this ceiling.

Disclosures

The code was written with the aid of Claude Sonnet 4.6, and the model training was performed on Google Colab Pro+ with a Tesla T4 GPU. The accuracy of the codes were verified by the author.

References

• Reim T, Hartebrodt A, Blumenthal DB, Bernett J, List M. Deep learning models for unbiased sequence-based PPI prediction plateau at an accuracy of 0.65. Bioinformatics. 2025 Jul 1;41(Supplement_1):i590-i598. doi: 10.1093/bioinformatics/btaf192. PMID: 40662806; PMCID: PMC12261406.
• Lin Z, Akin H, Rao R, Hie B, Zhu Z, Lu W, Smetanin N, Verkuil R, Kabeli O, Shmueli Y, Dos Santos Costa A, Fazel-Zarandi M, Sercu T, Candido S, Rives A. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science. 2023 Mar 17;379(6637):1123-1130. doi: 10.1126/science.ade2574. Epub 2023 Mar 16. PMID: 36927031.
• Luck K, Kim DK, Lambourne L, Spirohn K, Begg BE, Bian W, Brignall R, Cafarelli T, Campos-Laborie FJ, Charloteaux B, Choi D, Coté AG, Daley M, Deimling S, Desbuleux A, Dricot A, Gebbia M, Hardy MF, Kishore N, Knapp JJ, Kovács IA, Lemmens I, Mee MW, Mellor JC, Pollis C, Pons C, Richardson AD, Schlabach S, Teeking B, Yadav A, Babor M, Balcha D, Basha O, Bowman-Colin C, Chin SF, Choi SG, Colabella C, Coppin G, D’Amata C, De Ridder D, De Rouck S, Duran-Frigola M, Ennajdaoui H, Goebels F, Goehring L, Gopal A, Haddad G, Hatchi E, Helmy M, Jacob Y, Kassa Y, Landini S, Li R, van Lieshout N, MacWilliams A, Markey D, Paulson JN, Rangarajan S, Rasla J, Rayhan A, Rolland T, San-Miguel A, Shen Y, Sheykhkarimli D, Sheynkman GM, Simonovsky E, Taşan M, Tejeda A, Tropepe V, Twizere JC, Wang Y, Weatheritt RJ, Weile J, Xia Y, Yang X, Yeger-Lotem E, Zhong Q, Aloy P, Bader GD, De Las Rivas J, Gaudet S, Hao T, Rak J, Tavernier J, Hill DE, Vidal M, Roth FP, Calderwood MA. A reference map of the human binary protein interactome. Nature. 2020 Apr;580(7803):402-408. doi: 10.1038/s41586-020-2188-x. Epub 2020 Apr 8. PMID: 32296183; PMCID: PMC7169983.
• Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development