jaxfvld

jaxfvld is a small, fast, fully differentiable antibody Fv structure predictor in JAX/Flax. It folds an antibody variable region (paired VH and VL, or a single VHH domain) from sequence to an all-atom backbone, returning per-residue frames, side-chain torsions, atom14 coordinates, and a per-residue pLDDT. It takes no antigen and no input structure.

It is ~0.76 M parameters, BSD-3-Clause, and matches IgFold on CDR-H3 (see Accuracy). Provenance and citation are in ATTRIBUTION.md and CITATION.cff.

Why use it

• Fully differentiable: the sequence enters as a soft PSSM, so gradients flow from any function of the predicted structure back to the sequence, for gradient-based design.
• Small and fast: ~0.76 M parameters; a single seed folds a ~230-residue Fv in about 0.8 s on CPU (the 4-seed ensemble in about 2 s), no GPU required.
• Accurate on the hard loop: ~2.4 Å median CDR-H3 backbone RMSD on a leak-free, post-2021 test set.
• Permissive: BSD-3-Clause.
• JAX-native: jit / vmap / grad-compatible and composable inside a larger model.

Quickstart

Install:

pip install jaxfvld
pip install "jaxfvld[ablanx]"    # + the AbLang2 feature path the released checkpoint uses
# from source: pip install git+https://github.com/fabricagen/jaxfvld

Fold:

from jaxfvld import fold_fv

out = fold_fv(
    vh="QVQLVESGGGLVQPGGSLRLSCAAS...",
    vl="DIQMTQSPSSLSASVGDRVTITCRAS...",   # omit for a single-domain VHH
    weights="jaxfvld_weights.npz",
    ablanx_weights="ablang2_weights.npz",
)
out["atom14"]   # [N, 14, 3] all-atom coordinates
out["plddt"]    # [N] per-residue confidence

Command line:

jaxfvld fold --vh QVQL... --vl DIQM... --out fv.pdb \
    --weights jaxfvld_weights.npz --ablanx-weights ablang2_weights.npz
# --single folds the best single seed instead of the 4-seed ensemble (faster, ~0.25 Å worse on CDR-H3)

The released checkpoint was trained with the AbLang2 feature path and depends on it: pass the ablang2 weights (from the ablanx release) to fold at the reported accuracy. Without them CDR-H3 collapses (median ~18 Å vs ~2.5 Å), so the path is effectively required. Weights from weights= / --weights / JAXFVLD_WEIGHTS; ablanx weights from ablanx_weights= / --ablanx-weights / JAXFVLD_ABLANX_WEIGHTS (or ABLANG_WEIGHTS).

What it is

The forward pass is a featurizer, a structure module, and an all-atom decoder:

1. The soft PSSM is projected to a single representation with a per-chain embedding (VH = chain 0, VL = chain 1); the pair representation encodes relative position and a cross-chain indicator, mixed by a triangle-multiplicative update. An optional AbLang2 sequence-embedding + attention path feeds both; the released checkpoint uses it.
2. Invariant Point Attention iterates from a black-hole initialization, updating the single representation and a per-residue rigid frame. The block is weight-tied (n_layer = refinement iterations), and the trunk recycles n_recycle times with a stop-gradient.
3. The final frame and seven predicted torsions become atom14 coordinates via the vendored AlphaFold geometry tables; a head predicts pLDDT.

Training minimizes Frame-Aligned Point Error, so the loss is invariant to a global rigid transform and independent of centering. The released architecture is Fold(c_s=128, c_z=128, n_layer=4, n_recycle=3, tri_pair=True), max_len=280. The weights archive carries this config, so the loader reconstructs the exact model and the weight-tied n_layer/n_recycle trap cannot bite. See TECHNICAL_BRIEF.md for the architecture diagram, the AbLang2 path, and the geometry provenance.

Accuracy

The shipped model is a 4-seed ensemble at the converged step. CDR-H3 backbone RMSD is the headline metric (the hard loop that decides antibody structure quality), measured framework-aligned the standard IgFold / ABodyBuilder2 way. The test set is leak-free: no test antibody shares a CDR-H3 cluster with training, restricted to post-2021 depositions. Both predictors are raw (no OpenMM relaxation).

Per-region median backbone RMSD, set300 (n = 300):

Paired CDR-H3 on the 153 PDBs both predictors share, bootstrapping the per-structure paired delta directly (do not read small differences off overlapping CIs): median delta = −0.067 Å, 95% CI [−0.16, +0.08] (crosses zero), win-rate 54%. Both predictors are scored on the identical 153 PDBs, nothing dropped on either side.

On CDR-H3 the paired delta straddles zero: a statistical tie with the IgFold baseline, not a beat, from a 0.76 M-parameter model; the baseline keeps a small edge on the framework. See Validation for the reproduction, and reproduce the tables and figure from the committed fixtures:

python -m jaxfvld.eval.parity_report          # prints the tables, writes figures/parity.png

Speed and size

The trained folder is small and folds a variable region on CPU, no GPU. On one VH+VL Fv, one CPU, no refinement:

The 4-seed ensemble folds about 3x faster than IgFold's 4-model ensemble on the same CPU, from a trained folder about half the per-model size, at a statistical tie on CDR-H3; a single seed folds in about 0.8 s. Both models pair the folder with a frozen antibody language model as a sequence prior (jaxfvld: AbLang2, 44 M; IgFold: AntiBERTy, ~26 M), which is shared rather than trained here and is not included in the trained-folder rows above.

Reproduce the folder size and fold time (add IgFold with the [igfold] extra for the head-to-head):

JAXFVLD_WEIGHTS=jaxfvld_weights.npz ABLANG_WEIGHTS=ablang2_weights.npz \
    python -m jaxfvld.eval.speed_size --igfold

Validation

Folding the leak-free set300 through the released package reproduces the committed benchmark: the aggregate per-region medians match to within 0.01 Å, and every per-PDB CDR-H3 matches its recorded value to a mean of 0.013 Å. Predicted confidence tracks error (pLDDT calibration). See VALIDATION.md for the receipt, the figures, and the exact reproduction commands.

Differentiable sequence design

Because the sequence enters as a soft PSSM, gradients flow from any function of the predicted structure back to the sequence logits. examples/ has two worked scripts: design a soft sequence so the folded Fv matches a target CDR-H3 geometry, and maximize predicted foldability. This is the main reason to use jaxfvld over a non-differentiable folder.

Weights

Trained weights are attached to the GitHub release as a single archive, jaxfvld_weights.npz (4 seeds + the best single seed), which carries its own reconstruction config. The AbLang2 weights (ablang2_weights.npz) come from the ablanx release. export_ensemble.py builds the archive from per-seed checkpoints.

Tests

pip install "jaxfvld[dev]"
python -m pytest -q -k "not weights_reproduce"     # IPA geometry, inference, parity fixtures (no weights)
# the frame/FAPE/fold/padding gates need parsed data and self-skip without it; run them where the data lives:
JAXFVLD_TEST_DATA=<allatom shards> SABDAB_DIR=<sabdab2 dir> \
    python -m pytest -q tests/test_frames.py tests/test_fape.py tests/test_fold.py tests/test_padding_invariance.py
# the weights-gated reproduction folds set300 through the package; run it on a GPU box, not a laptop:
JAXFVLD_WEIGHTS=... ABLANG_WEIGHTS=... JAXFVLD_EVAL_DATA=... python -m pytest -q tests/test_parity.py

Evaluate, train, and ingest data

jaxfvld.eval scores any predictor on the same metric (cdr_rmsd, igfold_headtohead, plddt_calib); jaxfvld.train.train_folder trains the folder (pip install "jaxfvld[train]"); jaxfvld.data is the self-contained SAbDab / SAbDab2 ingestion path. Data comes from SAbDab (Oxford OPIG), used under its own terms and not redistributed. Locations come from environment variables with local defaults.

Ecosystem

jaxfvld is the structure folder in the seam bundle, which couples it with the ablanx sequence prior. For antibody developability screening and repair, see sift: https://sift.fabricagen.ai (coming soon, ca 08/26).

License

BSD-3-Clause. See LICENSE and ATTRIBUTION.md. jaxfvld/model/residue_constants.py is vendored from AlphaFold (Copyright 2021 DeepMind, Apache-2.0) with its original header retained; those terms are compatible with BSD-3-Clause redistribution.