Source-linked AI summary
SoftGene: Protein Language Model-Enhanced Soft Prompting for Interpretable Gene Set Annotation
Drew Ross, Arya Hadizadeh Moghaddam, Dongjie Wang, Xiaoyu Zhang, Zijun Yao
TL;DR
Gene set annotation relies on time-consuming manual curation, while existing LLM approaches often overlook protein structure. SoftGene combines protein-sequence representations and textual context through hierarchical attention-based soft prompting. It outperforms prompt-based, agent-based, and prompt-tuning baselines across reported metrics, while the contribution of sequence representations varies across biological domains.
Problem
Manual gene set annotation is time-consuming, curated databases can be incomplete or biased toward well-studied processes, and LLM approaches often overlook protein structure.
Source: §1 · p. 1Method
SoftGene uses a hierarchical attention-based protein-sequence encoder to condition an LLM with sequence-derived soft prompts and textual context for gene set annotation.
Source: §1 · p. 2Results
SoftGene outperforms prompt-based, agent-based, and prompt-tuning baselines across reported metrics on Gene Ontology and MSigDB, while protein-sequence contributions vary across biological domains.
Source: §5 · p. 9Takeaways & Limitations
The findings support sequence-aware soft prompting as a promising approach for interpretable gene set annotation and highlight the importance of domain-specific evaluation and biological validation.
Source: §5 · p. 9SoftGene is trained and evaluated on curated resources biased toward well-studied genes, and its reference-based evaluations do not fully capture biological correctness or experimental validity.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Gene set analysis is a cornerstone of functional genomics, yet it remains labor-intensive and heavily dependent on manual curation and expert biological interpretation. While Large Language Models (LLMs) have emerged as powerful tools for genomic reasoning and annotation, most existing approaches rely on symbolic gene names and fail to capture domain-specific biological structure, particularly protein sequence information that governs molecular activity, interactions, and downstream gene function. In this work, we propose SoftGene, a novel framework for LLM-based gene set annotation that leverages the hierarchical structure of gene sets. First, we use a hierarchical attention-based encoder built on ESM, a protein language model, to represent each gene set using protein-level amino acid sequence information. Second, we construct a hybrid prompting scheme that combines soft prompts derived from gene set embeddings with hard prompts containing auxiliary context generated by an LLM, and feed the resulting prompt into a local LLM for annotation. We evaluate our framework on two benchmark datasets: Gene Ontology (GO) and the Molecular Signatures Database (MSigDB). Our results show that integrating protein-sequence representations with textual context improves gene set annotation overall, while per-domain analyses reveal that the contribution of protein embeddings varies across biological domains.
1 Introduction
Gene set annotation remains labor-intensive, while current LLM approaches often overlook protein structure and lack a concrete path showing how genes contribute to annotations. SoftGene combines protein-sequence representations with textual context and uses hierarchical attention for feature-level explanations.
1 Introduction
- Manual annotation is time-consuming, and automated methods relying on curated databases are limited by incomplete coverage and bias toward well-studied genes and processes.Gene set annotations describe shared functional behaviors and metabolic roles, supporting the translation of raw data into biological insights.Source: §1 · p. 1
- Current LLM-based approaches often overlook constituent-protein structure, and hard prompting treats genes as semantic tokens while LLMs remain prone to hallucination.Protein-sequence representations can provide evolutionary and structural signals that may complement symbolic gene information; their utility can vary across biological domains.Source: §1 · p. 1
- SoftGene reports superior performance over existing baselines across metrics on the Gene Ontology and MSigDB benchmarks.The framework uses a hierarchical attention encoder over protein-language-model embeddings and conditions an LLM on sequence-derived representations and text context.Source: §1 · p. 2
- A hierarchical attention encoder aggregates pretrained protein-language-model embeddings to represent gene sets and weigh biological features.SoftGene conditions an LLM using both protein-sequence-derived representations and text-based context.Source: §1 · p. 2, Figure 1 · p. 2
- The hierarchical attention mechanism provides feature-level explanations that quantify the relative importance of biological features contributing to each annotation.These explanations are grounded in input data rather than relying solely on implicit language-model reasoning.Source: §1 · p. 2
2 Methodology
SoftGene frames gene-set annotation as conditional generation of concise functional text from gene sets represented through their constituent proteins. It combines hierarchical protein-sequence embeddings with soft prompts and textual context to condition a local language model.
2.1 Problem Formulation
- A gene set contains genes associated with shared function, pathways, or co-regulation, and the task is to predict a concise annotation describing its biological role.Gene sets may contain variable numbers of genes, and genes may encode multiple protein isoforms.Source: §2.1 · pp. 2–3
2.2 Model Overview
- SoftGene combines a hierarchical protein-based gene set encoder with a hybrid prompting module that conditions a local language model to generate free-text annotations.The model uses a frozen pretrained protein language model and a frozen language model; trainable attention pooling and a prompt-projector MLP provide the learned components.Source: §2.2 · p. 3, §2.5 · p. 5
2.3 Hierarchical Attention-Based Gene Set Encoder
- Protein sequences are encoded with ESM-2 and mean-pooled, then attention pooling aggregates isoform contributions into gene representations and gene contributions into a fixed-size gene-set representation.The learned attention pooling can differentially weight splice isoforms, downweighting noisy or irrelevant ones.Source: §2.3 · pp. 3–4
2.4 Hybrid Gene Set Soft Prompting
- The gene-set embedding is projected into continuous soft prompt tokens and combined with an instruction prompt and LLM-generated gene-set context for annotation generation.The context comes from a commercial pretrained language model and supplies knowledge that the smaller local model may not encode in its pretrained weights.Source: §2.4 · pp. 4–5, Figure 4 · p. 4
2.5 Training Procedure
- Training uses teacher forcing and token-level cross-entropy, while keeping the protein and language models frozen and updating only the prompt-projector MLP and hierarchical attention encoder.At each time step, the model conditions on the ground-truth output prefix.Source: §2.5 · p. 5
3 Experiments
Experiments compare SoftGene with prompting, prompt-tuning, and agentic baselines on GO and MSigDB, then examine its components, generalization, domain behavior, and interpretability. SoftGene outperforms the baselines under the reported evaluation, while analyses show that protein-embedding contributions vary across domains.
3.1 Datasets
- The experiments use GO and MSigDB gene sets, a 70/30 random split, and mean results with 95% confidence intervals over three runs.Evaluation uses METEOR, ROUGE, and BERTScore to assess lexical and semantic similarity.Source: §3.1 · p. 5, §3.2 · p. 5, §3.3 · p. 5
3.4.1 Baseline Comparison
- SoftGene outperforms prompting, prompt-tuning, and agentic baselines across the reported metrics on GO and MSigDB.The comparisons include zero- and few-shot prompting, P-Tuning and Prefix Tuning, and systems using external knowledge sources.Source: Table 2 · p. 6, §3.4.1 · p. 6
3.4.2 Ablation Study
- Removing soft prompts, LLM-generated context, or pretrained protein embeddings degrades performance, although the size of the embedding contribution varies by domain.SoftGene without pretrained protein embeddings still outperforms all baselines, supporting the contribution of multiple components rather than PLM embeddings alone.Source: §3.4.2 · p. 6, Table 3 · p. 7
3.4.3 Soft Prompt Length
- Across GO and MSigDB, performance varies only modestly across evaluated prompt lengths, so 10 tokens are selected to reduce computational cost without significant performance gains from longer prompts.Figures 5 and 6 show the prompt-length comparisons for GO and MSigDB.Source: §3.4.3 · p. 7, Figure 5 · p. 7, Figure 6 · p. 7
3.4.4 Gene-Disjoint Generalization
- Under a gene-disjoint split, SoftGene performs comparably to the random split across metrics on both datasets, with limited overall degradation.The stricter split tests generalization to genes excluded from training as much as possible.Source: §3.4.4 · p. 7, Table 4 · p. 8
3.4.5 Per-Domain PLM Embedding Ablation
- Protein embeddings yield larger improvements in GO Cellular Component than in Biological Process or Molecular Function.The per-domain ablation compares real PLM embeddings with randomized embeddings.Source: §3.4.5 · p. 8, Table 5 · p. 8
3.4.6 Case Verification: Zinc Ion Transmembrane Transport
- For zinc ion transmembrane transport, the highest attention weights go to known zinc transporter genes, while the less directly related SLC1A1 receives lower weight.The authors interpret this pattern as evidence that gene-level attention ranks genes by functional relevance rather than weighting all set members equally.Source: §3.4.6 · p. 8
3.4.7 Interpretability: Smooth Muscle Signaling
- For negative regulation of smooth muscle contraction, SoftGene assigns high attention to CALCA, ADRB2, ADRB3, and the PRKG1α isoform.The authors report that the isoform preference emerges without isoform-level supervision and suggest protein embeddings encode signals distinguishing functional splice variants.Source: §3.4.7 · p. 8
4 Related Works
Related work covers LLM-based gene set annotation and soft prompting. The discussion positions SoftGene among methods that combine language-model prompting with biological information.
Automatic Gene Set Annotation
- Gene set annotation approaches include LLMs, retrieval and database grounding, and self-verification to address factual consistency and hallucinated associations.The cited work describes TALISMAN, GeneAgent, and Discovera as examples.Source: §4 · p. 8
LLM Soft Prompting
- Figure 7 visualizes a three-level gene-set encoder hierarchy and the top five protein isoforms under each gene, with node size indicating learned attention weights.The visualization is for GO:0045986.Source: Figure 7 · p. 9
- Soft prompting steers frozen language models with learned vectors, with prior work extending the approach across input layers, deeper layers, and biological tasks.The review includes Prefix-Tuning, Prompt Tuning, P-Tuning, P-Tuning v2, and biological applications.Source: §4 · p. 9
5 Conclusion
SoftGene combines protein-sequence encoding with soft-prompt conditioning for gene set annotation. Across GO and MSigDB, it outperforms the reported baselines, while analyses find that protein-sequence contributions vary by biological domain.
5 Conclusion
- SoftGene integrates a hierarchical attention-based encoder over protein sequences with soft-prompt conditioning for language-model-driven gene set annotation.The conclusion describes this as a hybrid prompting framework.Source: §5 · p. 9
- On GO and MSigDB, SoftGene outperforms the reported prompt-based, agent-based, and prompt-tuning baselines across metrics, while protein-sequence contributions vary across domains.The authors conclude that sequence-aware soft prompting is promising and highlight domain-specific evaluation and biological validation.Source: §5 · p. 9
6 Ethical Considerations
Generated annotations are intended as hypothesis-generating tools, not definitive functional claims, and should not replace expert judgment or experimental validation, especially in high-stakes biomedical or clinical settings.
6 Ethical Considerations
- SoftGene should not replace expert biological judgment or experimental validation, particularly in high-stakes biomedical or clinical settings.The authors recommend treating annotations as hypothesis-generating rather than definitive functional claims.Source: §6 · p. 9
7 Limitations
SoftGene’s appendix analyses identify limits in dataset coverage and show that performance patterns vary across context models, gene-set properties, and biological domains. Protein-sequence contributions differ by GO domain, while gene uniqueness has limited overall effect and GO performance rises with ontology depth.
7 Limitations
- Curated benchmarks biased toward well-studied genes may limit generalization to poorly annotated or emerging functions; reference-based metrics also do not establish biological correctness or experimental validity.The authors call for evaluation on less curated gene sets, expert assessment, and downstream biological validation.Source: §7 · p. 9
A.1 LLM Backbone for Context Generation
- DeepSeek and GPT contexts yield the strongest overall scores across metrics, with DeepSeek highest on average; Claude is comparable but slightly lower, while Gemma underperforms.Figures 8 and 9 compare commercial context-generation models on GO and MSigDB.Source: §A.1 · p. 13, Figure 8 · p. 14, Figure 9 · p. 14
A.2 Encoder Protein Language Model
- Performance varies minimally across ESM-2 checkpoints and ProtBERT, suggesting SoftGene is not sensitive to the protein-embedding backbone.The comparisons cover GO and MSigDB and report minimal variation across random seeds.Source: §A.2 · p. 13, Figure 10 · p. 15, Figure 11 · p. 15
A.3 Performance by Description Length
- GO BERTScore F1 is stable for 1–6-word descriptions and modestly higher for descriptions of 7+ words, whereas MSigDB scores peak at 1–2 words and decline with longer descriptions.The authors suggest annotation complexity relates to performance differently across the two databases.Source: §A.3 · pp. 13–14, Figure 12 · p. 16
A.4 Performance by Gene Set Cardinality
- GO BERTScore F1 rises with gene-set size except for sets with 100+ genes, while MSigDB performance declines monotonically across cardinality bins.The authors note that annotation composition and difficulty may also contribute to these trends.Source: §A.4 · p. 14, Figure 12 · p. 16
A.5 Performance by Domain and Collection
- GO molecular-function annotations score highest and cellular-component annotations lowest, but protein embeddings yield negligible gains for biological process, modest gains for molecular function, and the largest gains for cellular component.In MSigDB, PID performs strongest despite n=64; BIOCARTA, KEGG LEGACY, and REACTOME are comparable and slightly lower, with higher variance.Source: §A.5 · pp. 14–15, Figure 12 · p. 16
A.6 Performance by Gene Uniqueness
- GO median performance gradually declines from the most unique genes (Q1) to the least unique (Q5), while MSigDB remains largely stable through Q4 and rises slightly at Q5.Across both datasets, gene uniqueness has a limited effect relative to the broader performance range.Source: §A.6 · pp. 15–16, Figure 12 · p. 16
A.7 Performance by Ontology Depth
- GO BERTScore tends to increase with ontology depth, with higher scores for more specific terms than for broad categories near the ontology root.The authors suggest specific terms may have more distinct, well-defined lexical profiles in training data.Source: §A.7 · p. 16
A.8 Additional Performance Metrics
SoftGene performs strongest on the supplied biologically grounded metrics, although its margins are smaller than on lexical metrics. The additional examples illustrate attention to pathway-relevant genes and show that low-scoring predictions often involve complex regulation or nuanced phenotypes.
A.8 Additional Performance Metrics
- SoftGene achieves the strongest scores on the biologically grounded metrics, but its margins over competing methods are smaller than for lexical metrics, making these complementary rather than equally strong evidence of biological-content gains.HSS is computed only for GO because MSigDB lacks an explicit term hierarchy.Source: §A.8 · pp. 16–17, Table 6 · p. 18
A.9.1 Regulation of Megakaryocyte Differentiation
- For regulation of megakaryocyte differentiation, attention is highest on MEF2C, intermediate on RBM15, and among the lowest on LOX and EIF6, matching their described pathway relevance.The authors interpret this as the encoder distinguishing central pathway genes from genes with more general cellular functions.Source: §A.9.1 · p. 18
A.9.2 Hypoxia-Inducible Factor in the Cardiovascular System
- For the HIF pathway, attention is highest on HIF1A, VHL, and VEGFA, while broadly expressed genes such as HSP90AA1 and P4HB receive lower weights.The authors characterize these as core pathway regulators and outputs versus genes with more general cellular functions.Source: §A.9.2 · p. 18
A.9.3 Glyoxylate and Dicarboxylate Metabolism
- For glyoxylate and dicarboxylate metabolism, attention is elevated on CS, ACO1/ACO2, MDH1/MDH2, HAO1, HAO2, and GRHPR, rather than being assigned uniformly across genes.The authors relate these weights to central carbon metabolism, interconnected metabolic processes, and glyoxylate metabolism or detoxification.Source: §A.9.3 · pp. 18–19
B.1 High- and Low-Scoring Completions
- High-scoring examples capture enzymatic activities and metabolic pathways, while lower-scoring examples often involve complex regulatory processes or nuanced phenotypic descriptions.Tables 7 and 8 show top- and bottom-scoring predictions for GO and MSigDB, respectively.Source: §B.1 · p. 19
B.2 Failure Mode Analysis
The failure analysis identifies recurring, interpretable errors rather than random mistakes, including missed substrate specificity, pathway pleiotropy, and semantic-granularity mismatches. The accompanying implementation notes specify the software, generation, and baseline settings used in the experiments.
B.2 Failure Mode Analysis
- Low-scoring predictions reveal systematic, interpretable failure patterns that can expose limitations in learned representations and in evaluation metrics.The authors introduce three recurring failure modes.Source: §B.2 · p. 19
Substrate Specificity Failure
- The sodium:sugar symporter prediction captures the transporter mechanism but misses the precise substrate, which the authors suggest may reflect smoothing of subtle active-site signals.The example contrasts the prediction with the reference “fucose transmembrane transport.”Source: Substrate Specificity Failure · p. 19
Pathway Pleiotropy Failure
- The cell-substrate adhesion prediction instead of fibroblast growth factor receptor signaling reflects a functionally coupled or downstream process; the authors interpret it as a broader pathway signature.They suggest the representation may capture functional relationships among cellular processes.Source: Pathway Pleiotropy Failure · p. 19
Semantic Granularity Failure
- A predicted synonym or parent term can be biologically valid yet receive a low ROUGE-L score when it differs from the reference’s wording or granularity.For the circadian example, the prediction has high BERTScore despite low ROUGE-L.Source: Semantic Granularity Failure · p. 19
C Implementation Details
- Local LLM inference and training use Transformers, distributed training uses Accelerate, and evaluation uses Evaluate; BERTScore uses DistilBERT.The implementation ran on three NVIDIA A6000 GPUs, 512 GB RAM, and a 48-core CPU.Source: C Implementation Details · p. 19
C.1 Hyperparameters; Prompt Engineering Methods; Prompt-Tuning Methods
- Auxiliary context generation uses default generation settings except for a 100-token maximum; prompt-engineering baselines use Llama-3.1-8B-Instruct via Replicate’s default configuration.Prompt-tuning baselines use base Llama-3.1-8B with 10 virtual tokens and omit commercial LLM context; they train for one GO epoch and two MSigDB epochs.Source: §C.1 · p. 19, Prompt Engineering Methods · p. 20, Prompt-Tuning Methods · p. 20
Agentic Systems
- Agentic baselines use GPT-4o-mini, with GeneAgent limited to three self-verification iterations and Discovera to 10 genes passed to INDRA to constrain inference costs.Source: Agentic Systems · p. 20
C.3 Runtime Analysis
SoftGene precomputes and caches protein embeddings, leaving a sparse aggregation and small projection in the runtime path alongside LLM inference. Its query-only attention and batched sparse aggregation keep large gene sets computationally tractable in the reported setting; the appendix also describes the source and construction of both benchmark datasets.
C.3 Runtime Analysis
- Protein embeddings are computed once over the human proteome and cached; the encoder’s sparse scatter aggregation and small MLP projection are negligible relative to the LLM forward pass.For novel gene sets, an additional commercial LLM inference pass generates auxiliary context, but its variable network latency is excluded from runtime reporting.Source: §C.3 · pp. 20–21, Table 9 · p. 21, Table 10 · p. 21
C.3.1 Scalability to Large Gene Sets
- Query-only attention runs in O(n) time and avoids pairwise attention’s quadratic memory bottleneck; sparse scatter operations batch variable-length sets, and the five largest evaluated GO sets remain tractable.The passage notes GO sets can contain up to 13,616 genes; Table 12 reports runtime for the five largest.Source: §C.3.1 · p. 21, Table 12 · p. 21
D Dataset Preparation
- The dataset-preparation section outlines how the benchmark inputs are prepared before detailing protein-sequence and dataset-specific processing.Source: D Dataset Preparation · p. 21
D.1 Protein Sequences and Isoforms
- Ensembl’s human proteome supplies protein sequences and gene-to-isoform mappings, providing the gene-to-protein hierarchy used to construct GO and MSigDB input representations.Source: §D.1 · p. 21
D.2 GO Dataset
- GO gene sets are built from the GO ontology and human gene annotation files, collecting genes by term and retaining each term’s name and namespace.The resulting JSON maps each GO ID to its genes, description, and domain.Source: §D.2 · p. 21
D.3 MSigDB Dataset
- MSigDB C2:CP sets are assembled from the GMT file and SQLite metadata database, filtered to sets present in both sources and exported with cleaned descriptions, gene symbols, and collection names.Some sub-collections are excluded to focus on canonical pathway sets.Source: §D.3 · p. 22
D.4 Dataset Contamination Analysis
The appendix examines possible dataset contamination and leakage through auxiliary context, then documents SoftGene and baseline prompts alongside protein-language-model background. It reports that recent releases reduce—but do not eliminate—contamination risk, while context-free SoftGene remains competitive and generated context has negligible lexical overlap with annotations.
D.4 Dataset Contamination Analysis
- For GPT-4.1 nano, its documented knowledge cutoff predates the selected recent GO and MSigDB releases, reducing but not eliminating contamination risk because underlying genes, pathways, and annotations may predate them.Source: §D.4 · p. 22
- SoftGene without auxiliary context still outperforms all competitive baselines across both datasets, while GPT-generated context has negligible lexical overlap with ground-truth annotations.The lexical-overlap comparison is reported in Table 13.Source: §D.4 · p. 22, Table 13 · p. 22
E.1 SoftGene Prompt
- SoftGene’s concise instruction prompt leaves task-relevant information to be learned during training, allowing the soft prompt to adapt to gene prediction.Source: §E.1 · p. 22
E.2 Zero-Shot Prompt
- The zero-shot baseline supplies task instructions without examples, leaving the model to rely on pretraining knowledge.Source: §E.2 · p. 22
E.3 Few-Shot Prompt
- The few-shot baseline pairs instructions with a small number of examples to convey the desired output format and style.Source: §E.3 · p. 22
E.4 One-Shot Chain-of-Thought Prompt
- The one-shot chain-of-thought baseline uses a single worked example to guide step-by-step analysis before the final output.Source: §E.4 · p. 23
F.1 Protein Language Models
- Protein language models learn protein structure and function from sequence data, with cited approaches spanning single-sequence, co-evolutionary, and multimodal representations.The discussion includes ESM, ProtTrans, MSA Transformer, and ESM3.Source: §F.1 · p. 23
F.2 Traditional Gene Set Annotation Methods
Traditional gene set annotation methods analyze gene membership statistically or incorporate pathway structure, but generally produce categorical labels rather than descriptive natural-language annotations. The reviewed methods also do not directly use protein-sequence representations.
F.2 Traditional Gene Set Annotation Methods
- Over-representation methods identify statistically enriched GO terms, while GSEA uses ranked gene lists and a weighted Kolmogorov-Smirnov statistic to detect coordinated activity shifts.Source: §F.2 · p. 23
- GSEA enrichment scores are cohort-level, whereas later approaches compute per-sample scores or generalize enrichment across microarray and RNA-Seq platforms.Source: §F.2 · p. 23
- Pathway-topology methods weight genes by connectivity in curated signaling graphs, but remain dependent on database completeness.Source: §F.2 · p. 23
- The reviewed methods map gene sets to predefined categorical labels rather than generating natural-language descriptions, and do not directly incorporate protein-sequence representations.Source: §F.2 · p. 23