RNAdiffusion

Natural RNAs come in wildly different lengths — a tRNA is 76 nucleotides, a long non-coding RNA can run to thousands — and a diffusion model working directly on nucleotide tokens has to fix its output length up front. The usual workarounds are to pad or truncate every sequence to a common window, or to bolt on a length-prediction module. Both put the most awkward engineering right where the biology is most variable.

RNAdiffusion sidesteps this by moving the diffusion process off the sequence entirely. A frozen RNA-FM encoder turns a sequence of any length into contextual per-nucleotide embeddings; a trainable Query Transformer cross-attends to them with a fixed set of learned query tokens, summarizing them into exactly K latent vectors regardless of input length. A causal Transformer decoder reads those latents as a soft prompt and reconstructs nucleotides autoregressively, so length is the decoder's business. The diffusion model only ever sees a fixed-size continuous latent.

The work comes from Mengdi Wang's group at Princeton University with collaborators at Tsinghua University and Le Cong's lab at Stanford University, posted to arXiv in September 2024 and presented at the NeurIPS 2024 AI for New Drug Modalities workshop. It applies the latent-diffusion-over-a-frozen-language-model recipe that AMP-Diffusion uses on the protein side, and its architecture carried forward into the group's later RNA foundation model RNAGenesis.

#Key Features

• Fixed-size latents from variable-length RNA: The Query Transformer maps a sequence of any length to K latent vectors, so the diffusion model needs no padding, truncation, or length predictor.
• Frozen pretrained encoder: RNA-FM supplies the representation and is never updated, keeping the latent space grounded in masked-language-model pretraining over non-coding RNA.
• Plug-and-play reward guidance: Reward-model gradients are added to the predicted noise during reverse diffusion, so a new objective needs a new reward head, not a retrained generator.
• Adaptation by score-network fine-tuning: Only the score network is fine-tuned onto 5' UTRs, on top of the frozen autoencoder, transferring to a new RNA class without retraining encoder or decoder.
• Explicit reward-versus-stability trade-off: Sampling is evaluated on a Pareto front of predicted reward against minimum free energy, making reward hacking visible.

#Technical Details

The autoencoder is built on RNA-FM, with a Query Transformer adapted from the ESM-2 architecture and a decoder following ProGen2-small with a custom RNA tokenizer. It was trained on 1.1M non-coding RNA sequences under 768 nucleotides from the Ensembl subset of RNAcentral. Sweeping query tokens K over {16, 32} and latent dimension D over {40, 80, 160, 320}, the best setting reconstructs held-out sequences at a length-normalized edit distance of 0.003%; ablating the Query Transformer for mean pooling degrades that to 34.78%. The denoiser is a 24-layer Transformer with hidden dimension 2048, trained for 10 epochs on one NVIDIA H100-80G GPU with the autoencoder frozen; (K, D) = (32, 160) was selected for generation.

Generated non-coding RNAs track held-out natural sequences on minimum Levenshtein distance, minimum 4-mer distance, G/C content, minimum free energy, and secondary-structure distance, while length- and G/C-matched random sequences do not. For the 5' UTR experiment, the score network was fine-tuned for one epoch on 205K five-species 5' UTRs from the 5' UTR-LM dataset. Reward models — 6-layer residual 1D CNNs over the latents — reach Spearman R of 0.69 for mean ribosome loading and 0.56 for translation efficiency, short of UTR-LM's 0.85 and 0.64. Guided sampling raises mean translation efficiency 166.7% and mean ribosome loading 52.6% over unguided generation, on a better reward-versus-free-energy front than UTRGAN or best-of-N search.

#Applications

The immediate use is therapeutic mRNA design, where the 5' UTR sets how efficiently a transcript is translated and is a standard tuning knob for vaccines and protein-replacement constructs. Because guidance is applied at sampling rather than baked into the weights, a lab with its own measured property — a reporter assay, a stability readout, a localization score — can train a small reward head on the frozen latents and steer generation toward it. The unguided sampler supplies realistic non-coding RNA sequences for downstream filtering.

#Impact

RNAdiffusion is an early demonstration that the latent-diffusion recipe familiar from images and proteins transfers to RNA, and that a Query Transformer is a clean answer to the variable-length problem that had forced earlier sequence diffusion models into padding or length prediction. That choice proved durable: the same stack, scaled up and extended to aptamer and sgRNA design, underpins RNAGenesis. The limitations are reported candidly. All results are computational — no generated sequence was synthesized or assayed — and the reward models the guidance depends on are weaker predictors than the published UTR-LM baselines, so gains measured against them carry that error. The authors also document reward hacking at high guidance strength, where sequences accumulate reward while their secondary structure drifts from natural RNA. No code or model weights accompany the paper.