CAST (Cross-subsample Amortised Site Tilting) steers a co-folding model's diffusion sampler toward a named antibody binding site rather than a pose. Training-free, no reference structure, frozen backbone.
The problem is site coverage. Co-folding models often dock the antibody on the wrong epitope, and more seeds deepen the mode the model already prefers instead of relocating it. A verifier cannot select an epitope the sampler never proposed. CAST adds a site-measurable reward from a candidate epitope map, carries the candidate site and the alignment-subsample index in the particle state so K conditionings share one pool, and tilts toward the site while leaving the pose to the base model.
This repository contains the running pipeline: preprocessing, input wiring, the methods, and the job templates to run them end to end.
- Linux, NVIDIA GPU.
- Python 3.11 (backends pin
>=3.11,<3.12), CUDA 12.6.
Clone the repository if you want a local checkout:
git clone https://github.com/.../CAST-AbAg.git
cd CAST-AbAgThen install with either Poetry:
poetry install # core
poetry install --with guidance # + torch, the differentiable reward path
poetry install --with cdr # + anarci, for the CDR mapor pip:
pip install .poetry install sets up this package only. It installs numpy, pandas, pyarrow,
omegaconf, biopython, biotite and tqdm. Two optional groups add the rest, torch in guidance and anarci in cdr. It deliberately
does not install a co-folding backend: nothing here does import protenix, the backends are
shelled out to in their own virtualenvs. Five things must be provisioned separately, and
./setup_backends.sh covers every one except the MSA database. Pick what you need, re-run it
freely, anything already on disk is skipped, and it prints the env-file lines for what it built.
It installs both backends and downloads OpenDDE's checkpoint, but leaves the Protenix weights to
you, for the reason in the Model weights row below. --cdr resolves HMMER only, since
ANARCI is already a poetry dependency.
./setup_backends.sh --protenix # Protenix v2 venv
./setup_backends.sh --opendde # OpenDDE venv + abag checkpoint
./setup_backends.sh --all --prefix /data/tts
./setup_backends.sh --all --dry-run # show what it would do| Component | How | Pointed at by |
|---|---|---|
| Co-folding backend | install Protenix v2 or OpenDDE in its own venv | PROTENIX_COMPONENT_PATH, or OPENDDE_ROOT_DIR + OPENDDE_BIN |
| Model weights | Protenix v2 is fetched by theprotenix package from protenix.tos-cn-beijing.volces.com/checkpoint/protenix-v2.pt, or pre-placed by hand when that origin is blocked. For a PRIVATE Protenix-v2 checkpoint, put it at $PROTENIX_ROOT_DIR/checkpoint/protenix-v2.pt, by baking it into your image or placing it by hand, because a checkpoint already on disk wins before any fetch is tried. OpenDDE weights are opendde_abag.pt from huggingface.co/aurekaresearch/OpenDDE |
PROTENIX_ROOT_DIR, OPENDDE_ROOT_DIR. The job template logs which checkpoint it resolved before taking a GPU. models/protenix_provision.py has the full artefact list, the manual-placement layout, a third-party mirror with its digest, and the caveats on using it |
| MSA database | a pre-built ColabFold-style DB, or generate per target withgenerate_msa |
MSA_DB |
| DockQ | install in a scoring venv, kept separate from the sampler's | SCORE_PY, SCORE_DOCKQ |
| ANARCI + HMMER | only for the CDR map. ANARCI comes frompoetry install --with cdr. HMMER is a C program and must be the real binary, not a wrapper script, which rules out the PyPI hmmer package's console scripts: HMMER writes argv[0] into its output header, a wrapper puts its own name there, Biopython then matches no program and leaves every alignment coordinate None, and ANARCI dies comparing them. Use a conda/apt/source HMMER, or point at the real executables the wheel bundles under site-packages/hmmer/data/bin |
--hmmer-path on token_maps cdr-map |
A pool run is configured by an env file passed with --env-file, sourced by the job template so
every export reaches the worker. The variables it must set:
| Variable | Meaning |
|---|---|
CODE, STATS_DIR |
repo root, and$CODE/jobs for the helper scripts |
BACKEND, TTS_CFG_SUFFIX |
which backend, and the config suffix that selects its variants |
RUN_DIR |
run outputs and the shared queue |
INPUTS_WIRED, GT_INDEX, MSA_DB |
prepared inputs, the target index, the MSA database |
EPI_HYBK12, PROTENIX_CDR_MAP |
candidate epitope map and paratope token map |
N_SAMPLES, SUBSAMPLE_POOL, N_SEEDS |
pool size N and conditioning count K, which must match |
UMBRELLA_HYBK12, SCALE |
arrival radii and guidance scale |
TARGETS, ARMS_ONLY, SKIP_CTRL |
what to run.SKIP_CTRL defaults to 1, so the matched control is opt-in with SKIP_CTRL=0 |
SCORE_PY, SCORE_DOCKQ, SCORE_JOBS |
scoring |
LAYERNORM_TYPE |
which Protenix layer-norm kernel to use. Every template defaults totorch, which is what all the measured arms ran. tts-b0b-pool.yaml also accepts fast_layernorm, the faster JIT-compiled CUDA kernel, which needs nvcc and ninja and is not bit-identical to torch in bf16 |
Cluster-specific paths default from APP_ROOT and TTS_ROOT, so nothing needs editing to run
elsewhere. The job templates also carry placeholders to fill in: <YOUR-NAMESPACE>,
<YOUR-PVC-CLAIM>, <YOUR-GPU-PRODUCT-LABEL>, <YOUR-PROTENIX-IMAGE>, <YOUR-PROTENIX-VENV>,
<YOUR-PROTENIX-CACHE> and <YOUR-VENV>.
The full pipeline, in order, for the paper's configuration: CAST steering to the hybK12
epitope map. Set PYTHONPATH to the repo root, or install the package.
Step 4 builds the epitope map, and is OPTIONAL if you already have one. The map is a
small JSON that travels with a run, so it is normally built once per dataset and reused. Nothing
builds it implicitly: the job templates only ever consume a map, so point them at the existing
file and go straight from step 3 to step 5. Which variable depends on the path: step 5a
(run_ablation) reads $PROTENIX_EPITOPE_MAP, while step 5b (the pool) takes EPI_HYBK12 and
exports PROTENIX_EPITOPE_MAP itself, per arm.
A reused map must pass two gates, both before any GPU is taken:
- its
_builderstamp must equal the config'sepitope_map_builder, which ishybrid_modes_tilesform5_cast.yaml. A map built bytiles-ensembleis stampedhybrid_modes_ensemble_tilesand is refused, which is what stops two different maps being reported under one label. dataset-checkrequires it to cover every target in$TARGETS, so a map built for one dataset cannot be silently reused on another.
Check both without launching anything:
PYTHONPATH=$PWD python jobs/tts_job_lib.py map-stamp --map $RUN/hybk12.json \
--source predicted --builder hybrid_modes_tiles
# exit 0 and "[map-stamp] ... OK ... over N targets", or exit 2 naming the mismatchIf you do need to build it, the ordering matters in one non-obvious place. The deployable map is derived from an unsteered pool, so the pool inside step 4 runs before the map is built and the steered run in step 5 cannot come first. That pool is not wasted either way, since it is also the Best-of-N control the steered arms are compared against.
Each map is one step here but several commands: step 3 builds cdr_map.json then
cdr_tokenmap.json, and step 4 builds modes.json and tiles.json then hybk12.json.
1 prepare_dataset 2 generate_msa (own MSA DB? same command, it only wires; | own wired inputs? skip it with REUSE_WIRED=1) ├─ 3 cdr-map ─> cdr ─> $PROTENIX_CDR_MAP ───────────────────┐ │ │ └─ 4 unsteered pool ─┬─ epitope(pool_modes) ─┐ │ (optional) └─ tiles ───────────────┴─> hybrid │ │ │ $PROTENIX_EPITOPE_MAP │ │ │ 5 steered run <─────────────────────────────┴─────────┘ │ 6 evaluate ─> 7 compute_table
Every block below shares this preamble.
DATASET=cfa50
RUN=runs/$DATASETStructures to backend-native inputs plus a target index.
python -m experiments.scripts.prepare_dataset \
--dataset $DATASET \
--backend protenix2 --runs-root runs
# -> $RUN/inputs/*.json and $RUN/target_index.csv (the latter is $GT_INDEX for scoring)
# cfa50 needs no --dataset-args, it reads the shipped data/cfa50_targets tree. The other
# loaders do, e.g. --dataset-args '{"bad_cases_dir": "'$DATA_ROOT'/bad_cases"}'--dataset takes one of cfa50, cofoldarena, foldbench, collected_bad_cases, opendde40.
Two jobs in one command: fill the DB from the server, then attach each input's per-chain MSA from the DB.
python -m experiments.scripts.generate_msa \
--inputs-dir $RUN/inputs --db-dir $DATA_ROOT/msa_db \
--wire-out $RUN/inputs_wired --backend protenix2ALREADY HAVE THE MSAs? Run the same command. Generation is per target and skipped when the DB
already holds it, so with a complete DB it contacts no server and only wires. Already have WIRED
inputs? Skip this step and point the run at them with REUSE_WIRED=1 and INPUTS_WIRED=<dir>,
otherwise it REWRITES every wired input the DB covers.
Stage (a) needs ANARCI, stage (b) does not.
python -m experiments.scripts.token_maps cdr-map \
--wired-dir $RUN/inputs_wired --out-json $RUN/cdr_map.json
python -m experiments.scripts.token_maps cdr \
--run-dir $RUN --cdr-map $RUN/cdr_map.json --out-json $RUN/cdr_tokenmap.json
export PROTENIX_CDR_MAP=$RUN/cdr_tokenmap.jsonOPTIONAL. Already have one? Then just export PROTENIX_EPITOPE_MAP=<existing>/hybk12.json and skip
to step 5.
4a, the unsteered pool the map is derived from, which is also the Best-of-N control.
CODE=$PWD TEMPLATE=jobs/tts-b0b-pool.yaml RUN_DIR=$RUN/bon \
bash jobs/submit_steer_pool.sh --gpus 4 --local --env-file env.sh
POOL=$RUN/bon/predictions4b, epitope candidates from that pool's own binding-site modes. --scorer defaults to pool_modes,
which is what hybK12 needs; pass --scorer discotope3 with --pool-dir pointing at the CSV dir for
that predictor instead.
python -m experiments.scripts.token_maps epitope \
--n-patches 3 --run-dir $RUN --pool-dir $POOL \
--wired-dir $RUN/inputs_wired --out-json $RUN/modes.json4c, exhaustive surface tiling of the antigen, 12 tiles per target.
python -m experiments.scripts.token_maps tiles \
--pool-dir $POOL --wired-dir $RUN/inputs_wired \
--k-want 12 --out-json $RUN/tiles.json4d, hybK12 is pool modes first, then tiles ranked by mode support, to a 15-candidate budget.
python -m experiments.scripts.token_maps hybrid \
--modes-json $RUN/modes.json --tiles-json $RUN/tiles.json \
--out-json $RUN/hybk12.json
export PROTENIX_EPITOPE_MAP=$RUN/hybk12.json # EPI_HYBK12 in the job templates5a, one CAST arm directly. Paths come from the config unless overridden, so --set them to keep
each unit's artefacts together, exactly as the job template does.
python -m experiments.scripts.run_ablation \
--config experiments/configs/diffusion_side/m5_cast.yaml \
--set model.run.msa_subsample.source_inputs_dir=$RUN/inputs_wired \
paths.inputs_dir=$OUT/inputs paths.predictions_dir=$OUT/pred \
paths.results_dir=$OUT/res5b, or a pool of workers, one per GPU. The env file exports CODE, STATS_DIR, GT_INDEX,
EPI_HYBK12 and PROTENIX_CDR_MAP; see the environment-file table above.
With only EPI_HYBK12 set, the queue is just the hybk12 arm: every other map arm is gated on its
own EPI_* variable, the oracle arms on WANT_GT, and the control on SKIP_CTRL.
CODE=$PWD TEMPLATE=jobs/tts-cast-pool.yaml RUN_DIR=$RUN/cast \
bash jobs/submit_steer_pool.sh --gpus 4 --local --env-file env.sh
# add --dry-run to stop after preflight, before any GPU workSKIP_CTRL defaults to 1, so NO control is queued. A ctrl unit depends only on
(target, seed, n_samples, backend), so it is built once and reused; set SKIP_CTRL=0 to run one
here when you do not already have a matching one. WANT_GT=1 adds the oracle ceiling arms, and then
EPI_TRUE is required.
REQUIRED after 5a, which writes structures and leaves the parquet's dockq column empty. After 5b
the pool has already run this command per unit as score_unit, so the only thing left for it to add
is --summary, the rank and oracle rollup that the pool never passes. Pool-written rows carry
--method-label inflight, so a later pass here is distinguishable rather than duplicated.
--from-per-sample reads the run's per_sample.parquet, and without it --predictions-dir is
required instead and the call will find nothing.
python -m experiments.scripts.evaluate --from-per-sample \
--target-index $RUN/target_index.csv \
--results-dir $OUT/res --backend protenix2 \
--metrics dockq --dockq-bin $SCORE_DOCKQ --summaryMeasured FLOPs and wall-clock across runs. Pass the Best-of-N pool too and --ref-arm to read each
steered arm as a ratio of it.
python -m experiments.scripts.compute_table $RUN/cast $RUN/bon --ref-arm b0bThe pool templates run steps 1, 2 and the pool in 4a themselves from $DATASET / $DATASET_ARGS, so a pool
launch needs only the two maps from steps 3 and 4 to exist and be exported. --dry-run stops
after the preflight, which is what checks that $PROTENIX_EPITOPE_MAP is present, carries a
_builder stamp matching epitope_map_builder, and that the per-cell budget is satisfiable.
| Arm | Config | What it is |
|---|---|---|
| CAST | diffusion_side/m5_cast.yaml |
this work |
| Best-of-N | baseline/b0b.yaml |
unsteered control |
| FK steering | diffusion_side/m1_fk.yaml |
same reward, no gradient |
| TDS | diffusion_side/m2_tds.yaml |
separate pool per conditioning |
| AF3-ReD | diffusion_side/m4_af3red.yaml |
RMSD repulsion |
| Boltz-Perturb | diffusion_side/b4_boltz_perturb.yaml |
conditioning perturbation |
| Noise search | diffusion_side/b2_noise_search.yaml |
evolutionary search over initial noise |
Three of these ship an OpenDDE variant with the _opendde suffix, b0b_opendde.yaml, m2_tds_opendde.yaml and m5_cast_opendde.yaml.
| Path | Contents |
|---|---|
config.py |
YAML loader, paths resolved relative to the config file |
results_io.py |
per_sample.parquet / per_target.csv schema |
data/ |
target.py, datasets/ loaders, msa/ generation and subsampling |
method/ |
diffusion_side/, trunk_side/, metrics/, verifier/ |
models/ |
backend adapters and inference hooks |
experiments/ |
configs/ and pipeline scripts/ |
jobs/ |
job templates and submission scripts |
The settings usually changed, all in the method.diffusion_search block of an arm config or
overridable with --set. Defaults are CAST's reference arm.
Budget and pool
| Knob | Default | Effect |
|---|---|---|
subsample_pool (K) |
4 | conditionings sharing one pool. K is the axis that scales; N alone does not |
n_samples_per_seed (N) |
720 | particles in the pool. Must satisfyN >= 3 * K * umbrellas * sites, or preflight refuses |
weight_level_correction |
true | required at K>1; restores the cross-latent baseline the difference potential divides out |
The site reward
| Knob | Default | Effect |
|---|---|---|
reward |
epitope |
the site collective variable. The joint pool requires this |
umbrella_sweep |
9,11,13,15 |
arrival radii d*, brackets the observed native paratope-epitope range |
site_hinge |
true | flat zero-reward ball inside d*. False gives a plain harmonic that keeps pulling inward |
site_huber_delta |
4.0 | bounds the far side; 0 leaves it unbounded quadratic |
epitope_source |
predicted |
ground_truth swaps in the oracle map, the ablation ceiling |
epitope_map_builder |
hybrid_modes_tiles |
pool modes first, then surface tiles ranked by mode support |
The tilt
| Knob | Default | Effect |
|---|---|---|
guidance_scale |
75.0 | gradient strength. 0 turns the tilt off entirely, leaving the joint pool unsteered |
guidance_max_rms_a |
6.0 | size-invariant displacement rail |
lam_adaptive |
true | calibrate lambda once from the observed reward spread |
lam |
4.0 | the cap on the calibrated lambda, not inert |
target_ess_frac |
0.4 | ESS target the calibration solves for |
Resampling
| Knob | Default | Effect |
|---|---|---|
potential |
difference |
score improvement rather than absolute reward |
resample / ess_fraction |
ess / 0.5 |
resample a cell when its ESS falls below half its size |
resample_windows |
0.25:0.35 |
the early window resampling is confined to, which bounds pool collapse |
schedule_start_frac, schedule_every |
0.25, 20 | when scoring events begin and how often they fire |
At the job level the same budget is set by N_SAMPLES, SUBSAMPLE_POOL, N_SEEDS (must equal
SUBSAMPLE_POOL), UMBRELLA_HYBK12 and SCALE.
jobs/submit_steer_pool.sh freezes a template (jobs/tts-*-pool.yaml) into a self-contained shell
body and starts N workers, one per GPU, all draining a shared queue under $RUN_DIR/queue/. A unit
is one target x arm. Workers are independent, so they may start at different times and the grid
still completes.
The template is not a thin wrapper. It carries the gates that decide whether a result is trustworthy:
| Stage | What it does |
|---|---|
| checkpoint route | resolves private-vs-public weights and logs which one,before a GPU is taken |
| dataset check | validates every target in the index and reports problems before any GPU work |
| preflight | builds the arm list and rejects invalid ones. Catches per-cell budget violations (n_samples too small for umbrellas x sites x K) and maps with no _source/_builder provenance stamp |
| queue | pending/ claimed/ done/ failed/ plus attempts/, heartbeat/, force_offload/. Claims are atomic mkdir, so workers never collide |
run_unit |
invokesrun_ablation --config <cfg> --set ... with the arm's overrides and per-unit paths |
| health gate | refuses to score a silently-failed run. ExpectsN+K-1 structures for a joint pool and N_SEEDS x N_SAMPLES for the control, because the two arms legitimately emit different counts |
| memguard | kills the unit atMEMGUARD_PCT (92) of the cgroup limit rather than letting the pod be OOM-killed |
| OOM classifier | separatesgpu_size (escalate to confidence offload) from gpu_contention (requeue, do not escalate) from ram, and records each in oom_events.jsonl |
score_unit |
runs DockQ concurrently atnice 19 ionice -c 3, so scoring never delays the GPU |
reclaim |
returns units whose worker died, via the heartbeat |
| dead-letter | retires a unit afterMAX_ATTEMPTS (2) so one bad unit cannot starve the pool |
Useful knobs: MAX_ATTEMPTS, MEMGUARD_PCT, DEADLINE_S/UNIT_EST_S (retire cleanly rather than
be killed mid-unit), SCORE_JOBS, ARMS_ONLY, SKIP_CTRL, TARGETS.
A backend is selected by name everywhere, so a new one is a registration plus wiring, not a change to any method or experiment code.
- Implement
ModelBackend(models/base.py) inmodels/<name>.py:emit_input,get_msa_sources,set_msa_sources,get_chain_sequences,run,iter_predictions,load_pae_and_coords.runtakes aBackendRunSpec(seeds,n_cycles,n_samples_per_seed,n_diffusion_steps,use_msa,extras). - Register it at module scope:
register_backend("<name>", <Name>Backend). - Auto-import it by adding
"<name>"to the tuple inmodels/__init__.py::_auto_register. Import failures there are non-fatal, so a missing runtime disables that backend only. - Add a peak-memory scale to
_PEAK_SCALE_BY_BACKENDinmodels/base.py, relative to the Protenix curve the model was fitted on. Measure it rather than assuming 1.0. - Install sampling hooks if the backend needs them.
models/protenix_hooks.pyholds the steering hook plus vendor-quirk patches, each self-gating oncofold_pkg()so it no-ops elsewhere. - Add per-backend configs as
<config>_<suffix>.yamlbeside the originals, e.g.m5_cast_opendde.yaml. The job template picks them up throughTTS_CFG_SUFFIXandcfg_for_backend. - Run it with
BACKEND=<name> TTS_CFG_SUFFIX=_<name>in the env file.SCORE_BACKENDdefaults toBACKEND, since the results reader must follow the model. - Check that sample chunking preserves coupling. This is the one correctness trap. A backend
may split the particle batch to fit memory, and that is exact for the denoiser, which has no
cross-particle term. But under a joint pool the chunk must be taken within one conditioning
group, never across groups, and every chunk's output must be scattered back into a single
N-particle tensor before anything else reads it. CAST's coupling lives in the SMC layer, the
reward, the weights and the resample, and that layer must see the whole pool. The reference
implementation does this by iterating groups first and chunking inside each
(
models/protenix_hooks.py, the_grpbranch withindex_copy_). A backend that chunks the flat batch will silently mix conditionings and break the cross-subsample comparison the joint pool exists for. - Verify with a small end-to-end run: preflight passes, the health gate sees
N+K-1structures, and DockQ is non-null inper_sample.parquet.
The memory model. models/base.py::peak_gib_for is the one cross-backend entry point: it scales
the Protenix-fitted curve by _PEAK_SCALE_BY_BACKEND (protenix2 1.0, the curve it was fitted on,
opendde 1.03). The curve itself, models/protenix2.py::_predicted_peak_gib, is the larger of two
measured paths:
inflight = base(N_token) + chunk*slope(N_token) + max(0, n_samples - chunk)*frac*slope(N_token)
peak = max(inflight, atom_attn_base(N_token) + chunk*atom_attn_slope(N_token))
The middle term is the retained-sample accumulation, which is why a large pool can OOM at a chunk that a small one survives, and the second path is why shrinking the chunk eventually stops helping.
The decisions built on it live in models/protenix2.py:
- Diffusion chunk,
PROTENIX_SAMPLE_DIFFUSION_CHUNK_SIZE=auto->auto_diffusion_chunk. It fits first and optimises second:max_diffusion_chunkasks whether anything fits on-GPU at all, thenoptimal_diffusion_chunkpicks the chunk to actually use, which is the smallest that reaches the speed knee, not the largest that fits. Below_MIN_ONGPU_CHUNKit falls back to an offload chunk. - Confidence offload,
PROTENIX_CONF_OFFLOAD_TOKENS=auto->should_offload. Free-GPU-first: it moves the PAE/PDE stack to host RAM only when the target cannot safely stay on-GPU, via three hard rails (a card-aware token ceiling, a sample ceiling, and the chunk floor) and then the chunk math. Offload is a fitting tool, not a speedup, and it trades GPU pressure for host pressure, so a pod that offloads needs a larger memory limit rather than a smaller chunk.
Opt-in controls for large joint pools. All three are exact, so they change memory and not results. The first two are unset by default and do nothing until you set them; the third already has a default of 5, which you lower.
| Variable | What it does | When you need it |
|---|---|---|
TTS_JOINT_CACHE_CPU=1 |
stages the joint conditioning cache through host RAM | the cache holds K trunk outputs, about 3.5 GB each, so K=16 will not fit on an 80 GB card beside the working set |
TTS_ODDE_SC_PAIR_CHUNK=16 |
shrinks OpenDDE's shape-complementarity token chunk | the vendor default of 128 materialises[n_sample, 128, n_atom, 3], about 37 GB for a 2880-particle pool on a 1171-token target |
PROTENIX_CONF_SAMPLE_CHUNK=2 |
shrinks the confidence sub-batch | the confidence peak is roughly27 + 9.5 * chunk GiB, so the default of 5 leaves only ~5 GiB of headroom on a large pool |
The last two fail after the diffusion has completed, which costs the whole unit, so set them ahead of a large run rather than waiting for the failure.
Apache-2.0. See LICENSE.md.
Release notes are in NEWS.md, and CONTRIBUTING.md covers issues and pull requests.