cs.AIOct 2, 2026
SavePredictor-Guided Latent Space Codon Optimization for Maximizing Protein Expression
Organizations: Johnson & Johnson Innovative Medicine · Ellison Institute of Technology Oxford · New Theory AI
Abstract
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.
Figures & tables
Figure 1: LSCO framework. (a) mRNA sequences are mapped to codon-latent space using a pretrained variational masked language model. (b) mRNA property predictor reuses the pretrained encoder in (a) and takes the soft and continuous version of the discrete sequence. (c) We generate a natural mRNA sequence using a pretrained protein-to-codon back-translation model, whose representation from the pretrained encoder in (a) is used as the naturalness prior. (d) Latent representations are optimized with respect to a data-driven expression predictor , an MFE predictor MFE , and a naturalness prior from (c). Constrained decoding ensures valid codon choices, preserving the target protein sequence. The result is an optimized mRNA that balances expression and stability.
Figure 2: Optimized sequence comparison on 54 mAbs. Top row: predicted A280 from three models: LSCO’s native 1D transformer, an external 1D transformer, and an external structure-aware 2D GNN. LSCO variants attain the highest or comparable expression across all three. Bottom row: biophysical metrics on heavy + light chains: MFE (lower = more stable), GC% (green band = host-appropriate range), CAI. LSCO-LCB+MFE variant recovers stability and nudges GC/CAI upward at a modest expression cost, staying within host-appropriate ranges.
| Configuration | Predicted A280 | Metrics (VH + VL) | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Study | Space | 1D Internal | 1D External | 2D External | MFE | GC% | CAI | |||
| MFE Sensitivity | LSCO | 0 | 0 | 0.01 | ||||||
| 0 | 1 | 0.01 | ||||||||
| 0.1 | 0 | 0.01 | ||||||||
| 0.1 | 1 | 0.01 | ||||||||
| 1 | 0 | 0.01 | ||||||||
Table 1: Sensitivity and ablation studies. Predicted protein yield (A280) from the internal 1D Transformer and two held-out evaluators, together with biophysical metrics for the joint heavy and light chains, on the same 54 starting mAbs as Figure 2 . Values are means with 95% confidence bands. LSCO searches in the pretrained VMAE latent space, whereas DISO optimizes codon logits directly in input space. The and columns respectively control MFE regularization and the uncertainty penalty; denotes removal of the naturalness prior. Overall, provides good expression–stability trade-off, and LSCO generally outperforms DISO.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Generative Codon optimization problem. Starting from a protein sequence, we navigate the combinatorial codon design space by mapping to and optimizing in a continuous latent space. LSCO integrates differentiable predictors for expression and stability, a naturalness prior, and constrained decoding to produce host-compatible, high-expression mRNA sequences.
Figure 4: Sequence-level biophysical proxies across generation methods. Minimum Free Energy (MFE; lower = more stable), GC Content (GC%; shaded band marks host-appropriate range), Codon Adaptation Index (CAI; higher is better), and codon Hamming distance ( ) from the reference sequence (lower = fewer codon edits), for non-contextual baselines ( Random , Freq Sample ), modern learned generators ( CodonBERT , CodonTransf , GEMORNA ), and our pretrained models ( VMAE–Prior , decoding from Gaussian latent prior; VMAE–BERT , decoding from CodonBERT generated sequence’s latent space).
Figure 5: LSCO optimization trajectories in latent space. UMAP projection of the pretrained VMAE latent space on a holdout fold, coloured by predicted A280 (left) and by predictor epistemic uncertainty (right). Starting sequences (red triangles) lie in a low-expression cluster (C7, ); LSCO-optimized sequences (gold stars) are relocated toward high-expression clusters ( – ). Trajectories (grey arrows) follow a consistent direction and remain in low-uncertainty regions.
Figure 6: Codon substitutions induced by LSCO optimization. Rows denote original codons, columns optimized codons; colours show row-normalized transition probabilities for synonymous substitutions produced by LSCO – LCB + MFE , over the 54 starting mAbs. Strong diagonal structure indicates codon conservation, while off-diagonal transitions reveal learned preferences. The Symmetrized KL (SKL) divergence quantifies the magnitude of codon redistribution per amino-acid family (higher SKL larger redistribution).
Figure 7: Predicted-A280 agreement across evaluators. Each dot is one optimized sequence, coloured by method; the dashed line marks perfect agreement. The horizontal axis is the native 1D predictor ; the vertical axis is an independent external predictor withheld from optimization, the 1D Shallow-AE model (a) and the structure-aware 2D GNN (b). LSCO variants (blue) are rated highest by the native model, but score favorably also under the two external models in terms of ordering/ranking preserving.
| Predicted A280 | Metrics (VH + VL) | ||||||
|---|---|---|---|---|---|---|---|
| 1D Transf (Internal) | 1D Transf (External) | 2D GNN | MFE | GC% | CAI | ||
| 0 | 0 | ||||||
| 1 | |||||||
| 0.1 | 0 | ||||||
| 1 | |||||||
| 1 | 0 | ||||||
Table 2: MFE sensitivity analysis for LSCO variants. Predicted A280 values from three surrogate models (1D Transformer (Internal), 1D Transformer (External), 2D GNN) and biophysical metrics (minimum free energy, GC content, codon adaptation index) reported for the joint heavy + light chain construct on the 54 starting sequences. The second row ( ) is the LCB-only LSCO baseline (‘LSCO – LCB’ in Figure 2 ). The fourth row ( ) is the LCB+MFE baseline (‘LSCO – LCB + MFE’ in Figure 2 ) that best balances A280 and MFE objectives. Bold indicates the best value per column: highest for A280 predictions and CAI, most negative for MFE. GC% is not bolded as “best” is target-dependent. Values are mean 95% confidence bands across starting sequences.
Figure 8: Epistemic uncertainty of the optimized sequences. Epistemic standard deviation (across-member disagreement, ) of the three probabilistic deep-ensemble predictors — 1D Transformer Internal/Native (left), 1D Transformer External (middle), 2D GNN External (right) — evaluated on the same 9 optimized mAbs as Figure 2 ; each dot is one mAb. Under the native predictor, LSCO variants (blue) show low, concentrated epistemic uncertainty, comparable to the unoptimized VMAE and well below the other optimizers, confirming the LCB utility keeps the search in-distribution. Under the external predictors — never seen during optimization — LSCO’s uncertainty is comparable to other high-expression methods, as expected for held-out evaluators.
Figure 9: Seed sensitivity on a single mAb (10 seeds). Boxes show the spread over different 10 seeded runs, reflecting the stochasticity behind different codon optimization methods. Top: predicted A280 from the native 1D transformer (internal), an external 1D transformer, and an external 2D GNN. Bottom: biophysical metrics for the joint heavy+light construct (VH+VL): MFE (lower = more stable), GC% (shaded band = host-appropriate – ), and CAI. Both LSCO variants remain above the stochastic baselines on the 1D predictors across all seeds with bounded spread, keep GC within the admissible band, and LSCO – LCB + MFE consistently lowers MFE relative to LSCO – LCB ; seed spread is largest under the more conservative 2D GNN evaluator.
| LSCO | DISO | |
| Optimized variable | latent code | per-position codon logits |
| Generative decoder in loop | Yes ( ) | No |
| Naturalness anchor | latent distance | KL to fixed BERT distribution |
| Naturalness weight | (0 to disable) | (0 to disable) |
| Initialization | BERT-natural warm start (or ) | |
| Temperature schedule |
Table 3: Configuration differences between LSCO and DISO. All model components (expression ensemble, MFE head, tokenizers, valid-codon mask, candidate gate) are shared; only the search geometry and its downstream consequences differ.