RiboSphere

<div align="center">

RiboSphere

Learning Unified and Efficient Representations of RNA Structures

Discrete geometric RNA tokens through a geometric Transformer, finite scalar quantization, and flow matching.

<p> <img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-2.4%2B-EE4C2C?logo=pytorch&logoColor=white"> <img alt="Safetensors" src="https://img.shields.io/badge/weights-safetensors-6C5CE7"> <img alt="ICML 2026" src="https://img.shields.io/badge/ICML-2026-2D6CDF"> <img alt="RNA 3D" src="https://img.shields.io/badge/domain-RNA%203D-168B72"> </p>

Overview | Checkpoints | Quick start | API | Evaluation | Citation

</div>

---

Overview

RiboSphere learns a discrete geometric alphabet for RNA structures. A geometric Transformer encodes mean-centered atomic coordinates, finite scalar quantization (FSQ) maps each nucleotide to a discrete structural token, and a flow-matching decoder reconstructs the full three-dimensional structure from the token sequence.

<div align="center"> <img src="assets/ribosphere-overview.png" alt="RiboSphere architecture and downstream integration" width="100%"> <br> <em>RiboSphere encodes RNA geometry into a discrete structural alphabet, reconstructs atomic coordinates with flow matching, and transfers the learned representation to downstream tasks.</em> </div>

This repository contains the nine reconstruction/tokenizer checkpoints reported for the 2-layer encoder / 8-layer decoder / 256-dimensional encoder setting.

Highlights

Checkpoint zoo

Checkpoint names describe the geometric representation and vocabulary size directly:

<representation>-vocab<size>

For example, backbone-base-vocab4375 uses the 11-atom representation and a 4,375-token FSQ vocabulary.

Atomic representations

Available checkpoints

The reconstruction metrics below are reproduced from Table 1 of the paper. RMSD is reported in angstroms; higher TM-score and lDDT are better.

Which checkpoint should I use?

• Best reported reconstruction: c4prime-vocab4375 achieves the lowest RMSD in the paper's reconstruction benchmark.
• Most detailed geometry: backbone-base-vocab4375 retains backbone geometry and a base-orientation anchor.
• Compact token space: choose a vocab240 checkpoint when vocabulary size matters more than reconstruction fidelity.
• Balanced vocabulary: vocab1000 offers an intermediate discrete space, but performance depends on atomic representation.

Quick start

Installation

Clone the model repository and install its runtime dependencies:

pip install -r requirements.txt

Load a checkpoint

import torch

from src.models import RiboSphere

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = RiboSphere.from_pretrained(
    ".",
    variant="backbone-base-vocab4375",
).to(device).eval()

To load from the Hugging Face Hub, replace "." with the repository ID:

model = RiboSphere.from_pretrained(
    "zz312/RiboSphere",
    variant="backbone-base-vocab4375",
).to(device).eval()

When variant is specified, only that checkpoint's configuration and weights are downloaded.

Tokenize and reconstruct an RNA structure

import torch
from biotite.structure.io.pdb import PDBFile

from src.datasets import kabsch_rmsd, prepare_rna_coordinates
from src.models import RiboSphere

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RiboSphere.from_pretrained(
    ".",
    variant="backbone-base-vocab4375",
).to(device).eval()

pdb_file = PDBFile.read("pdbs/7D7X_1_A.pdb")
atom_array = pdb_file.get_structure(model=1, altloc="first")
coordinates = prepare_rna_coordinates(
    atom_array,
    atoms=model.config.atoms,
).to(device)

torch.manual_seed(0)
with torch.inference_mode():
    encoder_states, quantized_states, token_ids = model.encode(
        coordinates,
        preprocess=True,
    )
    reconstructed_nm = model.decode(token_ids, num_steps=200)

# PDB/Biotite coordinates are in angstroms; decoder outputs are in nanometers.
reconstructed_angstrom = reconstructed_nm.cpu() * 10.0
rmsd = kabsch_rmsd(reconstructed_angstrom, coordinates.cpu())

print(f"Aligned RMSD: {rmsd:.4f} A")

API reference

RiboSphere.from_pretrained

model = RiboSphere.from_pretrained(
    model_path,
    variant="backbone-base-vocab4375",
)

• model_path: a local repository path or Hugging Face repository ID.
• variant: one of the canonical names listed in the checkpoint table.

model.encode

encoder_states, quantized_states, token_ids = model.encode(
    coordinates,
    preprocess=False,
)

Set preprocess=True for PDB coordinates: inputs are centered per structure and converted from angstroms to nanometers.

model.decode

coordinates_nm = model.decode(
    token_ids,
    num_steps=200,
    noise_weight=0.2,
    score_weight=1.0,
    guidance_weight=1.0,
)

The output has shape [B, L, A, 3] and is expressed in nanometers. Sampling is stochastic. Set a PyTorch random seed when reproducibility is required.

Repository layout

.
├── configs/                       # one JSON configuration per checkpoint
├── weights/                       # all safetensors checkpoints
├── src/
│   ├── datasets/                  # PDB preprocessing and aligned RMSD
│   └── models/                    # RiboSphere, FSQ, attention, and flow decoder
├── pdbs/                          # example RNA structures
├── variants.json                  # canonical names and checkpoint file paths
├── requirements.txt
└── README.md

Evaluation

Reconstruction data

The reconstruction split follows the single-state setting from gRNAde and structurally separates the test clusters from training clusters. After expanding multiple conformations associated with each sequence, the paper reports:

Metrics

• RMSD: rigid-body-aligned coordinate deviation in angstroms.
• TM-score: global fold similarity.
• lDDT: local geometric agreement.
• Codebook utilization: percentage of the implicit FSQ vocabulary.

Limitations and responsible use

• RiboSphere is a research model, not an experimentally validated structure-determination system.
• Generated structures may be biophysically implausible and should not be treated as experimental evidence.
• Non-standard residues and residues missing required atoms are removed by the provided preprocessing helper.
• Global translation is removed by construction. Random rotations were used during training, but downstream analyses should still use appropriate rigid-body alignment.
• Structural predictions should be validated with independent geometric, energetic, and experimental methods before biological interpretation.

Citation

If you use RiboSphere, please cite:

@inproceedings{zhang2026ribosphere,
title={RiboSphere: Learning Unified and Efficient Representations of {RNA} Structures},
author={Zhou Zhang and Hanqun Cao and Cheng Tan and Fang Wu and Pheng-Ann Heng and Tianfan Fu},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
}

---

<div align="center">

From continuous RNA geometry to a discrete, interpretable structural alphabet.

</div>