Table of Contents
Fetching ...

Reflection Pretraining Enables Token-Level Self-Correction in Biological Sequence Models

Xiang Zhang, Jiaqi Wei, Yuejin Yang, Zijie Qiu, Yuhan Chen, Zhiqiang Gao, Muhammad Abdul-Mageed, Laks V. S. Lakshmanan, Wanli Ouyang, Chenyu You, Siqi Sun

TL;DR

Bio sequence languages have limited expressiveness, quantified by $|\mathbb{S}_{\bm{L}}|$, which restricts Chain-of-Thought (CoT) reasoning. The authors propose reflection-based pretraining to expand the protein vocabulary to $|\mathbb{S}_{\bm{L}_{protein+}}|$ by adding thinking tokens and training with error-injection strategies (RPRE and RPLE) plus online dynamic perturbations and gradient blocking. Applied to de novo peptide sequencing on MassIVE-KB data, this approach yields substantial gains in amino acid and peptide precision, enables self-correction, and supports human-in-the-loop guidance, while finetuning alone is insufficient to elicit reasoning. Overall, the method narrows the gap between biological and natural language models by embedding explicit intermediate reasoning in biological sequence generation, with considerations for safety, ethics, and future refinements.

Abstract

Chain-of-Thought (CoT) prompting has significantly advanced task-solving capabilities in natural language processing with large language models. Unlike standard prompting, CoT encourages the model to generate intermediate reasoning steps, non-answer tokens, that help guide the model toward more accurate final outputs. These intermediate steps enable more complex reasoning processes such as error correction, memory management, future planning, and self-reflection. However, applying CoT to non-natural language domains, such as protein and RNA language models, is not yet possible, primarily due to the limited expressiveness of their token spaces (e.g., amino acid tokens). In this work, we propose and define the concept of language expressiveness: the ability of a given language, using its tokens and grammar, to encode information. We show that the limited expressiveness of protein language severely restricts the applicability of CoT-style reasoning. To overcome this, we introduce reflection pretraining, for the first time in a biological sequence model, which enables the model to engage in intermediate reasoning through the generation of auxiliary "thinking tokens" beyond simple answer tokens. Theoretically, we demonstrate that our augmented token set significantly enhances biological language expressiveness, thereby improving the overall reasoning capacity of the model. Experimentally, our pretraining approach teaches protein models to self-correct and leads to substantial performance gains compared to standard pretraining.

Reflection Pretraining Enables Token-Level Self-Correction in Biological Sequence Models

TL;DR

Bio sequence languages have limited expressiveness, quantified by , which restricts Chain-of-Thought (CoT) reasoning. The authors propose reflection-based pretraining to expand the protein vocabulary to by adding thinking tokens and training with error-injection strategies (RPRE and RPLE) plus online dynamic perturbations and gradient blocking. Applied to de novo peptide sequencing on MassIVE-KB data, this approach yields substantial gains in amino acid and peptide precision, enables self-correction, and supports human-in-the-loop guidance, while finetuning alone is insufficient to elicit reasoning. Overall, the method narrows the gap between biological and natural language models by embedding explicit intermediate reasoning in biological sequence generation, with considerations for safety, ethics, and future refinements.

Abstract

Chain-of-Thought (CoT) prompting has significantly advanced task-solving capabilities in natural language processing with large language models. Unlike standard prompting, CoT encourages the model to generate intermediate reasoning steps, non-answer tokens, that help guide the model toward more accurate final outputs. These intermediate steps enable more complex reasoning processes such as error correction, memory management, future planning, and self-reflection. However, applying CoT to non-natural language domains, such as protein and RNA language models, is not yet possible, primarily due to the limited expressiveness of their token spaces (e.g., amino acid tokens). In this work, we propose and define the concept of language expressiveness: the ability of a given language, using its tokens and grammar, to encode information. We show that the limited expressiveness of protein language severely restricts the applicability of CoT-style reasoning. To overcome this, we introduce reflection pretraining, for the first time in a biological sequence model, which enables the model to engage in intermediate reasoning through the generation of auxiliary "thinking tokens" beyond simple answer tokens. Theoretically, we demonstrate that our augmented token set significantly enhances biological language expressiveness, thereby improving the overall reasoning capacity of the model. Experimentally, our pretraining approach teaches protein models to self-correct and leads to substantial performance gains compared to standard pretraining.
Paper Structure (37 sections, 7 equations, 7 figures, 4 tables, 1 algorithm)

This paper contains 37 sections, 7 equations, 7 figures, 4 tables, 1 algorithm.

Figures (7)

  • Figure 1: Expressiveness of Languages.
  • Figure 2: Comparative framework for prompting in natural and biological language models.
  • Figure 3: De novo peptide sequencing workflow using tandem mass spectrometry.
  • Figure 4: Error injection and reflection training for augmenting reasoning in bio. sequence models.
  • Figure 5: Error Injection Rate on Validation Loss. Loss Plotting and Logging by Neptune.ai.
  • ...and 2 more figures