SurvIPTB is a PyTorch implementation of attention-based models for estimating individual treatment benefit from survival outcomes. The models operate on paired control and treatment observations and produce a value between 0 and 1 for whether the treatment survival time exceeds the control survival time.
The data preparation utilities account for right-censored observations. For pairing cases in which the treatment-benefit label cannot be determined directly from the observed times, labels are initialized from group-specific survival estimates and can be optimized during model training.
This repository is research software. The current implementation does not establish clinical validity or suitability for clinical decision-making.
- support for right-censored survival outcomes;
- attention-based treatment-benefit estimation from paired observations;
- scaled dot-product attention with learned query and key projections;
- Gaussian distance-based attention with a learned temperature;
- utilities for survival pairing, feature scaling, and PyTorch dataset creation;
- per-split loss histories and best model states returned by training;
- a PyTorch-based implementation with configurable CPU or CUDA execution.
SurvIPTB requires Python 3.10 or later. From the repository root, install the package and its declared dependencies in editable mode:
python -m pip install -e .To run the demonstration notebook with its plotting dependency, install the optional examples extra:
python -m pip install -e ".[examples]"AIPTBSurv/
├── __init__.py
├── AIPTBSurv.py
├── AIPTBSurvDotProduct.py
├── AIPTBSurvGaussian.py
└── AttentionSurvDataset.py
Datagens/
└── ...
pyproject.toml
AIPTBSurv/ is the installable package. Datagens/ contains synthetic-data utilities used by the research workflows.
The shared torch.nn.Module base class implements attention-weighted forward inference, treatment-benefit loss, optional entropy regularization, and the training loop. Its compute_kernel method is intentionally implemented by subclasses, so the base class is not a standalone model variant.
A concrete model using scaled dot-product attention. It learns query and key projection matrices and applies fixed feature weights before projection. The constructor accepts the original covariate dimension, a training dataset, and the optional no_pi flag.
A concrete model whose attention logits are negative squared Euclidean distances scaled by a learned temperature. Its constructor accepts a training dataset and an optional initial temperature.
A PyTorch dataset for paired, standardized feature vectors and their treatment-benefit labels. It stores feature, value, case3, and case4 tensors. Integer indexing returns one observation; string indexing returns a complete tensor used by the model.
Builds paired training and validation data, fits a StandardScaler on the paired training features, and returns a dataset dictionary plus the fitted scaler. The returned dictionary contains train and val1; with the input contract used by the current implementation, it also contains val2 based on the companion *_test columns.
create_attention_surv_datasets accepts a pandas.DataFrame, row indices for the training and validation splits, the number of original covariates, and an unfitted sksurv.nonparametric.SurvivalFunctionEstimator.
The observed record uses the following columns:
| Column | Meaning in the implementation |
|---|---|
a |
Treatment assignment. 0 is the control group; any nonzero value is treated as a treatment group. |
y |
Observed survival time. |
c |
Event indicator: values greater than 0 are events, and 0 denotes right censoring. |
x1 ... xN |
Covariates, where N equals the dim argument. |
The current helper also expects a_test, y_test, c_test, x1_test ... xN_test, and te_test. These companion columns are removed while constructing the paired train and val1 splits. They are rearranged into paired control-treatment records for val2; te_test itself is not used to form the model label. Because the implementation drops these columns unconditionally during survival pairing, they must be present even when only train and val1 are of interest.
For model input, each control row is paired with each treatment row within a split. The resulting feature order is x1, z1, x2, z2, ..., where x* contains control covariates and z* contains treatment covariates. The feature width seen by an attention kernel is therefore 2 * dim.
The helper fits the supplied survival estimator separately for each value of a. Training and validation subsets must consequently contain usable observations for every treatment group needed in that subset.
The following example creates one-covariate survival data, includes both observed events and right-censored observations, constructs all dataset splits, trains the dot-product model for one epoch on CPU, and obtains a treatment-benefit probability.
import pandas as pd
import torch
from sksurv.nonparametric import SurvivalFunctionEstimator
from AIPTBSurv import AIPTBSurvDotProduct, create_attention_surv_datasets
df = pd.DataFrame(
{
"a": [0, 0, 0, 0, 1, 1, 1, 1],
"y": [2.0, 4.0, 6.0, 8.0, 3.0, 5.0, 7.0, 9.0],
"c": [1, 0, 1, 1, 1, 0, 1, 1],
"x1": [0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9],
"a_test": [1, 1, 1, 1, 0, 0, 0, 0],
"y_test": [3.0, 5.0, 7.0, 9.0, 2.0, 4.0, 6.0, 8.0],
"c_test": [1, 1, 1, 1, 1, 1, 1, 1],
"x1_test": [0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9],
"te_test": [1.0, 1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0],
}
)
datasets, scaler = create_attention_surv_datasets(
df=df,
train_idx=[0, 1, 2, 4, 5, 6],
val_idx=[3, 7],
dim=1,
surv_model=SurvivalFunctionEstimator(),
)
model = AIPTBSurvDotProduct(dim=1, train_dataset=datasets["train"])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
history, best_states = model.fit(
datasets=datasets,
optimizer=optimizer,
num_epochs=1,
batch_size=8,
device="cpu",
)
model.eval()
with torch.no_grad():
benefit_probability = model(datasets["val1"]["feature"])
print(history["val1"][-1])
print(benefit_probability)The model has no separate public predict method. Inference uses the standard PyTorch module call shown above. The returned tensor is clamped to [0, 1]; for a paired query, it represents the model output for positive treatment benefit, defined by the data preparation code as treatment survival time greater than control survival time.
In this project, a row is right-censored when c == 0: the event has not been observed by time y, so y is the observed follow-up time rather than a known event time. A value c > 0 means that the event was observed. The same convention applies to c_test when companion records are provided.
During pairing, records for which both members are censored are removed. Two ambiguous one-censored pairing configurations are marked by the dataset tensors case3 and case4. Their initial te_prob labels are calculated from group-specific survival curves. In the base training behavior these labels are represented by trainable values; AIPTBSurvDotProduct(no_pi=True) keeps their initial values fixed instead.
Pass the event indicators unchanged to create_attention_surv_datasets; do not invert them to a 1 = censored convention.
All concrete models inherit fit from AIPTBSurv. The method receives datasets and a caller-created optimizer; its optional parameters are num_epochs=10, batch_size=10, an automatically selected device, metric_fns=None, and reg_entropy=0. datasets is a dictionary of named AttentionSurvDataset objects and must include train. The default device is CUDA when available and CPU otherwise.
history maps each split name to one dictionary per epoch, including epoch and Loss plus any requested metrics. Optional metrics are evaluated only on pairs not marked as case3 or case4. best_states maps each split to the state dictionary with its lowest recorded loss. The training dataset must be large enough for at least one batch to contain non-empty query and key subsets.
Open the demonstration notebook for a reproducible walkthrough of censored survival data generation, model training, loss convergence, predictions, evaluation, and inline visualizations.
Citation information will be added after publication details are finalized.