[Paper Review] TaxDiff: Taxonomic-Guided Diffusion Model for Protein Sequence Generation

Lin Zongying, Hao Li+6 more|arXiv (Cornell University)|Feb 27, 2024

Microbial Metabolic Engineering and Bioproduction4 citations

TL;DR

TaxDiff is a novel diffusion model for controllable protein sequence generation that integrates taxonomic information into each transformer block to guide the generation of structurally stable proteins. By combining global and local attention via a patchify mechanism, it achieves state-of-the-art performance in both unconditional and taxonomic-guided generation, surpassing prior models in structural metrics and reducing inference time to just 25% of baseline diffusion models.

ABSTRACT

Designing protein sequences with specific biological functions and structural stability is crucial in biology and chemistry. Generative models already demonstrated their capabilities for reliable protein design. However, previous models are limited to the unconditional generation of protein sequences and lack the controllable generation ability that is vital to biological tasks. In this work, we propose TaxDiff, a taxonomic-guided diffusion model for controllable protein sequence generation that combines biological species information with the generative capabilities of diffusion models to generate structurally stable proteins within the sequence space. Specifically, taxonomic control information is inserted into each layer of the transformer block to achieve fine-grained control. The combination of global and local attention ensures the sequence consistency and structural foldability of taxonomic-specific proteins. Extensive experiments demonstrate that TaxDiff can consistently achieve better performance on multiple protein sequence generation benchmarks in both taxonomic-guided controllable generation and unconditional generation. Remarkably, the sequences generated by TaxDiff even surpass those produced by direct-structure-generation models in terms of confidence based on predicted structures and require only a quarter of the time of models based on the diffusion model. The code for generating proteins and training new versions of TaxDiff is available at:https://github.com/Linzy19/TaxDiff.

Motivation & Objective

• To address the limitation of unconditional protein sequence generation in biological applications, where random generation requires extensive filtering.
• To enable fine-grained, controllable protein design by incorporating biological taxonomy as a conditioning signal.
• To improve structural consistency and modeling accuracy in generated protein sequences using a hybrid attention mechanism.
• To reduce inference time while maintaining or improving performance compared to existing diffusion-based protein generation models.

Proposed method

• Integrates taxonomic control features into every layer of the denoising transformer block to enable conditional generation based on species or family classification.
• Proposes a patchify attention mechanism that combines global attention over the full sequence with local attention over fixed-size amino acid patches (e.g., 16 residues).
• Reclassifies protein sequences at the family and species levels to consolidate overly granular UniProt classifications, improving training stability and control fidelity.
• Adopts a diffusion-based generative framework similar to EvoDiff, enabling joint generation of sequences and structures in a shared latent space.
• Uses a shared architecture for both unconditional and conditional generation, allowing seamless adaptation to different biological design tasks.
• Employs OmegaFold for structure prediction and evaluation, using pLDDT, TM-score, RMSD, and Fident as key metrics to assess generated sequence quality.

Experimental results

Research questions

• RQ1 Can taxonomic information be effectively used to guide the conditional generation of biologically plausible protein sequences?
• RQ2 How does the combination of global and local attention mechanisms affect the structural and sequence-level quality of generated proteins?
• RQ3 To what extent does the patchify attention mechanism improve generation efficiency and structural consistency compared to standard self-attention?
• RQ4 How does TaxDiff compare to state-of-the-art models in both unconditional and taxonomic-guided protein sequence generation benchmarks?
• RQ5 Can TaxDiff generate protein sequences with higher structural confidence and lower inference time than existing diffusion-based models?

Key findings

• TaxDiff achieved a 7.13% improvement in Fident score and an 11.93% improvement in TM-score over baseline models in unconditional protein sequence generation, outperforming even structure-based models.
• In taxonomic-guided generation, TaxDiff's pLDDT scores were significantly higher than other sequence models and approached those of natural proteins, indicating high structural confidence.
• The model required only 24 minutes to generate 1,000 protein sequences, which is 1/4 to 2/3 the time of other diffusion-based models, demonstrating superior inference efficiency.
• The combination of global and local attention (Method A) achieved the best performance across all metrics, with local-only attention (Method E) also outperforming standard attention configurations.
• Dividing sequences into 16-amino acid patches yielded optimal results for structural metrics like TM-score and RMSD, while larger patches (64 residues) improved sequence-level metrics such as Fident.
• Empirical results confirmed that local attention is more effective than full-sequence attention for protein sequences, suggesting inherent local structural patterns akin to 'short sentences' in language.