PaperlayerView exact revision on arXiv

Source-linked AI summary

ProtTeX: Structure-In-Context Reasoning and Editing of Proteins with Large Language Models

Zicheng Ma, Chuanliu Fan, Zhicong Wang, Zhenyu Chen, Xiaohan Lin, Yanheng Li, Shihao Feng, Jun Zhang, Ziqiang Cao, Yi Qin Gao

arXiv:2503.08179v3q-bio.BMcs.AI

TL;DR

Protein LLM tokenization has largely relied on amino-acid sequences, although many protein challenges are structure-dependent. ProtTeX applies Chain-of-Thought reasoning in a framework for multimodal protein reasoning and generation. ProtTeX reports a twofold enhancement in accuracy for protein function prediction, and decoder-only LLMs can understand, predict, and design proteins.

  • Problem

    Protein LLM tokenization has largely relied on amino-acid sequences, although many protein challenges are structure-dependent.

    Source: §1 · pp. 2–3

  • Method

    ProtTeX applies Chain-of-Thought reasoning in a framework for multimodal protein reasoning and generation.

    Source: §1 · p. 4, §4 · p. 26

  • Results

    ProtTeX reports a twofold enhancement in accuracy for protein function prediction, and decoder-only LLMs can understand, predict, and design proteins.

    Source: §1 · p. 5

  • Takeaways & Limitations

    The authors conclude that decoder-only LLMs can address diverse protein-related tasks, including understanding, predicting, and designing proteins.

    Source: §1 · p. 5

    The authors suggest BiomedGPT’s training may be insufficient for comprehensive protein understanding.

    Source: §3.1 · p. 15

Abstract

from arXiv · show

Large language models have made remarkable progress in the field of molecular science, particularly in understanding and generating functional small molecules. This success is largely attributed to the effectiveness of molecular tokenization strategies. In protein science, the amino acid sequence serves as the sole tokenizer for LLMs. However, many fundamental challenges in protein science are inherently structure-dependent. The absence of structure-aware tokens significantly limits the capabilities of LLMs for comprehensive biomolecular comprehension and multimodal generation. To address these challenges, we introduce a novel framework, ProtTeX, which tokenizes the protein sequences, structures, and textual information into a unified discrete space. This innovative approach enables joint training of the LLM exclusively through the Next-Token Prediction paradigm, facilitating multimodal protein reasoning and generation. ProtTeX enables general LLMs to perceive and process protein structures through sequential text input, leverage structural information as intermediate reasoning components, and generate or manipulate structures via sequential text output. Experiments demonstrate that our model achieves significant improvements in protein function prediction, outperforming the state-of-the-art domain expert model with a twofold increase in accuracy. Our framework enables high-quality conformational generation and customizable protein design. For the first time, we demonstrate that by adopting the standard training and inference pipelines from the LLM domain, ProtTeX empowers decoder-only LLMs to effectively address diverse spectrum of protein-related tasks.

1. Introduction

Protein LLMs rely on sequence tokens even though functional inference often needs structural information, and many protein tasks have prompted specialized models. ProtTeX unifies sequence and structure tokenization for LLMs and reports improved protein-function prediction alongside structure generation and design capabilities.

1. Introduction

  • Protein science spans diverse biological questions that often require problem-specific models rather than a single unified framework.Source: §1 · p. 2
  • Protein tokenization largely uses one-letter amino-acid codes, which can be ambiguous with text, while functional inference often requires structural information.Source: §1 · p. 3
  • ProtTeX tokenizes protein sequences and 3D structures for LLMs, enabling protein-data formatting and editing and interleaving text with protein modalities.Source: §1 · p. 4
  • ProtTeX combines Chain-of-Thought reasoning with a unified model trained through next-token prediction, using structural or textual information as reasoning components.Source: §1 · p. 4
  • Compared with domain-expert LLMs BiomedGPT22 and Llama2-molinst-protein-7B19, ProtTeX reports a twofold enhancement in protein function prediction accuracy.Source: §1 · p. 5
  • The authors report high-accuracy structure prediction and conformational sampling, controllable protein generation from text prompts, and decoder-only LLMs’ ability to understand, predict, and design proteins.Source: §1 · p. 5

2. Methods

ProtTeX represents protein sequences, structures, and natural-language text as reversible discrete tokens and uses autoregressive transformers. Prompts arrange protein and text in different orders to support tasks from unimodal inference to multimodal Chain-of-Thought generation.

2. Methods

  • ProtTeX represents protein sequences, structures, and natural-language text as reversible discrete tokens for autoregressive transformers.Source: §2 · p. 5
  • Training prompts combine protein and text in arbitrary orders to support different downstream tasks.Source: §2 · p. 5
  • The prompts support tasks ranging from unimodal inference to multimodal Chain-of-Thought generation.Source: §2 · p. 5

2.1. Tokenizing All-Atom Protein Structures

ProtTeX converts protein structures into discrete residue tokens and represents sequence with added special tokens. Its structure-token pipeline combines an encoder, tokenizer, and decoder.

2.1. Tokenizing All-Atom Protein Structures

  • An SE(3)-invariant encoder maps protein structures to per-residue representations within a framework that tokenizes and reconstructs encoded structures.The framework comprises an encoder, tokenizer, and decoder.Source: §2.1 · pp. 5–6
  • Each residue representation is assigned to its nearest code among 512 codebook entries, and the code index becomes its LLM token.The codebook is initialized with 512 codes.Source: §2.1 · pp. 6–7
  • An SE(3)-equivariant sandwich-like decoder samples structures from the tokenized string’s metastable ensemble, while alignment and uniformity losses optimize the embedding space.The decoder updates single and pairwise representations.Source: §2.1 · pp. 6–7
  • The original amino acid sequence is represented with 20 added special tokens rather than abbreviations or compositional encoders.These tokens are added to the LLM vocabulary.Source: §2.1 · p. 7

2.2. Task-Unifying Prompts for Structure-In-Context Learning

ProtTeX combines protein sequences, structures, and text in document-like prompts, with modality order allowed to vary. These prompts let structures serve as inputs, outputs, or reasoning intermediates, and support task-specific extensions.

2.2. Task-Unifying Prompts for Structure-In-Context Learning

  • Training inputs combine protein sequences, structures, and natural-language prompts into document-like inputs using special tokens.The format includes separate protein-sequence and protein-structure blocks alongside question or description text.Source: §2.2 · p. 7
  • The three modalities can be arranged in different orders according to downstream tasks.The design is inspired by training on varied orderings of images and text.Source: §2.2 · p. 7
  • Four prompt templates support multimodal reasoning with structures as inputs, outputs, or chain-of-thought intermediates.Examples include generating a structure description from sequence before generating the structure, or generating a structure and then describing the protein using both modalities.Source: §2.2 · p. 8
  • The framework focuses on function prediction, structure generation, and controllable design, and supports exploring inverse folding or structure design with task-specific fine-tuning.The authors characterize it as supporting any-to-any modality transformation and a general foundational framework.Source: §2.2 · p. 8

2.3. Data and Models

This section describes ProtTeX’s data, task sampling, baselines, and answer-evaluation metrics, as well as an EMJI-computing approach for efficient, consistent evaluation and precise keyword matching.

2.3.1 Dataset

  • ProtTeX curated 3.36 million proteins with sequences and structures, processed them into sequence and structure tokens, and split them into training, validation, and test sets.The proteins were drawn from clustered AFDB v4, Swiss-Prot, and RCSB PDB, and were released before July 25, 2022.Source: §2.3.1 · p. 8
  • The main training dataset combines PFUD, PSAD, PDD, and PSPD, with QA-derived tasks sourced from Mol-Instruction [19] and ProteinLMBench [42].PFUD and PDD were adapted from Mol-Instruction, PSAD from ProteinLMBench, and proteins without corresponding QA pairs formed PSPD.Source: §2.3.1 · pp. 8–9

2.3.2 Training

  • Because protein sequences and structures are discrete tokens, ProtTeX trains with autoregressive next-token prediction and cross-entropy, using continuous pre-training and supervised fine-tuning of existing LLMs.Protein-modality and natural-language tokens receive equal weight.Source: §2.3.2 · pp. 9–10

2.3.3 Inference and sampling

  • ProtTeX uses greedy search for protein-understanding tasks, Beam Search with Lowest perplexity for structure analysis or prediction, and nucleus sampling for multi-conformation generation.Source: §2.3.3 · p. 10

2.3.4 Baseline Setups

  • ProtTeXLlama is fine-tuned on Meta-Llama-3-8B-Base with the ProtTeX tokenizer and combined task data, replacing PFUD’s textual sequences with tokenized sequences and structures.The comparisons include general, biomedical, and protein-focused baselines; prompts are standardized across models except BioMedGPT.Source: §2.3.4 · p. 12

2.3.5 Metrics

  • Because generative LLMs do not output class scores required for Fmax, the paper evaluates fluency with BLEU [50] and ROUGE [51] and biological-keyword overlap with Exact Match Jaccard Index (EMJI).DeepSeek-V3 [52] extracts keywords from reference labels and model outputs for exact matching.Source: §2.3.5 · pp. 12–13
  • The EMJI-computing approach supports efficient, consistent evaluation while maintaining high precision in keyword matching.Source: §2.3.5 · p. 13

3. Results and discussion

The results and discussion examine ProtTeX across protein-understanding tasks, structure-involved reasoning, and multi-conformation generation and protein design. They also report residue conservation in designed proteins and evidence relevant to controllable design and potential catalytic activity.

3.1. ProtTeX Enables Structure-In-Context Protein Understanding

  • ProtT3 performs comparably to ProtTeX on motif recognition and multi-attribute tasks, while ProtTeX outperforms competing approaches across other domains; BioMedGPT performs suboptimally.Approximately 2% of BioMedGPT’s responses were unknown answers.Source: §3.1 · p. 15, Figure 2 · p. 16
  • ProtTeXLlama3 achieves optimal PFUD performance in linguistic fluency and accuracy on domain-specific protein queries.The authors attribute this to ProtTeX and overlapping task subspaces that mutually enhance training.Source: §3.1 · p. 14, Table 2 · p. 16
  • ProtTeX’s tokenizer significantly outperforms models using one-letter amino-acid sequences or compositional protein encoders, supporting the benefit of structural tokens for protein-function comprehension.The comparisons include Llama2-molinst-protein-7B and Llama3-AAseq-FT, as well as ProtT3-FT and BioMedGPT.Source: §3.1 · p. 14, Table 4 · p. 17
  • Ablations indicate that pretraining without explicit functional information enhances functional comprehension, and comprehension and generation training improve one another.More comprehension data improves generative-task performance, while more generative data strengthens interpretative capabilities.Source: §3.1 · p. 15, Table 4 · p. 17

3.2. ProtTeX Enables Structure-Involved Reasoning for Proteins

  • ProtTeX supports multi-round multimodal reasoning that analyzes a sequence, generates a structure, and produces functional descriptions conditioned on the structure and prior descriptions.The framework also uses bidirectional protein-text generation to support cross-modal understanding and reasoning.Source: §3.2 · pp. 17–18, Figure 3 · p. 18
  • On the PFUD test-set prediction task using proteins shorter than 400 residues, CoT prompting improved exact-match accuracy by 49.6% over direct prompting.CoT first generates an intermediate reasoning step involving protein structure prediction from the input sequence, then uses both the original sequence and generated structure to predict function.Source: §3.2 · pp. 18–19
  • A statistically significant negative correlation between output perplexity and TM-score motivated Beam Search with Lowest PPL, which improves structure-prediction accuracy over greedy search.The model accomplishes structure prediction solely through next-token prediction.Source: §3.2 · p. 19, Figure 4 · p. 20
  • CoT prompting consistently outperforms direct prompting in structure-generation accuracy on the PSAD test set.The authors also report that this demonstrates transfer of reasoning techniques from natural language processing to biology.Source: Figure 4 · p. 20, §3.2 · p. 21
  • CoT makes model reasoning more transparent and enables it to explicitly connect protein structures with natural-language descriptions.The authors characterize this as an initial, promising demonstration using a limited dataset.Source: §3.2 · p. 21

sign

  • Generated protein sequences had similarity scores below 0.8 against the training set, while refolding showed high sequence–structure self-consistency and folds resembling natural enzymes despite low sequence identity.These results suggest designability, not experimentally verified function.Source: sign · p. 25
  • Interaction-involved side-chain amino acids are well conserved across designed and natural proteins, suggesting the designed proteins retain potential catalytic activity and functional integrity.The analysis compared co-crystallized PDB structures with corresponding AlphaFold3-predicted structures.Source: sign · p. 25
  • The results indicate potential for controllable protein design, and suggest fine-tuning LLMs could help them comprehend relationships between human and biological language.Source: sign · p. 25

4. Concluding Remarks

ProtTeX unifies protein understanding, structure generation, and design in a single-model, task-agnostic framework. It introduces multimodal Chain-of-Thought reasoning, while acknowledging performance gaps on some tasks compared with domain-specific models.

4. Concluding Remarks

  • ProtTeX unifies core protein challenges through a single-model, task-agnostic formulation with a unified-loss paradigm.The framework also uses a mixed-modal fusion strategy to process protein-related tasks.Source: §4 · pp. 25–26
  • The framework pioneers multimodal Chain-of-Thought reasoning for protein reasoning and generation, making deductive processes more transparent and interpretable.Source: §4 · p. 26
  • ProtTeX demonstrates capabilities across protein tasks, including protein understanding, structure generation, and protein design.Source: §4 · p. 26
  • The model may have slight performance gaps on some tasks compared with task-specific or domain-expert models such as ESMFold.Source: §4 · p. 26
  • The authors identify larger language models, more biological data, alignment, and inference-time self-improvement as future directions for improving ProtTeX.Source: §4 · pp. 26–27

A.2. Dataset

The dataset combines protein sequences and structural information with task-specific datasets for function understanding, structure analysis, design, and prediction. The authors describe how these datasets are curated, adapted, and partitioned for training and evaluation.

A.2. Dataset

  • The dataset comprises 3.36 million protein sequences paired with structural information, incorporating experimentally determined structures from RCSB PDB,41.The PDB component includes 551,957 single-chain structures released before October 13, 2021.Source: §A.2 · p. 35
  • The structural dataset is tokenized and filtered to proteins with reconstructed TM-score above 90, then split into 90% training, 5% validation, and 5% test sets.Questions without a matching accession in the database are excluded.Source: §A.2 · pp. 35–36
  • PSAD adapts proteinLM-Bench prompts to generate structures after structural analysis, while PDD adds structure generation after sequence generation to support multimodal CoT.Because PDD data are limited, its sequences are restricted to sequence-generation tasks without structure prediction, and the dataset remains relatively small.Source: §A.2 · pp. 36–37
  • PSPD maps protein sequences to structures using AFDB and Swiss-Prot structures while excluding accessions used to train the three previously described datasets.Its data sources include 2.27 million AFDB single-chain structures and 541,327 curated Swiss-Prot single-chain structures.Source: §A.2 · pp. 35, 37

A.3. Training Details

The authors compare full-parameter fine-tuning with LoRA and select full-parameter fine-tuning after finding that LoRA underperforms on task-specific metrics. They report the optimizer, learning-rate schedule, batch size, training duration, and hardware.

A.3. Training Details

  • LoRA significantly underperforms full-parameter fine-tuning on task-specific metrics, so the authors use full-parameter fine-tuning as their primary strategy.Source: §A.3 · p. 37
  • For parameter fine-tuning, training uses AdamW with weight decay 0.1 and a cosine annealing learning-rate schedule that gradually decreases the rate from 5e-6 to 1e-7.Source: §A.3 · p. 38
  • The model is trained for four epochs with batch size 3 per device, requiring about five days on 16 NVIDIA A100 GPUs.Source: §A.3 · p. 38

A.4. Ablation details

The ablations examine training data, training strategy, and model scale. They report improved functional comprehension with sequence-structure pretraining, while comparing fine-tuning strategies and model sizes.

Ablation on different training data

  • On PFUD, the ablations fine-tune the base model and test sequence-structure pretraining; the final model is supervised across downstream tasks, including PSPD, without redundant PSPD pretraining.Source: §A.4 · p. 38
  • A single epoch of sequence-structure pretraining significantly improves performance on functional comprehension tasks.The analysis reports this finding in Table 4 and Supplementary Figure S2.Source: §A.4 · p. 38

Ablation details

  • Token representations provide critical inductive biases for biological function prediction, and different tasks can mutually reinforce one another without explicit correlations.Source: §A.4 · p. 38

Ablation on model scaling

  • The 8B model slightly outperforms the 1B model, but the gap is relatively small given the limited dataset; 8B is selected for expanded data and downstream computational scalability.Source: §A.4 · p. 39

Ablation on training strategy

  • LoRA converges significantly more slowly than full-parameter fine-tuning and, with DeepSpeed Zero Stage 3, offers no significant multi-GPU memory or training-speed advantage; full tuning is therefore standard.Source: §A.4 · p. 39, Figure S2 · p. 40
Loading 2503.08179v3…