Cross-task interpretability through unified modeling reveals a universal shortcut bias in neoantigen prediction

This repository contains the implementation of ImmUni, the unified transformer-based framework introduced in our manuscript “Cross-task interpretability through unified modeling reveals a universal shortcut bias in neoantigen prediction.”

ImmUni models three key stages of neoantigen processing — peptide–HLA binding, antigen presentation, and T cell immunogenicity activation — within a single architectural framework. By enforcing a consistent representational scheme across tasks, ImmUni enables cross-task interpretability, allowing upstream, well-characterized processes to serve as internal comparators for diagnosing shortcut learning and dataset-driven biases in downstream prediction tasks.

An earlier version of this framework was presented under the name ImmuBPI at RECOMB 2024 (https://doi.org/10.1007/978-1-0716-3989-4_28). To maintain continuity and avoid fragmenting the codebase, we retain the original repository name ImmuBPI, while the current manuscript refers to the unified model as ImmUni. No functional changes were made to the implementation for this renaming; the code here corresponds to the version used in both the RECOMB presentation and the current manuscript.

⭐️ Background

There are three main steps in the neoantigen immune process, i.e., binding with MHCs, extracellular presentation, and induction of immunogenicity. While these stages operate on the same underlying biological entities, the available datasets differ greatly in quality, scale, and bias, making downstream immunogenicity prediction particularly vulnerable to shortcut learning. Existing models often perform well on isolated benchmarks yet fail to generalize, raising questions about whether they truly capture biological determinants or merely reflect dataset-specific artifacts.

ImmUni was designed to address this challenge by modeling all three stages within a unified transformer architecture. This enables cross-task interpretability: well-understood upstream processes serve as internal references that reveal when the downstream immunogenicity model deviates from mechanistic expectations. Through this unified perspective, ImmUni uncovers a pervasive shortcut bias driven by intra-HLA label imbalance and provides a principled framework for diagnosing and mitigating such biases in multi-step biological modeling.

⭐️ Setup

To use our model or reproduce our work, You need to install some dependency python package. We suggest you to do the dependency installtion in sperated conda environment or container.

To install the dependency, you could do it by

pip install -r requirements.txt

Note: to faithfully reproduce the training process, we recommend using cuda 11.8 (by conda install cudatoolkit=11.8). In our project setup checking, there is a slight difference in the training process and final performance with cuda 11.7 or different device driver version. It is due to the operator implementation and precision shifting in cuda. We recommend you follow this step to step up

1. Use nvidia GeForce RTX 3090 with driver version 525.125.06
2. conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.8 -c pytorch -c nvidia
3. pip install -r requirement.txt

For inference-only usage, there is no special version requirement on cuda or other driver lib and system lib. Just do the pip install command and have fun.

⭐️ Train

You can reproduce the training all of our models, by the all-in-one training bash script.

./train_bash.sh

We will introduce the training details of different models, including dataset preparation and training conifg in following subsections.

Dataset & Checkpoints

You can download the datasets and model checkpints from https://drive.google.com/drive/folders/1F0lrMcukm1E_aarWf_Xx1vA0SM-VRFlg?usp=sharing

After downloading, unzip the downloaded datasets/checkpoints in the ./dataset/ and ./models_saved/ folders

Training script

To train a single model, just run

python run_main.py -c <path to config file>

The config file contains the settings of model type, dataset type, training strategy and hyperparameter.

In this repo, we provide a example config file and config files for our best binding, presentation, immunogenicity models in ./configs. For all the parameter setting, like training learning rate, parts to be masked and so on, please check the commonts in ./configs/example_config.yaml.

⭐️ Inference

Scripts with specific functions are provided in ./scripts/:

• run_inference.py
• get_model_results.py
• plot_attention_map.py
• generate_embedding.py

All of these scripts have some same arugments to choose the model and dataset

-p  --path  <path to the model checkpoint saved dir(with at least .pt weights and config.yaml file)>
-e --ensemble [True/False] to use model ensemble if `True`, else use the single model for inference
-t --test_dataset <abbreviation of test set>

🧬 Inference on your own dataset

1.📄 Put your .csv file in dataset with format like below (e.g. dataset/example.csv):

HLA	peptide	label / immunogenicity (optional)
HLA-A*24:02	AYLEAIHHF	1
HLA-A*24:02	AYLEAIHKF	1
HLA-A*24:02	AYLEAIHNF	0

2.📎 Pick an abbreviation you like for your custom dataset and register it in dataset/dataset_register_table.csv

abbreviation	path
my_test	example.csv

3.🚀 Run inference! You can look up the results file model_score.csv at the folder named after your abbreviation saved under the model folder.

PYTHONPATH=./ python ./scripts/run_inference.py -p ./models_saved/ImmuBPI_I_Model -e True -t my_test

🧬 Inference in GUI

Run

python -m webui.webui

The open the localhost:7860 in your browser, you can get a intuitive interface, to help you choose the model, run the inference on a single data sample or your uploaded dataset.

🧬 Inference directly from command line

You can also directly use the command line as input for conveniency.

PYTHONPATH=./ python ./scripts/get_model_results.py -p ./models_saved/ImmuBPI_I_Model  -e True --peptide AYLEAIHNF AYLEAIHHF AYLEAIHKF AYLEAIHDF --HLA HLA-A24:02

🧬 Generate attention map

Run plot_attention_map.py to generate attention map. You can look up the results file layers_heatmap.png at the folder named after your abbreviation saved under the model folder.

PYTHONPATH=./ python ./scripts/plot_attention_map.py -p ./models_saved/ImmuBPI_B_Model/Fold_0  -t  TransPHLA_external

🧬 Generate model embedding

Run generate_embedding.py to generate model embedding.

PYTHONPATH=./ python ./scripts/generate_embedding.py -p ./models_saved/ImmuBPI_I_Model  -t my_test -e 

You can look up the results file transformer_feature.npy and transformer_feature.npy at the folder named after your abbreviation saved under the model folder.