Survival and causal tasks, served by SurvivalPFN and CausalPFN

.github/workflows/docker_image.yml

@@ -51,7 +51,7 @@ jobs:
         # (a missing wheel or broken dependency fails here, on CPU).
         run: |
           docker run --rm --entrypoint python tabctx-under-test -c \
-            "import torch, tabicl, tabpfn, tabdpt, tabctx; print('torch', torch.__version__)"
+            "import torch, tabicl, tabpfn, tabdpt, causalpfn, survivalpfn, tabctx; print('torch', torch.__version__)"
 
       - name: Locked versions and a C compiler for Triton
         # Every package but torch's own stack is the uv.lock version, and
@@ -63,7 +63,8 @@ jobs:
             cc --version | head -1
             cd /opt/tabctx
             uv export --frozen --no-dev --no-hashes --no-emit-project \
-              --extra standalone --extra tabicl --extra tabpfn --extra tabdpt -q -o /tmp/locked.txt
+              --extra standalone --extra tabicl --extra tabpfn --extra tabdpt \
+              --extra causalpfn --group survivalpfn -q -o /tmp/locked.txt
             python - <<"PY"
           import importlib.metadata as md, os, re, sysconfig
           assert os.path.exists(os.path.join(sysconfig.get_paths()["include"], "Python.h"))

CHANGELOG.md

@@ -1,6 +1,20 @@
 # Changelog
 
 ## [Unreleased]
+### Added
+- **Two new tasks, survival and causal, with SurvivalPFN and CausalPFN** (`TABCTX_BACKEND=...,survivalpfn,causalpfn`). Both models are about 20M parameters (75 MB of weights) and fit beside the others on one GPU.
+  - **Survival** (`survivalpfn-v0.1`, [SurvivalPFN](https://github.com/rgklab/SurvivalPFN), Apache-2.0): time to an event with right censoring. A fit sends each row's time as `train_y` and its event flag as `train_event` (or a file's `event_column`); a predict returns each row's median event time, `quantiles` of the event time, with `times` the probability of no event by each (`survival`, `survival_times`), and with `rmst_horizon` the restricted mean survival time up to it (`rmst`). Nothing is estimated past the longest training time (the fit response's `summary.max_time`; that row is always kept in the context): those values, `rmst` included, are `null` (with heavy censoring the model put near-zero survival there, against a true 0.74).
+  - **Causal** (`causalpfn-v0`, [CausalPFN](https://github.com/vdblm/CausalPFN), Apache-2.0): a binary treatment's effect from observational data. A fit sends each row's outcome and its treatment flag (`train_treatment` / `treatment_column`); a predict returns each row's conditional average treatment effect, with `potential_outcomes` the expected outcome if not treated and if treated, and `quantiles` of a treated minus an untreated outcome (CausalPFN's interval method). The quantiles are calibrated as CausalPFN's `calibrate=True` does (a temperature by 3-fold cross-validation, computed exactly rather than from samples, at a context's first predict that asks for quantiles, so fits that never ask pay nothing): on simulated data a 90% range went from 5x the outcome noise's spread to its spread. They read each row's two distributions at fixed levels instead of sampling them, so they repeat exactly. Missing values are filled with the training medians. The effect matches CausalPFN's own estimator when the table fits the context; for a larger table, CausalPFN picks a context per group of test rows, which made a row's answer depend on the rest of its request, so tabctx uses one subsampled context instead (below).
+  - **One context per fit.** Each model keeps at most 4,096 training rows as its context (`TABCTX_SURVIVALPFN_CONTEXT_SIZE`, `TABCTX_CAUSALPFN_CONTEXT_SIZE`; a larger table is subsampled at fit, seeded, stratified on both targets: the share of events or of each arm, and the spread of times or outcomes within it), and test rows attend only to it, so a row's answer is the same alone, merged with other requests or split into chunks. The models take any number of rows; a larger context costs predict time (roughly with its square: 1,000 survival rows on an A40 took 0.12 s at 4,096 and 50 s at 100,000) and, on four of SurvivalPFN's benchmark datasets larger than 4,096 rows (up to 36,695: hdfail, prostateSurvival, Dialysis, dataDIVAT1), added at most 0.01 C-index.
+  - **Checked against SurvivalPFN's published benchmark** (`tests/integration/test_survivalpfn_reference.py`, `uv sync --group survivalpfn --group survivalpfn-reference`): on the authors' own splits of six small public datasets (veteran, mgus, whas500, stagec, Melanoma, zinc), SurvivalPFN's medians reproduce the published C-index, tabctx's medians equal them wherever given, and ranking rows by `rmst` up to the longest training time scores within 0.01 of it (0.015 tolerance in the test). With heavy censoring many medians are past the follow-up (`null`; zinc, 81% censored: C-index 0.621 from the medians, 0.784 from `rmst`), so `rmst` is the way to rank rows (a higher `rmst` is a lower risk). SurvivalPFN is not calibrated: its authors found CausalPFN's temperature calibration makes it worse.
+  - **What a predict returns** (`tabctx.outputs`): beyond the predictions, `quantiles`, `times`, `rmst_horizon` and `potential_outcomes`, each for the tasks whose models give it, checked before any work (a 422 otherwise, and for `return_proba` on a survival or causal context). One validated value carries them through the engine, devices, worker processes and the coalescer, so adding an output is a field there and its backends.
+  - A survival or causal fit's response has a `summary`: `max_time`, `events` and `censored`, or `treated` and `control`.
+  - Each fit's task picks its model: a backend lists its `tasks`, `GET /v1/models` shows them with each model's `outputs` (now with a response schema), and a fit that names no model goes to the default if it fits the task, else the first that does. A model that doesn't fit the task is a 422 naming one that does. An event or treatment column implies its task.
+  - Both networks load once per device, as tensors only (`torch.load(weights_only=True)`; the models' own loaders unpickle arbitrary objects), from pinned Hugging Face revisions; the spill tier stores a context without its network.
+  - On a CPU, on simulated data (1,500 training rows, 500 held out): SurvivalPFN's medians had a C-index of 0.922 against the observed times (0.942 for the true risk itself); CausalPFN's effects correlated 0.998 with the true ones. Not yet run on a GPU or on the papers' benchmarks.
+  - `pip install "tabctx[causalpfn]"`. SurvivalPFN is not on PyPI, so it is a uv dependency group (`uv sync --group survivalpfn`), pinned to a commit; the image includes both.
+  - Event and treatment columns in files may be 0/1, true/false or yes/no; with no target named, the target is the last column that isn't the flag.
+  - The client: `fit(..., event=)` / `treatment=` (sending the task they imply, so a server without these tasks refuses the fit rather than fitting a regression), `fit_file` and `fit_uploaded` with `event_column` / `treatment_column`, `predict(..., times=, rmst_horizon=, potential_outcomes=)`, and `tabular_models(task)` / `default_model(task)`; the scikit-learn estimators and the MCP tools only pick models that fit their task.
 
 ## [0.20.0] - 2026-11-10
 ### Added

Dockerfile

@@ -1,4 +1,5 @@
-# The standalone serving image: tabctx + tabicl + tabpfn + tabdpt + torch, serving
+# The standalone serving image: tabctx + tabicl + tabpfn + tabdpt + causalpfn +
+# survivalpfn + torch, serving
 # `python -m tabctx.serve` (serve/standalone.py) -- no Ray. Built and
 # pushed by .github/workflows/docker_image.yml on each release tag, and
 # smoke-tested there against the fake backend (no GPU in CI).
@@ -52,7 +53,8 @@ COPY src ./src
 # resolution for both, so the extras' unpinned `torch` and the pin cannot
 # diverge. The project itself goes in last, without re-resolving.
 RUN uv export --frozen --no-dev --no-hashes --no-emit-project \
-        --extra standalone --extra tabicl --extra tabpfn --extra tabdpt -o /tmp/locked.txt \
+        --extra standalone --extra tabicl --extra tabpfn --extra tabdpt \
+        --extra causalpfn --group survivalpfn -o /tmp/locked.txt \
     && grep -v -E '^(torch|triton|nvidia-)' /tmp/locked.txt > /tmp/requirements.txt \
     && uv pip install --system --no-cache \
         "torch==${TORCH_VERSION}" -r /tmp/requirements.txt \

README.md

@@ -54,10 +54,11 @@ own GPUs.
 - **Many tables, one GPU.** Contexts share a byte-budgeted cache. The
   coldest one is evicted (or spilled to disk and restored on next use),
   and a fit that would run the GPU out of memory is refused *before* it runs.
-- **Three models, one endpoint.** `tabicl-v2`, `tabpfn-3.5` and
-  `tabdpt-1.3` behind `GET /v1/models`; pick one per request. Two of them
-  are licensed for commercial use. Adding a model is one backend class
-  and a few lines of wiring.
+- **Five models, one endpoint.** `tabicl-v2`, `tabpfn-3.5` and
+  `tabdpt-1.3` for classification and regression, `survivalpfn-v0.1` for
+  time to an event and `causalpfn-v0` for treatment effects, behind
+  `GET /v1/models`; pick one per request. Adding a model is one backend
+  class and a few lines of wiring.
 - **Prediction intervals.** Regression predicts a distribution, so you can
   ask for its quantiles (say a 90% interval) at no extra cost:
   `reg.predict_quantiles(X, [0.05, 0.95])`.
@@ -95,7 +96,8 @@ clf.score(X_test, y_test)          # 0.77, no training loop, no tuning
 clf.predict_proba(X_test[:5])      # reuses the cached context: ~50 ms on a laptop
 ```
 
-On a GPU server, use the image (TabICL, TabPFN and TabDPT included,
+On a GPU server, use the image (TabICL, TabPFN, TabDPT, SurvivalPFN and
+CausalPFN included,
 weights downloaded when the server starts):
 
 ```bash
@@ -162,9 +164,12 @@ users and tables, on GPUs you control.
 | `tabicl-v2` | [TabICLv2](https://github.com/soda-inria/tabicl) (Inria) | BSD-3-Clause: commercial use OK |
 | `tabpfn-3.5` | [TabPFN-3.5](https://github.com/PriorLabs/TabPFN) (Prior Labs) | Non-commercial; production use needs a [license from Prior Labs](https://priorlabs.ai/pricing) |
 | `tabdpt-1.3` | [TabDPT](https://github.com/layer6ai-labs/TabDPT-inference) (Layer 6) | Apache-2.0: commercial use OK |
+| `survivalpfn-v0.1` | [SurvivalPFN](https://github.com/rgklab/SurvivalPFN) (University of Toronto): time to an event, with censoring | Apache-2.0: commercial use OK |
+| `causalpfn-v0` | [CausalPFN](https://github.com/vdblm/CausalPFN) (University of Toronto, Layer 6): treatment effects from observational data | Apache-2.0: commercial use OK |
 
 `TABCTX_BACKEND=tabicl,tabdpt` serves both commercially usable models from
-one cache. TabPFN weights download after you accept Prior Labs' terms (set
+one cache; add `survivalpfn,causalpfn` for the survival and causal tasks
+([Python guide](https://vectorinstitute.github.io/tabctx/getting-started/python/#survival-and-treatment-effects)). TabPFN weights download after you accept Prior Labs' terms (set
 `TABPFN_TOKEN`). tabctx itself is Apache 2.0; each model's weights keep
 their own license.
 

deploy/helm/tabctx/values.yaml

@@ -14,7 +14,10 @@ imagePullSecrets: []
 # Models to serve (TABCTX_BACKEND); the first is the default. tabicl and
 # tabdpt are licensed for commercial use (tabdpt is in images from
 # 0.19.0); tabpfn needs a token (below) and is licensed for non-commercial
-# use.
+# use. survivalpfn (time to an event) and causalpfn (treatment effects)
+# add those tasks, in images from 0.21.0. List a classification and
+# regression model first: clients older than 0.21 use the default model
+# for every fit.
 models: [tabicl]
 
 # Must stay 1 with the chart's volumes: they are ReadWriteOnce, and a

docs/contributing/adding-a-model.md

@@ -23,7 +23,27 @@ models: Mitra ([#102](https://github.com/VectorInstitute/tabctx/issues/102)).
   at predict. The model can treat every column as numeric.
 - **`y`**: the labels. For classification, any values (tabctx returns
   them as strings); for regression, numbers.
-- **`task`**: `"classification"` or `"regression"`.
+- **`task`**: `"classification"` or `"regression"`, or `"survival"` or
+  `"causal"` for a model that serves those (below).
+
+### Which tasks
+
+A backend serves classification and regression unless it says otherwise
+in a `tasks` attribute (e.g. `tasks = frozenset({"survival"})`); the
+engine sends it only those fits, and a fit that names no model goes to
+the first one that serves its task. A survival or causal fit's `y` has
+two columns (`tabctx.tasks.split_targets`): time and event flag, or
+outcome and treatment flag.
+
+### Outputs beyond the predictions
+
+`tabctx/outputs.py` lists what a predict can ask for besides its
+predictions (`quantiles`, `times`, `rmst_horizon`, `potential_outcomes`)
+and which tasks each is for. A backend gives an output by taking a
+keyword argument of that name in `predict` and filling the matching
+`PredictOutcome` field; it is passed only the outputs asked for, and a
+request for one its `predict` doesn't take is a 422 before any work. `backends/survivalpfn.py`
+and `backends/causalpfn.py` are the examples.
 
 ### The rules that matter
 
@@ -65,9 +85,9 @@ version needs none of them:
 2. **The server**: in `src/tabctx/serve/factory.py`, add its kind to
    `BackendKind` and `BACKEND_KINDS`, and a branch for it in
    `_build_real_backend`.
-3. **The client**: add its id to `KNOWN_MODELS` in
+3. **The client**: add its id and tasks to `KNOWN_MODELS` in
    `src/tabctx/client.py`, so the client finds it among a gateway's
-   models.
+   models (a gateway's listing has no `tasks`).
 4. **Packaging**: an extra in `pyproject.toml` (`tabctx[<model>]`), added
    to the `Dockerfile`'s export and to the image check in
    `.github/workflows/docker_image.yml`, and a row in the README's models

docs/getting-started/http.md

@@ -48,6 +48,31 @@ can also ask for `quantiles` of the predicted distribution, e.g.
 gets `quantiles` (one row per test row) and `quantile_levels` back. When the training table
 was a file, a test file's columns are matched by name.
 
+A survival fit sends each row's time as `train_y` and its event flag as
+`train_event` (1: the event happened then; 0: censored), or names the
+columns of a file (`target`, `event_column`); its predicts return each
+row's median event time and, with `times`, the probability of no event
+by each time (`survival`), and with `rmst_horizon` the expected
+event-free time up to it (`rmst`). Past the longest training time (the
+fit response's `summary.max_time`) nothing is estimated, so to rank rows
+by risk ask for `rmst` with `rmst_horizon` at that time: it is always
+defined (a higher `rmst` is a lower risk). A causal fit sends each row's outcome
+and its treatment flag (`train_treatment`, or `treatment_column`); its
+predicts return each row's treatment effect, and with
+`potential_outcomes: true` the expected outcome if not treated and if
+treated. Predicting a survival or causal context with `return_proba` is a
+422: there are no class probabilities. `GET /v1/models` lists each model's
+`tasks` and `outputs`. A survival or causal fit's response has a
+`summary`: `max_time` (outputs past it are `null`), `events` and
+`censored`, or `treated` and `control`.
+
+```bash
+curl localhost:8000/v1/tabctx/fit -F dataset_id=retention -F target=days \
+  -F event_column=cancelled -F file=@policies.csv
+curl localhost:8000/v1/tabctx/predict -F dataset_id=retention \
+  -F times=90,365 -F file=@current.csv
+```
+
 ## Large tables
 
 A table file sent with the fit or predict is fine up to

docs/getting-started/python.md

@@ -72,6 +72,61 @@ client.models(), client.limits()
 `TABCTX_API_KEY`. The key is sent as a bearer token, never appears in
 the client's `repr`, and is left out when the client is pickled.
 
+## Survival and treatment effects
+
+Two more tasks, each served by its own model (`GET /v1/models` lists
+every model's `tasks`; a fit that names no model goes to one that fits
+its task).
+
+**Time to an event** (`survivalpfn-v0.1`): each training row has a time
+and an event flag, 1 if the event happened at that time and 0 if the row
+was censored then (still event-free when follow-up ended).
+
+```python
+client.fit(X_train, days, event=cancelled, dataset_id="retention")
+r = client.predict("retention", X_now, times=[90, 365], quantiles=[0.25, 0.75])
+r.predictions   # each row's median time to the event
+r.survival      # per row, the probability of no event by day 90 and day 365
+client.predict("retention", X_now, rmst_horizon=365).rmst  # expected event-free days in the year
+```
+
+**Treatment effects** (`causalpfn-v0`): each training row has an outcome
+and a treatment flag (1 treated, 0 control), from observational data.
+
+```python
+client.fit(X_train, spend, treatment=got_offer, dataset_id="offer")
+r = client.predict("offer", X_now, potential_outcomes=True, quantiles=[0.05, 0.95])
+r.predictions          # per row: the expected outcome if treated minus if not
+r.potential_outcomes   # per row: [expected outcome if not treated, if treated]
+sum(r.predictions) / len(r.predictions)   # the average effect over these rows
+```
+
+Nothing is estimated past the longest time in the training data: a
+median, quantile or survival probability out there is `None`. To rank
+rows by risk, ask for `rmst` with `rmst_horizon` at that longest time
+(`days.max()`, or the time column's max for `fit_file`): it is always
+defined, while with heavy censoring many medians are `None`. A higher
+`rmst` is a longer expected event-free time, so a lower risk. Each model
+uses at most 4,096 training rows as its context (a larger table is
+subsampled, keeping the share of events or of each arm and the spread of
+times or outcomes;
+`TABCTX_SURVIVALPFN_CONTEXT_SIZE`, `TABCTX_CAUSALPFN_CONTEXT_SIZE`), and
+every test row is predicted against that same context, so a row's answer
+never depends on the other rows sent with it.
+
+From a file, name the columns: `fit_file(path, target="days",
+event_column="cancelled")` or `treatment_column=`; a flag column may be
+0/1, true/false or yes/no. What the models
+assume: the survival model, that censoring is independent of the event
+given the features; the causal model, that every confounder of treatment
+and outcome is a feature, and that every kind of row could have been
+treated or not. No data can confirm the causal one. The causal model's
+`quantiles` are of a treated outcome minus an untreated one for the row,
+calibrated (on the first predict that asks for them) to cover held-out
+outcomes at their stated rates: the
+spread of an individual's effect, outcome noise included, not a
+confidence interval for the average effect.
+
 ## Errors
 
 Server errors come back as tabctx exceptions, importable from `tabctx`.

docs/getting-started/quickstart.md

@@ -37,8 +37,8 @@ categories and missing values go as they are.
 
 ## 3. On a GPU server
 
-The container image includes TabICL, TabPFN and TabDPT; each
-model's weights download when the server starts.
+The container image includes TabICL, TabPFN, TabDPT, SurvivalPFN and
+CausalPFN; each model's weights download when the server starts.
 
 ```bash
 docker run --gpus all -p 8000:8000 -e TABCTX_BACKEND=tabicl ghcr.io/vectorinstitute/tabctx