README.md

---
library_name: pytorch
tags:
- rna
- rna-structure
- structure-generation
- molecular-generation
- flow-matching
- finite-scalar-quantization
---

<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](#overview) | [Checkpoints](#checkpoint-zoo) | [Quick start](#quick-start) | [API](#api-reference) | [Evaluation](#evaluation) | [Citation](#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

| | |
| --- | --- |
| **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:

```text
<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`](weights/c4prime-vocab240.safetensors) | `(8, 6, 5)` | 240 | 2.14 | 0.70 | 0.73 | 100.0% |
| [`c4prime-vocab1000`](weights/c4prime-vocab1000.safetensors) | `(8, 5, 5, 5)` | 1,000 | 1.88 | 0.75 | 0.76 | 80.7% |
| [`c4prime-vocab4375`](weights/c4prime-vocab4375.safetensors) | `(7, 5, 5, 5, 5)` | 4,375 | **1.25** | **0.84** | **0.83** | 39.8% |
| [`backbone-vocab240`](weights/backbone-vocab240.safetensors) | `(8, 6, 5)` | 240 | 1.80 | 0.76 | 0.77 | 100.0% |
| [`backbone-vocab1000`](weights/backbone-vocab1000.safetensors) | `(8, 5, 5, 5)` | 1,000 | 2.33 | 0.69 | 0.74 | 84.2% |
| [`backbone-vocab4375`](weights/backbone-vocab4375.safetensors) | `(7, 5, 5, 5, 5)` | 4,375 | **1.58** | **0.80** | 0.76 | 39.4% |
| [`backbone-base-vocab240`](weights/backbone-base-vocab240.safetensors) | `(8, 6, 5)` | 240 | 2.05 | 0.71 | 0.73 | 100.0% |
| [`backbone-base-vocab1000`](weights/backbone-base-vocab1000.safetensors) | `(8, 5, 5, 5)` | 1,000 | 1.60 | 0.78 | 0.79 | 88.0% |
| [`backbone-base-vocab4375`](weights/backbone-base-vocab4375.safetensors) | `(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:

```bash
pip install -r requirements.txt
```

### Load a checkpoint

```python
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:

```python
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

```python
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`

```python
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`

```python
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`

```python
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

```text
.
├── 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:

```bibtex
@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>