SoftAlign

Structure prediction has produced hundreds of millions of protein models, and comparing them at that scale is now the bottleneck. The classical structural aligners, TM-align and DALI, are accurate but expensive: their similarity scores are non-local, which rules out the dynamic-programming shortcuts that make sequence search fast. Foldseek solved the throughput problem by discretizing each structure into a 20-letter 3Di alphabet and handing the result to a sequence search engine, trading alignment quality for orders of magnitude in speed.

SoftAlign takes the opposite route: train the aligner end to end. It maps a pair of 3D coordinate sets to an alignment in one differentiable pipeline — a retrained ProteinMPNN encoder produces a per-residue embedding, the scalar products of those embeddings form a similarity matrix, and a differentiable alignment operator turns that matrix into a soft alignment that gradients can flow through. The model comes from Jeanne Trinquier at Sorbonne Université and Sergey Ovchinnikov at MIT, with co-authors including Martin Steinegger and Michel van Kempen of the Foldseek team and Johannes Söding at the Max Planck Institute for Multidisciplinary Sciences, and was posted to bioRxiv in May 2025.

The more portable contribution is the softmax pseudo-alignment. Applying row- and column-wise softmax to the similarity matrix and multiplying the two produces something that is not a valid alignment — it does not preserve residue order — but is a markedly better discriminator for structure search, and the operator is architecture-agnostic.

#Key Features

• Fully end-to-end training: Both alignment operators are differentiable, so the encoder is optimized directly against alignment quality; Foldseek's 3Di alphabet, by contrast, is learned by a separate autoencoder before a substitution matrix is fitted.
• Softmax pseudo-alignment as a filter: By aligning only positions with mutual row-and-column agreement, the softmax variant refuses to align anything when no real correspondence exists, driving false-positive scores toward zero where Smith-Waterman forces a plausible-looking path.
• Smith-Waterman rescue at inference: A model trained with softmax can be run with Smith-Waterman at inference, lifting mean LDDT from 0.44 to 0.52 and combining the discriminative training signal with a well-formed output alignment.
• Differentiable LDDT loss: A sigmoid-smoothed local distance difference test removes the dependence on TM-align reference alignments, making the objective self-contained.
• Categorical variant for prefiltering: An optional discretization step, initialized by k-means and updated with a straight-through estimator, produces a structural alphabet analogous to 3Di for use with sequence-search prefilters.

#Technical Details

Training uses SCOPe40, which holds 11,211 protein domains clustered at 40% identity. Splits are by fold rather than sequence identity — 6,506 proteins across 681 folds for training, 4,507 across 513 folds for testing — so structurally similar but sequence-divergent pairs cannot leak. Training pairs are restricted to TM-score above 0.6. The encoder reuses ProteinMPNN's three-layer message-passing architecture over interatomic distances, backbone dihedrals, and the 64 nearest neighbours per residue, but is retrained from scratch; ProteinMPNN's design weights and decoder are not used. The alignment temperature is annealed from 5 to 1 during training and set to 1e-4 at inference.

On 1,000 held-out SCOPe40 pairs, the Smith-Waterman variant reaches mean LDDT 0.56 against TM-align's 0.57, ahead of Foldseek's best configuration and PLMAlign, both at 0.47. The categorical variant reaches 0.53, beating Foldseek despite using the same discrete strategy. On all-versus-all SCOPe40 search, measured by fraction of true positives before the first false positive, the continuous softmax model outperforms every structural aligner tested, including DALI. That run over 11,211 proteins takes about 6,425 seconds including roughly 180 seconds of embedding generation — slower than Foldseek, competitive with TM-align and DALI.

#Applications

The immediate use is remote-homology detection and fold classification over large structure databases, where the softmax variant serves as a high-precision reranker behind a fast prefilter; the repository ships a Colab notebook that queries the AlphaFold Database through the Foldseek API and reranks the hits. Beyond that, the softmax operator is the transferable piece for anyone building similarity search on learned embeddings, since it supplies a discriminative score without the false-positive inflation that forced alignment produces.

#Impact

SoftAlign is a clean demonstration that end-to-end training closes most of the accuracy gap between learned and classical structural aligners, and that continuous embeddings support search better than the discretized alphabet they were compared against. The authors are candid about the limits. The softmax variant is a poor aligner in the conventional sense: normalized by query length its LDDT and TM-scores fall well short of TM-align, and only when normalized by the number of aligned positions do they match. Softmax is also incompatible with the categorical variant, since identical discrete similarity values yield identical probabilities. Throughput remains behind Foldseek, positioning the method as a reranking stage rather than a drop-in replacement. The preprint has not yet been peer reviewed, and the repository carries no license file, so reuse terms for the code and released weights are unstated.