zPocket: ESMFold2 pairformer2

zPocket heads predict which residues of a protein line a ligand-binding pocket. They read the frozen pair representation z of a cofolding model. This repository holds one head: pairformer2 on the ESMFold2 trunk, the main-text model of the zPocket paper. Given one protein's (L, L, 256) ESMFold2 pair tensor, it returns (L,) per-residue logits. sigmoid(logits) is a per-residue pocket score.

Weights only

This release contains weights and a configuration, with no code. zPocket's inference package is being released separately. Until then, the head can be rebuilt from config.json. architecture.forward lists every operation, and architecture.state_dict.shapes lists every tensor. The head has 330,177 parameters.

Architecture

• proj: Linear 256 → 64 over z.
• Two pairformer blocks, each with three residual updates:
  • outgoing triangle multiplication;
  • incoming triangle multiplication;
  • SwiGLU transition, 64 → 256 → 64.
• Per-residue pooling: the concatenation of residue i's row mean, column mean and diagonal entry x[i, i], giving 192 features.
• Readout: LayerNorm, then a two-hidden-layer GELU MLP (192 → 256 → 256 → 1).
• All LayerNorms use eps 1e-5. Dropout (p = 0.1) is on the triangle branches and is inactive at inference.

Precision. The head was trained and evaluated on GPU:

• Parameters were float32.
• The two pairformer blocks ran under torch.autocast(..., dtype=torch.bfloat16), with their LayerNorms computed in float32.
• The projection and readout ran in float32.

Running entirely in float32 on CPU moves probabilities by at most 0.010 on the proteins we checked (see Verification).

Input: the exact feature

The head reads one array, z, and nothing else. It does not use the single representation, the ESMC embeddings, the distogram, MSA features or coordinates.

• z has shape (L, L, 256) for a single protein chain of length L. z[i, j] is indexed by 0-based residue position, and index i is canonical UniProt position i + 1.
• It was stored as float16 and is cast to float32 before the head.
• It is fed as is:
  • no normalisation (entries routinely reach the hundreds to thousands);
  • no symmetrisation;
  • no cropping or padding;
  • one protein per forward pass.
• Training proteins were 47–1,200 residues long.

How to generate it. These steps match how the training and evaluation features were made:

1. Fold the canonical UniProt sequence as a monomer with no ligands. Residues without a standard amino-acid code (B, Z, U, O) were input as X. Use biohub/ESMFold2, the MSA-conditioned model, not ESMFold2-Fast, with its ESMC-6B language model loaded in bfloat16. We used:

   • esm at commit af8ef5cead4388c4d7af04168157c98898c78805 and the Biohub transformers fork at commit 3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf;
   • ESMFold2InputBuilder().fold(model, spi, num_loops=10, seed=0, ...), which runs 11 recurrent trunk iterations;
   • the checkpoint's own language-model dropout setting, left unchanged.

   config.json records the weight file's sha256.

2. MSA. Attach one merged A3M per chain, built from jackhmmer searches of UniProt, UniRef90, reduced BFD and MGnify:

   • up to 8,192 deduplicated UniProt hits;
   • then up to 8,191 hits pooled from the other three databases, deduplicated across them;
   • the query first, at most 16,384 rows in total.

   Parse it with MSA.from_a3m(path, remove_insertions=True, max_sequences=16384). The exact merge order is in config.json under feature_provenance.msa.

3. Capture z as the output of model.parcae_coda with a forward hook. This is the final trunk pair tensor, which feeds the distogram and structure heads. Drop the batch axis:

   captured = {}
   hook = model.parcae_coda.register_forward_hook(
       lambda module, inputs, output: captured.__setitem__("z", output.detach())
   )
   result = ESMFold2InputBuilder().fold(model, spi, num_loops=10, seed=0, ...)
   hook.remove()
   z = captured["z"][0].float()        # (L, L, 256); then: logits = head(z)
   

ESMFold2's trunk is stochastic at inference, because of its random initial pair state and its per-loop language-model dropout. Regenerated features are therefore a fresh draw, even at seed 0, and will not match ours bit for bit. The diffusion-sampling settings act after z is computed and do not affect it.

Output

logits: float32, shape (L,). Higher means a residue is more likely to lie within 5.0 Å of a biologically relevant ligand in some PDB entry of that protein. Apply sigmoid for a score in [0, 1].

The head was trained with unweighted binary cross-entropy. No post-hoc calibration was applied or evaluated, so the scores are best used to rank residues within a protein. Turning them into discrete pockets or centroids is not part of this release.

Training data

The labels are PocketWatch, version def5a_lightly_curated_20260827: Q-BioLiP / BioLiP2 ligand contacts at 5.0 Å, aggregated per canonical UniProt sequence.

The split is temporal, by the first release date of a protein's pocket sites. A protein goes to test if any site first appeared on or after 2024-01-01, to validation if any first appeared after 2021-09-30, and to train otherwise.

Training setup

• Data: all 11,016 training proteins, at full length.
• Positives included interface-only contacts.
• Loss: per-residue binary cross-entropy.
• Optimiser: Adam with lr 2e-3, 8 proteins per step, 20 epochs.
• Seed: 0, a single seed.
• Checkpoint selection: best validation per-target AUPRC, reached at the 16th epoch.

Validation and test residues in pockets first seen before the split's own cutoff are masked from scoring.

Evaluation

These are test-split results from the zPocket paper. The protocol:

• Interface-only pockets, pockets seen before the cutoff, and pockets subsumed by seen pockets are masked.
• Scoring covers the 1,275 test targets and 1,562 scored pockets for which every compared method returned a prediction.
• Per-target metrics are computed within each protein and averaged with equal weight; per-pocket metrics likewise per pocket.
• AUPRC is average precision.

The random baseline for per-target AUPRC is 0.049. Per-pocket AUPRC is 0.751 on proteins with one scored pocket and 0.599 on proteins with more than one.

The training run's own test evaluation gives per-target AUPRC 0.734 over 1,559 targets. The gap to 0.745 is almost entirely the target set. That scorer counts contacts that exist only between two chains as positives, so it also scores 280 proteins whose only new pockets are interface pockets, which a single-chain model cannot fully represent. The head reaches 0.676 on those and 0.747 on the rest. The benchmark excludes them for every single-chain method, giving 0.747, plus 4 proteins one baseline did not predict, which has no effect. Masking interface-only residues in the remaining proteins then gives 0.745 (−0.001).

Verification. We loaded model.safetensors into a module built from config.json and ran it on 24 random test proteins' ESMFold2 features. It reproduces the paper's stored test predictions:

• GPU, bfloat16 path (as trained):
  • largest per-residue probability difference 2.5e-3;
  • 21 of the 24 proteins agree to within 1e-6;
  • per-target AUPRC over the 24 proteins is identical (0.655).
• CPU, float32:
  • largest difference 0.010;
  • per-target AUPRC moves by 2e-4.

Limitations

• ESMFold2-specific. The head was trained on z from this ESMFold2 checkpoint, captured at this point and generated with these settings. It is not expected to transfer, without retraining, to:
  • other trunks;
  • ESMFold2-Fast;
  • other capture points;
  • substantially different MSA or recycling settings.
• Familiar pockets dominate the headline number. Coverage here is the largest share of a pocket's residues that a ligand-lining pocket in a structurally similar pre-cutoff PDB structure covers. Per-pocket AUPRC is 0.803 for the 1,119 test pockets at 80–100% coverage, and 0.327 for the 133 pockets below 20%. On the 40 cryptic pockets it is 0.547.
• Single chain, no ligand. Features come from monomer folds. Pockets that exist only at interfaces between chains were masked in evaluation, but they counted as positives in training.
• Scope.
  • Lengths of 47–1,200 residues were seen in training; longer chains are untested.
  • Compute is O(L³) and memory O(L²).
  • This is one model from one seed.
  • Scores are per residue, not discrete pockets, and their calibration was not evaluated.

Licence

The weights in this repository are released under the MIT licence (see LICENSE; Copyright (c) 2026 Genesis Molecular AI). Generating inputs requires ESMFold2 and ESMC-6B, which Chan Zuckerberg Biohub distributes under its own terms (MIT, plus the third-party notices linked from those model cards). The training labels, PocketWatch, are MIT.

Citation

@misc{portnoi_zpocket,
  title = {{zPocket}: Better Pocket Finding with Frozen Co-folding Features},
  author = {Portnoi, Tally and Keys, Allison M and DiFrancesco, Benjamin R and Gruver, Nate},
  year = {2026},
  eprint = {2610.17696},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url = {https://arxiv.org/abs/2610.17696}
}