Source-linked AI summary

Predictor-Guided Latent Space Codon Optimization for Maximizing Protein Expression

Alberto Caron, Tianyu Cui, Dmytro S. Lituiev, Mangal Prakash, Artem Moskalev, Amina Mollaysa, Bo Zhai, Hirsh Nanda, Daniel M. Poole, Zhongyin Liu, Iman Farasat, Robert Davidson, Nikolay V. Manyakov, Tommaso Mansi, Scott Oloff, Rui Liao

arXiv:2610.03098v1cs.AIcs.LG

TL;DR

Codon optimization must search a vast discrete space, while common proxies do not reliably represent protein expression. LSCO uses continuous latent-space search with an uncertainty-aware expression objective and regularizers; on antibody data, it outperforms frequency-based and deep generative baselines in predicted expression while retaining suitable biophysical properties.

  • Problem

    Codon optimization has a discrete, combinatorially large search space, and scarce expression labels make accurate mRNA-to-protein expression modeling difficult, leaving common heuristic proxies that often miss true protein expression.

    Source: §1 · p. 1
  • Method

    LSCO optimizes continuous representations in a pretrained mRNA language model’s latent space using an uncertainty-aware expression objective, stability and naturalness regularizers, and constrained decoding.

    Source: §1 · p. 2, §3 · p. 3
  • Results

    On antibody data, LSCO outperforms frequency-based and modern deep generative methods on predicted expression from internal and two external models.

    Source: §1 · p. 2, §5 · p. 6
  • Takeaways & Limitations

    LSCO’s antibody designs retain structural stability and host-appropriate GC content while balancing predicted expression and stability.

    Source: §1 · p. 2, §5 · p. 6

    Because reliable mRNA–protein expression data are scarce, the evaluation relies on predicted rather than empirical yield.

    Source: §5 · p. 6

Abstract

from arXiv · show

Codon optimization, the process of selecting synonymous codons to improve mRNA translation efficiency and protein expression, is central to therapeutic protein production and mRNA vaccines, yet it remains a hard problem. The design space is discrete and combinatorially large, precluding gradient-based methods, and existing tools rely on heuristic proxies (e.g., Codon Adaptation Index or GC-content) that poorly capture true expression. We introduce Latent-Space Codon Optimization (LSCO), which recasts this discrete problem as a continuous one by mapping sequences into the latent space of a pretrained mRNA language model, enabling efficient gradient-based search. LSCO combines four components: a data-driven expression objective from an uncertainty-aware predictor, a Minimum-Free-Energy regularizer for structural stability, a naturalness prior from a protein-to-codon back-translation model, and constrained decoding for protein fidelity. On a real-world, wet-lab antibody expression dataset, LSCO outperforms simple frequency-based, as well as modern deep generative baselines in predicted expression, while retaining suitable biophysical properties.

1 Introduction

Codon optimization offers a large design space for improving protein yield, but that space is hard to search and true expression is difficult to model. LSCO addresses both challenges with continuous latent-space optimization and reports stronger predicted expression than existing baselines while retaining suitable biophysical properties.

1 Introduction

  • Synonymous codons create a combinatorially vast design space for tailoring mRNA without changing the encoded protein.The paper notes that trastuzumab’s heavy and light chains admit on the order of 10300 synonymous encodings.Source: §1 · p. 1
  • Search is difficult because synonymous choices are discrete and numerous, while scarce labels and incomplete biological understanding limit accurate expression models.Existing methods therefore use proxies such as CAI and GC content, which often fail to capture true protein expression.Source: §1 · p. 1
  • LSCO replaces discrete search with continuous latent optimization, combining an uncertainty-aware expression objective, MFE and naturalness regularizers, and constrained decoding for protein fidelity.Its naturalness prior comes from a protein-to-codon back-translation model trained on Observed Antibody Space data.Source: §1 · pp. 1–2
  • On a real-world antibody expression dataset, LSCO outperforms frequency-based and modern deep generative baselines on predicted expression while retaining structural stability and host-appropriate GC content.Independent held-out predictors confirmed the predicted-expression result.Source: §1 · p. 2

2 Background and Related Work

Codon optimization seeks an mRNA encoding a target protein that maximizes predicted expression, but the search space is combinatorial and the expression surrogate is imperfect. Prior approaches use sequence proxies or learned models, while LSCO places optimization in a latent space and adds realistic codon-use and uncertainty controls.

2 Background and Related Work

  • Codon optimization searches synonymous codon combinations for an mRNA sequence that maximizes a surrogate for downstream protein expression.The design space is the product of the synonymous-codon choices available at each amino-acid position.Source: §2 · p. 2
  • The problem resists exhaustive or gradient-based search because its synonymous-codon space is discrete and combinatorial, and its expression model is inaccessible or unreliable.Underlying biology is partly understood and labelled data are scarce.Source: §2 · p. 2
  • Earlier tools optimize frequency proxies, generative mRNA models learn codon usage, and post-hoc optimizers target sequence features rather than protein expression.The paper identifies a potential gap between those optimization objectives and the desired biological outcome.Source: §2 · p. 2
  • LSCO extends latent-space optimization with realistic codon-usage priors, a domain-trained expression surrogate, and ensemble uncertainty estimates to penalize out-of-distribution drift.Its framework also balances stability, naturalness, and protein fidelity.Source: §2 · p. 2, Figure 1 · p. 3

3 Latent-Space Codon Optimization

LSCO searches for synonymous codon choices by optimizing continuous representations in a pretrained mRNA language model’s latent space. Its expression objective is balanced with uncertainty control, predicted structural stability, natural codon usage, and decoding constraints that preserve the target protein.

3 Latent-Space Codon Optimization

  • LSCO optimizes codon-level embeddings in a pretrained mRNA language model’s latent space, using predicted yield as its main objective and regularizers for stability and naturalness.Constrained decoding returns synonymous choices so the target protein sequence is preserved.Source: §3 · p. 3

Appendix B).

  • The pretrained variational mRNA model encodes codons continuously and decodes them through a temperature-controlled softmax relaxation suitable for differentiable optimization.The model is trained on curated OAS-mRNA sequences.Source: Appendix B · p. 3
  • An ensemble expression predictor takes the relaxed sequence as input, and its lower-confidence-bound utility penalizes high-uncertainty designs.This is intended to discourage out-of-distribution designs during optimization.Source: Appendix B · pp. 3–4
  • The MFE regularizer discourages unstable sequences, with lower predicted MFE corresponding to more stable structure.Its predictor is trained on ViennaRNA-computed MFE labels.Source: Appendix B · p. 4
  • A back-translation model supplies a natural-codon reference, and a latent-distance penalty keeps optimized sequences near that reference.The reference is generated conditionally from the target protein sequence.Source: Appendix B · p. 4
  • Constrained decoding masks non-synonymous codons, while Figure 2 compares predicted A280 with MFE, GC%, and CAI across 54 antibodies.The caption reports that LSCO variants attain the highest or comparable expression across three predictors; adding MFE regularization improves stability and nudges GC/CAI upward at modest expression cost.Source: Appendix B · p. 4, Figure 2 · p. 4

4 Experiments

On antibody data, LSCO outperforms frequency-based and modern deep generative methods on predicted expression, while experiments examine expression–stability trade-offs and the effect of the pretrained latent space. Evaluation relies on predicted rather than empirical yield.

4 Experiments

  • The evaluation compares host-preferred frequency selection, CodonBERT back-translation, modern generative methods and optimizers, and lower-bound random and VMAE designs.Protein yield is assessed using predicted A280 from the native predictor and two held-out evaluators, alongside MFE, GC%, and CAI.Source: §4 · pp. 4–5
  • Table 1 reports A280 from three predictors and biophysical metrics for 54 mAbs; LSCO generally outperforms direct input-space optimization (DISO).The table also varies MFE and uncertainty penalties and includes a naturalness-prior ablation; values are means with 95% confidence bands.Source: Table 1 · p. 5, §4 · p. 5
  • LSCO attains the highest predicted A280 on internal and external 1D predictors and scores comparably on the 2D GNN; adding MFE regularization improves stability at modest expression cost.The expression-only variant has weaker MFE than strong baselines; LSCO–LCB+MFE lowers MFE and nudges GC and CAI upward within host-appropriate ranges.Source: §4 · p. 5
  • With λMFE = 0.1, MFE falls and GC and CAI rise into host-appropriate ranges at modest expression cost; λMFE = 1 yields no further stability.Enabling λLCB = 1 only slightly lowers A280 while keeping optimization from drifting off-distribution, motivating the selected settings λMFE = 0.1 and λLCB = 1.Source: §4 · p. 5
  • Removing the naturalness prior leaves native-predictor A280 largely unchanged but degrades the held-out 2D-GNN score.The authors describe this as the expected signature of an on-manifold regularizer that keeps designs plausible to independent evaluators.Source: §4 · p. 5

5 Conclusion

  • On antibody data, LSCO outperforms frequency-based and modern deep generative methods on predicted expression across internal and two external models.The evaluation uses 54 starting monoclonal antibodies; reliable mRNA–protein expression data are scarce, so yield is predicted rather than empirically measured.Source: §4 · p. 4, §5 · p. 6

A Extended Related Work

Prior codon optimization ranges from host-frequency heuristics to pretrained generators and sequence optimizers, while discrete-sequence methods differ in search geometry and trust mechanisms. LSCO uses a pretrained mRNA model’s latent space and an uncertainty-aware expression surrogate.

Codon Optimization

  • Traditional codon optimization tailors gene sequences to host codon preferences, while frequency-based tools emphasize CAI and codon-pair bias but may neglect mRNA stability and expression.Deep-learning approaches include pretrained mRNA language models and optimizers that combine multiple sequence objectives.Source: A Extended Related Work · p. 9

Discrete Sequence Optimization Methods

  • Discrete-sequence optimization methods differ by whether they search sequences or a learned latent space, and by how they guard against unreliable surrogate predictions.The paper uses this search-geometry and trust-mechanism distinction to organize prior work.Source: A Extended Related Work · p. 9

Sequence-Space Search

  • Sequence-space methods evolve or resample candidates directly, but gradient ascent on relaxed sequences can reward-hack when the surrogate overpredicts off-distribution inputs.Prior remedies include conservative surrogates, reinforcement-learning model selection, and Bayesian optimization with uncertainty or trust regions.Source: A Extended Related Work · p. 9

Latent-Space Optimization

  • Latent-space optimization offers smooth, low-dimensional search and a realism prior, but requires a decoder that can mis-decode and move designs off-distribution.Latent-space Bayesian optimization adds probabilistic surrogates to balance exploration and exploitation.Source: A Extended Related Work · p. 9

Codon Optimization

  • Figure 3 depicts LSCO mapping protein codon choices into continuous latent optimization, combining expression and stability predictors, a naturalness prior, and constrained decoding.The decoding step produces mRNA sequences compatible with the target protein.Source: Figure 3 · p. 10

Our approach

  • LSCO searches a pretrained mRNA foundation model’s latent space, directly targets predicted expression, and uses ensemble uncertainty to steer away from unreliable regions.Its multi-criteria objective also weighs stability and naturalness while maintaining target-protein fidelity.Source: A Extended Related Work · p. 10

B Implementation details

The implementation combines pretrained mRNA and back-translation models with differentiable expression and MFE predictors, constrained decoding, and annealed-temperature gradient optimization. The expression predictor reaches a test-split Spearman rank correlation of 0.65.

Back-translation model

  • The back-translation BERT maps amino-acid tokens to codon tokens, masking amino-acid inputs with probability 0.2 during training.It uses a joint 91-token vocabulary and has a maximum sequence length of 2048 tokens.Source: B Implementation details · p. 10

Variational mRNA language model

  • The 12-layer ModernBERT VMAE uses a 70-token vocabulary, splits into six-layer encoder and decoder, and is pretrained with a Gaussian latent prior.Its masking ratio is 0.15; pretraining runs for 40 epochs.Source: B Implementation details · p. 10

Expression predictor

  • The differentiable dual-chain expression predictor aggregates heavy- and light-chain representations with MIL attention; predicted and ground-truth expression have test-split Spearman correlation 0.65.It reuses the pretrained encoder and is trained on antibody expression measured by A280.Source: B Implementation details · pp. 10–11

Algorithm 1

  • LSCO updates latent variables by gradient ascent on expression, MFE, and naturalness terms, then converts the constrained decoder output to a discrete codon sequence.The algorithm anneals decoding temperature from a high to a low value during optimization.Source: B Implementation details · p. 11

MFE predictor

  • The MFE predictor mirrors the expression predictor and reaches 0.91 test-split Spearman rank correlation with ground-truth MFE.It is trained on 50,000 OAS mRNA sequences using MFE labels computed with ViennaRNA.Source: B Implementation details · p. 11

Constrained Decoding

  • Constrained decoding masks out codons that do not encode the target amino acid before softmax, allowing synonymous-codon exploration while preserving the protein sequence.The optimizer evaluates the relaxed sequence rather than the latent representation because direct latent optimization caused expression shifts without corresponding sequence changes in early experiments.Source: B Implementation details · p. 11

Decoding Temperature

  • Temperature annealing balances broad synonymous-codon exploration against rounding errors at high temperatures and hindered optimization progress at low temperatures.The implementation uses τhigh = 3 and τlow = 0.01 and runs for 3000 iterations.Source: B Implementation details · p. 11

C Additional Results and Experiments

The comparison tests several codon-generation strategies against biophysical measures, and modern pretrained mRNA models score favorably across those measures while making fewer changes to reference sequences.

C.1 Pretrained mRNA language models comparison

  • The comparison evaluates Random, Freq Sample, CodonBERT, and two VMAE generation variants.VMAE – Prior samples from the prior latent space; VMAE – BERT starts from CodonBERT-generated sequences.Source: §C.1 · p. 12
  • All generation strategies are assessed using MFE, GC content, CAI, and codon-level Hamming distance from the reference sequence.Source: §C.1 · p. 12
  • Modern pretrained models score favorably on MFE, GC content, CAI, and codon Hamming distance versus other baselines.CodonTransf, GEMORNA, and VMAE – BERT show lower MFE, higher CAI, and lower Hamming distance; GEMORNA’s GC content is just above the ideal host-specific range.Source: §C.1 · p. 12

C.2 LSCO Optimization Trajectory

LSCO moves starting sequences from a low-expression latent-space cluster toward clusters with higher predicted A280, following consistent trajectories that remain in low-uncertainty regions.

C.2 LSCO Optimization Trajectory

  • LSCO-optimized sequences move from C7, with average predicted A280 of 1.808, toward clusters averaging approximately 3.0–3.4.Source: §C.2 · p. 12
  • Optimization trajectories follow a similar direction, consistent with ascent along the learned fitness landscape rather than arbitrary sequence perturbations.Source: §C.2 · p. 12
  • Optimized sequences remain in low-epistemic-uncertainty regions, consistent with the LCB utility keeping the search in-distribution.Source: §C.2 · p. 12

C.3 Codon-Usage Analysis

LSCO makes amino-acid-specific codon substitutions and shifts sequences toward higher-expression regions; across evaluators, its expression ranking is largely preserved.

C.3 Codon-Usage Analysis

  • Starting sequences in C7 have average predicted A280 of 1.808, while optimized sequences move toward clusters averaging approximately 3.0–3.4 and remain in low-uncertainty regions.Source: Figure 5 · p. 13
  • LSCO redirects codon usage toward preferred alternatives, with moderate shifts toward GC-richer codons but amino-acid-specific patterns rather than simple GC maximization.The magnitude of redistribution varies across amino-acid families and is concentrated in specific synonymous groups.Source: §C.3 · p. 13, Figure 6 · p. 13

C.4 A280 Predictors and Agreement

  • The external 1D evaluator uses a different, shallower autoencoder backbone, while the 2D GNN evaluates secondary-structure graphs rather than raw sequence embeddings.Both external predictors are distinct from the native predictor used to guide LSCO.Source: §C.4 · p. 14
  • LSCO variants are rated highest by the native predictor and comparably or slightly higher by external predictors, with agreement mainly in ranking rather than absolute values.The external evaluators are a 1D Shallow-AE model and a structure-aware 2D GNN.Source: §C.4 · p. 14, Figure 7 · p. 14

C.5 Deep-Ensemble and Epistemic Uncertainty Aware Objective

The uncertainty-aware ensemble objective keeps LSCO’s search low-uncertainty under its guiding predictor, while MFE regularization trades some predicted expression for improved stability and host-appropriate metrics.

C.5 Deep-Ensemble and Epistemic Uncertainty Aware Objective

  • Each of the ten probabilistic ensemble members predicts a Gaussian distribution over A280 rather than a single point estimate.Source: §C.5 · p. 14
  • The LCB utility penalizes epistemic uncertainty, which is higher in underrepresented sequence regions, to steer optimization away from out-of-distribution designs.Epistemic uncertainty is measured by disagreement across ensemble members’ means; aleatoric uncertainty is the average of their predicted variances.Source: §C.5 · p. 15
  • LSCO attains the highest predicted A280 across all three predictors while keeping low epistemic uncertainty under the predictor guiding optimization.Under external predictors not used in optimization, LSCO’s uncertainty is comparable to other high-expression methods.Source: §C.5 · p. 16

C.6 MFE regularization sensitivity

  • Increasing the MFE weight lowers MFE and raises GC and CAI into host-appropriate ranges, at a modest predicted-A280 cost; stability gains saturate beyond λMFE = 0.1.The authors identify λMFE = 0.1 as the sweet spot used throughout.Source: Table 2 · p. 15, §C.6 · p. 16

C.7 Seed sensitivity study

Across 10 seeds on one representative mAb, both LSCO variants remain above stochastic baselines on the 1D predictors with bounded spread, while maintaining stable biophysical metrics.

C.7 Seed sensitivity study

  • The 10-seed comparison uses one representative mAb chosen near the dataset’s median A280, and its spread measures generator stochasticity rather than antibody-to-antibody variation.Deterministic methods are omitted because they are seed-inert; evaluators are deterministic for a given sequence.Source: §C.7 · p. 16
  • Both LSCO variants exceed stochastic baselines on predicted A280 for every seed under the two 1D predictors, with bounded spread.The LCB variant is tightest; LCB + MFE has slightly more variability and slightly lower absolute A280.Source: §C.7 · p. 16
  • Both LSCO variants keep GC within the 50–60% band on every seed, while LCB + MFE consistently yields lower MFE than LCB.Source: §C.7 · p. 16
  • Under the more conservative, less aligned 2D GNN, method boxes overlap more and seed spread is larger for all methods.Source: §C.7 · p. 16

D LSCO vs. Direct Input-Space Optimization (DISO)

LSCO optimizes in a pretrained mRNA language model’s latent space, whereas DISO optimizes codon logits directly. They share downstream components and optimization settings, and LSCO generally matches or exceeds DISO across evaluators.

D LSCO vs. Direct Input-Space Optimization (DISO)

  • LSCO generally matches or exceeds DISO across evaluators, with the comparison identifying the pretrained latent space as driving the gains.The empirical comparison is reported in Table 1.Source: D LSCO vs. Direct Input-Space… · p. 17
  • LSCO optimizes a continuous latent code decoded into codons at each step; DISO instead optimizes per-position codon logits directly in sequence space.The downstream softmax, valid-codon mask, expression predictor, and MFE predictor are identical.Source: D LSCO vs. Direct Input-Space… · p. 17
  • LSCO anchors naturalness by penalizing distance from a latent code for a BERT-natural sequence, while DISO penalizes KL divergence from its fixed codon distribution.For a two-chain antibody, DISO averages heavy- and light-chain KL divergences; setting λ_nat = 0 removes either anchor.Source: D LSCO vs. Direct Input-Space… · p. 17
  • The methods share model components and optimization settings, leaving search geometry and its induced naturalness anchor as the comparison’s differences.The matched schedule includes temperature annealing, protein-consistency constraints, step count, learning rate, optimizer, and MFE-weight range.Source: Table 3 · p. 18, D LSCO vs. Direct Input-Space… · p. 18
  • DISO’s initialization can be BERT-natural or random Gaussian, but its naturalness anchor remains the BERT-natural distribution in either case.The random-Gaussian variant masks invalid codons so optimization starts from a protein-consistent point.Source: D LSCO vs. Direct Input-Space… · p. 18
Loading 2610.03098v1…