RCLM

RCLM (RNA Composite-Likelihood Language Model) is an Evoformer-style RNA language model that encodes a single RNA sequence into a per-residue embedding s and a pairwise (residue-residue) embedding z. These embeddings are used as input features for downstream structure-prediction models such as DRfold2.

Method details are described in the DRfold2 paper: DRfold2 is a deep learning-based tool that enables efficient and accurate RNA structure prediction*.

Installation

pip install torch numpy

The pretrained checkpoint (weight/epoch_67000, ~190MB) isn't in this repo — download it into weight/:

mkdir -p weight
pip install gdown
gdown https://drive.google.com/uc?id=1IM3nINUYxJSncsjKhoPw5ZbC-WswoaIv -O weight/epoch_67000

Repository layout

RCLM/
├── rclm/               # model source
│   ├── Model.py         # RNAembedding + RNA2nd (top-level model)
│   ├── Evoformer.py      # stack of Evoformer blocks
│   ├── EvoMSA.py         # MSA row/col attention, transition, outer-product-mean
│   ├── EvoPair.py        # triangle update/attention, pair transition
│   └── basic.py          # shared building blocks (Linear, LayerNorm helpers, etc.)
├── weight/
│   └── epoch_67000       # pretrained checkpoint (state_dict)
└── example.py           # runnable example, see "Example usage" below

Model

rclm.Model.RNA2nd is the top-level module. It is built from a config dict:

lmcfg = {
    's_in_dim':  5,    # one-hot RNA alphabet size (A, G, C, U, gap)
    'z_in_dim':  2,
    's_dim':     512,
    'z_dim':     128,
    'N_elayers': 18,
}

Loading the pretrained weights:

import torch
from rclm import Model

RNAlm = Model.RNA2nd(lmcfg)
RNAlm.load_state_dict(torch.load('weight/epoch_67000', map_location='cpu'), strict=False)
RNAlm.to(device)
RNAlm.eval()

Input

RNA2nd.embedding(in_dict) takes a dict with:

• aa: L x 5 one-hot tensor over {A, G, C, U, -} (- = gap/unknown)
• idx: L LongTensor of residue indices (1-based sequence position, used for relative positional encoding)
• mask: L FloatTensor, 1 marks a masked/unknown position, 0 otherwise

Output

s, z = RNAlm.embedding(in_dict)
# s: L x s_dim     (512) per-residue embedding
# z: L x L x z_dim (128) pairwise embedding

These embeddings can be used by downstream folding/structure/function models

Pairwise representation

Unlike a plain per-residue (sequence-only) language model, RCLM explicitly maintains and updates a full L x L x 128 pairwise representation z alongside s, via triangle-update/triangle-attention layers borrowed from the Evoformer design (rclm/EvoPair.py). This z tensor captures residue-residue coupling — the RNA analogue of a co-evolution/contact signal a single sequence normally lacks — and is exactly what downstream structure modules (e.g. DRfold2) consume to initialize their own pair/distance-map representations (lm_layer_z in EvoMSA2XYZ.py). This pairwise output is the main thing that distinguishes RCLM's embedding from a standard single-sequence LM embedding.

Training

RCLM is pretrained on ~30M RNA sequences from RNAcentral (Release 22) with a masked composite-likelihood objective, not standard per-token masked-language-modeling. Concretely, the loss is a negative log-composite-likelihood over randomly masked position subsets, and it includes a second-order term: in addition to predicting each masked nucleotide individually (first-order/unary term), the model also predicts the joint identity of pairs of masked positions (second-order/pairwise term). Combining both terms lets the pairwise representation z learn residue-residue coupling directly during pretraining, rather than only as a side effect of the per-residue objective. See the DRfold2 paper (linked above) for the full derivation.

The checkpoint shipped here (weight/epoch_67000) corresponds to the architecture in the lmcfg above (18 Evoformer blocks, s_dim=512, z_dim=128, ~47.5M parameters), trained for 67,000 batches (batch size 128) on a single GPU.

Example usage

Run example.py for a minimal, runnable end-to-end example (load the checkpoint, embed a sequence, print the s/z shapes):

python example.py cpu

Note: RNALM2 inside DRfold2 is the same model code as rclm/ here (module renamed only), so the DRfold2 config doubles as a working reference implementation for this repo. Also note seq_idx is 1-based (np.arange(len(seq)) + 1, as in DRfold2/cfg_95/test_modeldir.py) — example.py follows this convention.