RiboSphere

Learning Unified and Efficient Representations of RNA Structures

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

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

---

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.

RiboSphere encodes RNA geometry into a discrete structural alphabet, reconstructs atomic coordinates with flow matching, and transfers the learned representation to downstream tasks.

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

Highlights

Discrete structural tokens	One integer token per nucleotide through an implicit FSQ codebook
Three geometric resolutions	C4'-only, 10-atom backbone, and 11-atom backbone-plus-base representations
Three vocabulary sizes	240, 1,000, and 4,375 structural tokens
Safe serialization	All checkpoints are distributed in safetensors format

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

Name component	Atoms per nucleotide	Atom channels
c4prime	1	C4'
backbone	10	P, C5', C4', C3', C2', C1', O5', O4', O3', O2'
backbone-base	11	The 10 backbone atoms plus N9 for purines or N1 for pyrimidines

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.

Checkpoint	FSQ levels	Vocabulary	RMSD ↓	TM-score ↑	lDDT ↑	Utilization
c4prime-vocab240	(8, 6, 5)	240	2.14	0.70	0.73	100.0%
c4prime-vocab1000	(8, 5, 5, 5)	1,000	1.88	0.75	0.76	80.7%
c4prime-vocab4375	(7, 5, 5, 5, 5)	4,375	1.25	0.84	0.83	39.8%
backbone-vocab240	(8, 6, 5)	240	1.80	0.76	0.77	100.0%
backbone-vocab1000	(8, 5, 5, 5)	1,000	2.33	0.69	0.74	84.2%
backbone-vocab4375	(7, 5, 5, 5, 5)	4,375	1.58	0.80	0.76	39.4%
backbone-base-vocab240	(8, 6, 5)	240	2.05	0.71	0.73	100.0%
backbone-base-vocab1000	(8, 5, 5, 5)	1,000	1.60	0.78	0.79	88.0%
backbone-base-vocab4375	(7, 5, 5, 5, 5)	4,375	1.35	0.82	0.82	39.9%

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:

Split	Structures
Training	11,183
Validation	551
Test	239

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},
}

---

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