Reimplementing Unirep in JAX
Abstract
UniRep is a recurrent neural network model trained on 24 million protein sequences, and has shown utility in protein engineering. The original model, however, has rough spots in its implementation, and a convenient API is not available for certain tasks. To rectify this, we reimplemented the model in JAX/NumPy, achieving near-100X speedups in forward pass performance, and implemented a convenient API for specialized tasks. In this article, we wish to document our model reimplementation process with the goal of educating others interested in learning how to dissect a deep learning model, and engineer it for robustness and ease of use.
02 /Specification & Classification
Click any tag to browse matching papers- Assay
- Research Focus
- Approach
Reproducibility ledger
Nobody has reported trying to reproduce this paper yet. Did you replicate an experiment or run the code? Your report saves the next lab a week.
Sign in to file a reproducibility report.
Related papers
- Unified rational protein engineering with sequence-only deep representation learning
Ethan C. Alley, Grigory Khimulya, … George McDonald Church·Nature MethodsPeer-reviewed·Oct 21, 2019
- Model reconstruction from Small angle X-ray Scattering data using deep learning methods
Hao He, Can Liu, Haiguang Liu·iSciencePeer-reviewed·Feb 12, 2020
- NeRFax: An efficient and scalable conversion from the internal representation to Cartesian space
Oliver Dutton, Falk Hoffmann, Kamil Tamiola·bioRxiv·May 25, 2022
- Jaxley: Differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics
Michael Deistler, Kyra L. Kadhim, … Jakob H. Macke·Nature MethodsPeer-reviewed·Nov 13, 2025
- jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models
David Lung, Max Haberbusch·bioRxiv·Jul 31, 2026
- DeepImageJ: A user-friendly environment to run deep learning models in ImageJ
E. Gomez-de-Mariscal, C. Garcia-Lopez-de-Haro, … D. Sage·bioRxiv·May 6, 2021