UNAAGI: Atom-Level Diffusion for Generating Non-Canonical Amino Acid Substitutions

TL;DR

UNAAGI is a diffusion-based model that predicts amino acid substitutions at the atomic level, including non-canonical amino acids, improving mutation effect prediction and enabling broader protein engineering applications.

Contribution

Introduces UNAAGI, a novel E(3)-equivariant diffusion model for residue reconstruction that includes non-canonical amino acids, expanding the scope of structure-based mutation prediction.

Findings

Outperforms existing methods on NCAA substitution benchmarks.

Enables exploration of non-canonical amino acids in protein design.

Suggests a unified approach for protein engineering and drug design.

Abstract

Proposing beneficial amino acid substitutions, whether for mutational effect prediction or protein engineering, remains a central challenge in structural biology. Recent inverse folding models, trained to reconstruct sequences from structure, have had considerable impact in identifying functional mutations. However, current approaches are constrained to designing sequences composed exclusively of natural amino acids (NAAs). The larger set of non-canonical amino acids (NCAAs), which offer greater chemical diversity, and are frequently used in in-vivo protein engineering, remain largely inaccessible for current variant effect prediction methods. To address this gap, we introduce \textbf{UNAAGI}, a diffusion-based generative model that reconstructs residue identities from atomic-level structure using an E(3)-equivariant framework. By modeling side chains in full atomic detail rather than as discrete tokens, UNAAGI enables the exploration of both canonical and non-canonical amino acid substitutions within a unified generative paradigm. We evaluate our method on experimentally benchmarked mutation effect datasets and demonstrate that it achieves substantially improved performance on NCAA substitutions compared to the current state-of-the-art. Furthermore, our results suggest a shared methodological foundation between protein engineering and structure-based drug design, opening the door for a unified training framework across these domains.