JAX-RNAfold

JAX-RNAfold implements a scalable differentiable RNA folding algorithm that optimizes probabilistic sequence representations via gradient descent for RNA sequence design and integration into deep learning pipelines.

Key Features:

• Scalable Differentiable Folding: The algorithm scales to 1,250 nucleotides on a single GPU, a drastic improvement over previous proof-of-concept algorithms limited to 50 nucleotides.
• Deep Learning Integration: The software permits natural inclusion of differentiable folding as a module in larger deep learning pipelines.
• Complex RNA Design: The software supports complex RNA design procedures with flexible objective functions, including mRNA design.
• Open-Source Implementation: JAX-RNAfold is implemented as an open-source Python package.

Scientific Applications:

• RNA Sequence Design: Design of RNA sequences via gradient descent optimization of probabilistic sequence representations.
• mRNA Design: Complex mRNA design procedures employing flexible objective functions.
• Deep Learning Pipelines: Integration of differentiable RNA folding modules into larger deep learning architectures.

Methodology:

JAX-RNAfold computes the expected partition function over all RNA sequences and differentiates through this computation to enable gradient-based optimization of probabilistic sequence representations.