Pi-SAGE

Whether a point mutation strengthens or weakens a protein-protein interface is governed by what happens on the molecular surface: the electrostatics, hydrophobicity, and curvature that two partners present to each other. State-of-the-art predictors of binding free-energy change work from all-atom structural graphs and are left to infer those surface properties implicitly. Pi-SAGE — Permutation-Invariant Surface-Aware Graph Encoder — tests the opposite premise, that giving a structural model an explicit, discretized description of the surface is worth more than making it rediscover one.

Developed at Amazon with a collaborator at MIT and presented at the ICML 2025 Workshop on Generative AI for Biology, Pi-SAGE learns a protein surface codebook directly from structure and assigns every surface-exposed residue a token drawn from that vocabulary. The idea parallels the discrete structural alphabets used elsewhere in protein modeling — the Foldseek 3Di states that SaProt pairs with amino acid tokens, for instance — but the vocabulary here describes local chemical and geometric context on the solvent-accessible surface rather than backbone conformation. The tokens are then used as extra node features for GearBind, an all-atom geometric network built on GearNet message passing, to predict how mutations shift binding affinity.

Pi-SAGE is a short workshop paper rather than a full conference or journal publication, and it is evaluated as an augmentation to an existing predictor on a single benchmark. No public code or weights release accompanies it, and no license is stated.

#Key Features

• Surface codebook via finite scalar quantization: A fixed-grid quantizer maps continuous surface-patch embeddings to discrete codes with a straight-through estimator for gradients, producing a reusable token vocabulary instead of a single-purpose embedding.
• Permutation invariance by construction: A permuter module learns a soft alignment between input and reconstructed graphs using a continuous relaxation of argsort, so reconstruction loss does not depend on arbitrary node ordering.
• Explicit chemical and geometric surface features: Each patch carries five MaSIF-derived descriptors — electrostatic potential, hydrophobicity, hydrogen bond propensity, shape index, and distance-dependent curvature — plus two geometric terms, the patch centroid-to-C-alpha distance and the C-alpha-to-C-beta angle.
• Two-stage training that separates general from task-specific: General surface representations come from a large unlabeled structure corpus; only a domain adaptation stage touches the affinity data, keeping the codebook broadly reusable.

#Technical Details

The encoder is a graph transformer based on GRPE, which combines dot-product attention with learnable relative positional encodings for topological and edge relationships. Each residue is represented by 32 sampled surface patches — 70% from high-, 20% from medium-, and 10% from low-scoring patches — connected by edges wherever patch centroids fall within 3 Å, and each patch is described by a 7-dimensional feature vector. The encoder output passes through the finite scalar quantizer, whose implicit codebook size is the product of the per-dimension level counts, and a mirrored decoder reconstructs node features and the adjacency matrix under a combined cosine-similarity and MSE objective with a permutation entropy penalty. Four configurations are reported, from a 2-layer, 512-dimensional 13M-parameter model to a 16-layer, 1280-dimensional 378M-parameter one. Stage one pretrains on roughly 200,000 experimentally determined structures from the RCSB Protein Data Bank; stage two domain-adapts on SKEMPI, about 6,000 binding affinity measurements over 340 unique complexes. Augmenting GearBind with the resulting tokens raises per-complex Pearson correlation from 0.525 ± 0.106 to 0.600 ± 0.084, a 14.3% relative gain, with per-structure Pearson 0.423 ± 0.091, Spearman 0.428 ± 0.038, RMSE 1.817 ± 0.241, and AUROC 0.691 ± 0.026.

#Applications

The natural use is antibody and protein engineering, where in-silico affinity maturation depends on ranking large sets of candidate point mutations before any of them reach a wet lab. A surface tokenizer is also a general-purpose component: because the codebook is learned from unlabeled structures, its tokens can in principle feed any structure-based model of interfaces, including epitope and binding-site analysis, without retraining the tokenizer for each downstream task.

#Impact

Pi-SAGE contributes a clean ablation to an ongoing question in protein representation learning — whether surface information is fully recoverable from all-atom structure or deserves its own encoder. The measured improvement over GearBind on SKEMPI supports the latter, and the discrete-token formulation makes the surface representation portable in a way that continuous surface fingerprints are not. The evidence base is deliberately narrow: one downstream model, one benchmark, and, as the authors note, ablations and sensitivity analysis still to come for a method with several interacting components. Without a code or weights release, independent evaluation of the tokenizer on other interface tasks is not yet possible.