Inference-Time Toxicity Mitigation in Protein Language Models

Manuel Fernández Burda; Santiago Aranguri; Iván Arcuschin Moreno; Enzo Ferrante

cs.LGMar 4, 2026

Inference-Time Toxicity Mitigation in Protein Language Models

Authors: Manuel Fernández Burda, Santiago Aranguri, Iván Arcuschin Moreno, Enzo Ferrante

Organizations: Laboratory of Applied Artificial Intelligence (LIAA), Institute of Computer Sciences (ICC), CONICET - Universidad de Buenos · AI Safety Argentina (AISAR) · Goodfire · APOLO Biotech

Abstract

Protein language models (PLMs) are becoming practical tools for de novo protein design, yet their dual-use potential raises safety concerns. We show that domain adaptation to specific taxonomic groups can elicit toxic protein generation, even when toxicity is not the training objective. To address this, we adapt Logit Diff Amplification (LDA) as an inference-time control mechanism for PLMs. LDA modifies token probabilities by amplifying the logit difference between a baseline model and a toxicity-finetuned model, requiring no retraining. Across four taxonomic groups, LDA consistently reduces predicted toxicity rate (measured via ToxDL2) below the taxon-finetuned baseline while preserving biological plausibility. We evaluate quality using Fréchet ESM Distance and predicted foldability (pLDDT), finding that LDA maintains distributional similarity to natural proteins and structural viability (unlike activation-based steering methods that tend to degrade sequence properties). Our results demonstrate that LDA provides a practical safety knob for protein generators that mitigates elicited toxicity while retaining generative quality.

Figures & tables

Figure 1: LDA reduces predicted toxicity across taxa. Percentage of generated sequences classified as toxic by ToxDL2 (lower is better) for four taxonomic finetunes. Baseline denotes the corresponding taxon finetune, Toxic denotes models finetuned on toxin-enriched data from within that taxon and LDA denotes the intervened models at inference time. Bars report mean ± s.e.m. across three independent generation runs under identical sampling and perplexity filtering.

Taxon	$\bm{\Delta}$ FED $\downarrow$	$\bm{\Delta}$ pLDDT $\uparrow$
Arthropoda	$+0.03$	$+1.59\pm 21.20$
Gastropoda	$-0.30$	$+0.10\pm 11.78$
Lepidosauria	$-0.26$	$-6.95\pm 23.61$
Arachnida	$-0.09$	$-1.55\pm 12.69$

Table 1: Biological quality under LDA at optimal (least predicted toxicity)

\alpha

values. Changes in Frechet ESM Distance (lower is better) and pLDDT (higher is better) for LDA-mitigated generations relative to the taxon baseline.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Supplementary Figure 1: Taxon finetuning elicits toxic generation. Toxicity rates for baseline ProGen2 versus taxon-finetuned models across four taxonomic groups. Error bars show $\pm$ 1 standard deviation.

Taxa	#Toxic Proteins	#Non-Toxic Proteins	% Toxicity
Arthropoda	2644	9508	21.76%
Arachnida	2136	534	80.00%
Lepidosauria	1830	652	73.73%
Gastropoda	1306	284	82.14%

Appendix

Supplementary Table 1: Taxonomic fine-tuning datasets statistics.

Supplementary Figure 2: Steering intensity unveils mitigation regimes. Toxicity rate versus amplification strength $\alpha$ for (a) Arthropoda (log scale), (b) Gastropoda, (c) Lepidosauria, and (d) Arachnida. Dashed lines indicate taxon-finetuned baseline (gray) and toxic-finetuned model (red). For all taxa, there exists an $\alpha$ range where toxicity drops below baseline.

Supplementary Figure 3: LDA maintains sequence quality. $\Delta$ FED (left) and $\Delta$ pLDDT (right) versus $\alpha$ for LDA across taxa.

Supplementary Figure 4: Linear probing reveals toxicity encoding across layers. Classification metrics (Accuracy, AUC-ROC, F1) for linear probes trained on layer-wise activations. Performance increases with depth, indicating toxicity-related information emerges gradually and becomes linearly accessible in intermediate-to-final layers. Metrics statistics are calculated over 5 random species-stratified splits.

Supplementary Figure 5: Effect of linear steering on generative toxicity in ProGen2. Comparison between Direct Steering (a, b) and Affine Steering (c, d) under interventions applied to all layers (All Layers, first column) and layer-by-layer (Layer by Layer, second column). The vertical axis indicates the percentage of toxic sequences generated (logarithmic scale). The dashed black line shows baseline model performance ( $\pm$ sd). Colors represent steering intensity (red: addition/ $\alpha$ =1.0; blue: ablation/ $\alpha$ =0.0; intermediate tones: $\alpha$ = $\pm$ 0.5).

Supplementary Figure 6: Direct steering degrades biological quality. $\Delta$ FED (upper) and $\Delta$ pLDDT (bottom) for direct steering interventions.

Supplementary Figure 7: Affine steering degrades biological quality. $\Delta$ FED (upper) and $\Delta$ pLDDT (bottom) for affine steering interventions.

Supplementary Figure 8: ESMFold as structure predictor for ToxDL2. Performance comparison of ToxDL2 using ESMFold-predicted structures versus original AlphaFold2/ColabFold structures on the authors’ benchmark datasets. Left: held-out test set (1579/1710 sequences within ESMFold context). Right: independent test set (4099/4480 sequences). Metrics used on the original paper where mirrored: Accuracy, AUC-ROC, F-max and Matthew’s Correlation Coefficient.

Protein Language Models

Protein language models learn patterns in amino-acid sequences to predict protein properties, interpret mutations, and design new proteins.

See all 69 papers

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference
Feb 24, 2026Anna Hart, Chi Han, Jeonghwan Kim +2Protein Language Models Protein Representation Learning
Siebel School of Data Science and Computing, University of Illinois Urbana-Champaign, 201 N Goodwin Ave, Urbana, IL 61801, USA · DOE Center for Advanced Bioenergy and Bioproducts Innovation, University of Illinois Urbana-Champaign, 1206 W Gregory Drive, Urbana, IL 61801, USA · Chemical and Biomolecular Engineering, University of Illinois Urbana-Champaign, 600 S Matthews Ave, Urbana, IL 61801, USA
AgentPLM: Agentic Protein Language Models with Reasoning-Augmented Decoding for Protein Sequence Design
Jun 1, 2026Sahil Rahman, Maxx Richard RahmanRL for Language Models Direct Preference Optimization
Bedford College, London, United Kingdom · Saarland University, Saarbrücken, Germany
PFArena: Benchmarking Language Models for Protein Modification
Sep 24, 2026Yawen Ouyang, Xinbo Zhang, Ziyuan Ma +15Protein Fitness Prediction Protein Language Models
Shanghai Artificial Intelligence Laboratory · Generative Symbolic Intelligence Lab (GenSI), Tsinghua University · Institute for AI Industry Research (AIR), Tsinghua University +1

Feb 24, 2026J/K move · Enter open · S save

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference

Anna Hart, Chi Han, Jeonghwan Kim, Huimin Zhao+1

Siebel School of Data Science and Computing, University of Illinois Urbana-Champaign, 201 N Goodwin Ave, Urbana, IL 61801, USA · DOE Center for Advanced Bioenergy and Bioproducts Innovation, University of Illinois Urbana-Champaign, 1206 W Gregory Drive, Urbana, IL 61801, USA · Chemical and Biomolecular Engineering, University of Illinois Urbana-Champaign, 600 S Matthews Ave, Urbana, IL 61801, USA

Protein Language Models Protein Representation Learning Bioinformatics Molecular Property Prediction

Open paper AskPDF arXiv

Modern Protein Language Models (PLMs) apply transformer-based model architectures from natural language processing to biological sequences, predicting a variety of protein functions and properties. However, protein language has key differences from natural language, such as a rich functional space despite a vocabulary of only 20 amino acids. These differences motivate research into how transformer-based architectures operate differently in the protein domain and how we can better leverage PLMs to solve protein-related tasks. In this work, we begin by directly comparing how the distribution of information stored across layers of attention heads differs between the protein and natural language domain. Furthermore, we adapt a simple early-exit technique-originally used in the natural language domain to improve efficiency at the cost of performance-to achieve both increased accuracy and substantial efficiency gains in protein non-structural property prediction by allowing the model to automatically select protein representations from the intermediate layers of the PLMs for the specific task and protein at hand. We achieve performance gains ranging from 0.4 to 7.01 percentage points while simultaneously improving efficiency by over 10 percent across models and non-structural prediction tasks. Our work opens up an area of research directly comparing how language models change behavior when moved into the protein domain and advances language modeling in biological domains.

Abstract

Protein Language Models

Explore similar work

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference

AgentPLM: Agentic Protein Language Models with Reasoning-Augmented Decoding for Protein Sequence Design

PFArena: Benchmarking Language Models for Protein Modification