RNA_model

EviRNA-NOVA

EviRNA-NOVA is a multimodal RNA 3D structure prediction research pipeline that combines nucleotide sequence information, RibonanzaNet2 representations, MSA-derived features, secondary-structure/base-pair priors, a Pairformer-style backbone, and flow-based coordinate generation.

The current codebase is designed primarily for:

• Stanford RNA 3D Folding–style datasets.
• Multi-stage training on NVIDIA T4/L4 GPUs.
• Checkpoint-based training continuation.
• Standalone inference on a new RNA test folder.
• Generation of multiple 3D conformers for each RNA.

Status: Research / experimental. This is not a clinical model or a validated RNA structure tool for high-stakes biological decision-making.

---

1. Model Objective

The main input is an RNA sequence:

AUGCGGAUACCGUAAGCU...

EviRNA-NOVA builds multiple evidence sources:

RNA sequence
    │
    ├── Sequence embedding
    ├── RibonanzaNet2 single features
    ├── RibonanzaNet2 pair features
    ├── MSA features
    ├── Secondary-structure / base-pair priors
    └── Template evidence (when available)
            │
            ▼
     Dynamic Evidence Router
            │
            ▼
        Pairformer
            │
            ▼
      Geometry heads
            │
            ▼
     E(3)-aware / flow module
            │
            ▼
       3D coordinates

The final output is one or more RNA structures:

[N, 3]

where N is the number of nucleotides and the three dimensions correspond to x, y, z coordinates.

---

2. High-Level Architecture

2.1 Sequence Branch

The RNA sequence is encoded into nucleotide indices and residue-level embeddings.

Main nucleotide symbols:

A
C
G
U
N / unknown

The sequence branch provides residue representations and basic pair priors.

---

2.2 RibonanzaNet2 Branch

RibonanzaNet2 is used as a frozen feature extractor.

For an RNA of length N, the pipeline produces:

single feature : [N, D_single]
pair feature   : [N, N, D_pair]

The current NOVA configuration typically uses projected dimensions:

D_single = 128
D_pair   = 16

RibonanzaNet2 features must be generated for the actual new RNA sequences during inference. They must not be replaced with random or reused features from unrelated targets.

---

2.3 MSA Branch

If the test folder contains:

MSA/

the pipeline can use multiple sequence alignments to construct:

• nucleotide frequencies,
• conservation statistics,
• entropy,
• co-variation / pair compatibility,
• base-pair compatibility.

If no MSA is available, the inference-only pipeline can use a query-only fallback so that the model can still run.

---

2.4 Secondary-Structure / Base-Pair Evidence

The pipeline constructs priors such as:

A-U
U-A
G-C
C-G
G-U
U-G

These can be combined with MSA compatibility and an external base-pair map when available.

This information becomes an independent evidence source for each nucleotide pair (i, j).

---

2.5 Dynamic Evidence Trust / Router

The model does not simply concatenate all modalities.

For each nucleotide pair (i, j), the router learns weights for evidence sources such as:

Sequence
RibonanzaNet2
MSA
Template
Secondary structure

Conceptually:

[ z_{ij} = \sum_k w^{(k)}{ij} E^{(k)}{ij} ]

with:

[ \sum_k w^{(k)}_{ij} = 1 ]

This allows the model to dynamically change how much it trusts each modality for different nucleotide pairs.

---

3. Pairformer Backbone

The fused pair representation is processed through multiple Pairformer blocks.

The backbone may include:

• pair transitions,
• triangle-style updates,
• axial attention,
• node-to-pair and pair-to-node interactions,
• recycling.

A typical Tesla T4 profile is approximately:

d_node       ≈ 160
d_pair       ≈ 96
pair blocks  ≈ 8
recycles     ≈ 2

Larger profiles can use wider representations and more blocks on stronger GPUs.

---

4. Geometry Heads

The Pairformer representation is supervised through several geometry-related tasks:

distance
distogram
contact
base-pair
stacking
orientation
pair confidence

For example, the predicted pairwise distance is:

[ \hat d_{ij} ]

and the ground-truth distance is:

[ d_{ij} = |x_i - x_j| ]

The distogram head predicts a probability distribution over distance bins rather than only a single continuous value.

---

5. Coordinate Generation with Flow

NOVA uses a flow-style coordinate generation module.

During training:

native coordinates x1
        +
starting/noisy coordinates x0
        ↓
intermediate xt
        ↓
predict velocity
        ↓
flow loss

A simple interpolation is:

[ x_t = (1-t)x_0 + tx_1 ]

with target velocity:

[ v^* = x_1 - x_0 ]

The model learns:

[ v_\theta(x_t,t) ]

to move a noisy or initial structure toward the target RNA structure.

---

6. Training Curriculum

The default training curriculum contains five stages:

Default total:

48 epochs

If the code is launched with --epochs N, the curriculum may be proportionally rescaled.

---

7. Mixed Precision and Numerical Stability

The implementation is designed to avoid several common mixed-precision failures on Tesla T4 GPUs.

Neural Layers

FP16 / AMP

is used for Transformer, Pairformer, and MLP operations to reduce memory usage.

Geometry-Sensitive Operations

The following operations are forced to FP32:

SVD
Kabsch alignment
torch.cdist
coordinate geometry
distance-target construction

This avoids failures such as:

svd_cuda_gesvdjBatched not implemented for 'Half'

Confidence-related binary classification uses:

binary_cross_entropy_with_logits

instead of applying BCE directly to a sigmoid output inside autocast.

---

8. Gradient Checkpointing

Pairformer blocks support activation checkpointing to reduce VRAM usage.

An important checkpointing issue was addressed in later versions: checkpoint closures should not capture loop variables such as blk through a lambda.

Avoid:

checkpoint(lambda x, z: blk(x, z, ...), ...)

when the captured block may change before backward recomputation.

The implementation instead checkpoints the actual module/block directly so that forward and backward recomputation use the same graph.

---

9. Training Data

The training pipeline was originally designed around files such as:

train_sequences.csv
train_labels.csv
validation_sequences.csv
validation_labels.csv
MSA/
PDB_RNA/
extra/

The split policy is:

1. Split at the original RNA target level first.
2. Do not crop before splitting.
3. Long RNAs are cropped/windowed only after they belong to FIT or CALIB.
4. Crops from the same source RNA must not appear in both FIT and CALIB.

This is intended to reduce data leakage.

---

10. Long-RNA Handling

Pairwise representations become expensive as sequence length grows because memory scales approximately with:

[ O(N^2) ]

Training commonly uses:

crop_len    = 256
crop_stride = 192

For inference, long RNAs can be processed using a windowed strategy:

1. Generate new RibonanzaNet2 features for the RNA.
2. Split NOVA inference into overlapping windows.
3. Generate coordinates for each window.
4. Align overlapping regions using Kabsch alignment.
5. Stitch the windows into a full-length structure.

---

11. Standalone Inference on a New Folder

Recommended inference scripts include:

EviRNA_NOVA_INFERENCE_ONLY_NEW_FOLDER_V19.py

or the direct-path variant:

EviRNA_NOVA_INFERENCE_DIRECT_PATHS_V20.py

Inference-only execution does not require:

NOVA_history.csv
split_fit_targets.csv
split_calib_targets.csv
curriculum_fit_windows.csv
curriculum_calib_windows.csv
training evidence_cache.pkl
optimizer state
scheduler state

Only the following are required:

1. new test folder
2. NOVA checkpoint
3. RibonanzaNet2 checkpoint
4. output directory / GPU runtime

---

12. New Test Folder

Minimum structure:

MY_NEW_TEST/
└── test_sequences.csv

Recommended structure:

MY_NEW_TEST/
├── test_sequences.csv
├── MSA/
│   ├── RNA001.a3m
│   ├── RNA002.a3m
│   └── ...
├── BP/                   # optional
└── sample_submission.csv # optional

Example test_sequences.csv

target_id,sequence
RNA001,GGGAAACCCUUUGGG
RNA002,AUGCGGAUACCGUAAGCU
RNA003,GGCUAUAGCUCAGUUGGU

---

13. NOVA Checkpoint

The trained model checkpoint is typically named:

NOVA_last.pt

or:

NOVA_best.pt

For inference, the checkpoint should contain at least:

state_dict
nova_cfg
foundation_single_dim
foundation_pair_dim

Additional fields such as:

optimizer_state
scheduler_state
scaler_state
stage
stage_epoch

are useful for training continuation but are not required for inference.

---

14. RibonanzaNet2 Checkpoint

The NOVA checkpoint does not include the full RibonanzaNet2 model.

Inference on new RNA sequences therefore requires a separate RibonanzaNet2 checkpoint.

A compatible bundle may look like:

ribonanzanet2_checkpoint/
├── Network.py
├── pairwise.yaml
└── pytorch_model_fsdp.bin

If Kaggle is offline, this bundle must be added as an input before inference.

---

15. Direct-Path Inference

The script:

EviRNA_NOVA_INFERENCE_DIRECT_PATHS_V20.py

is designed so that the main paths are edited directly near the top of the file.

Example:

DIRECT_TEST_FOLDER = (
    "/kaggle/input/datasets/USERNAME/MY_NEW_TEST"
)

DIRECT_NOVA_CHECKPOINT = (
    "/kaggle/input/datasets/USERNAME/NOVA_MODEL/NOVA_last.pt"
)

DIRECT_RNET2_CHECKPOINT = (
    "/kaggle/input/datasets/USERNAME/RIBONANZANET2/"
    "ribonanzanet2_checkpoint"
)

DIRECT_OUT_DIR = (
    "/kaggle/working/EVIRNA_NOVA_INFERENCE"
)

Then run:

!CUDA_VISIBLE_DEVICES=0 \
python EviRNA_NOVA_INFERENCE_DIRECT_PATHS_V20.py

---

16. Inference Outputs

A typical output directory is:

EVIRNA_NOVA_INFERENCE/
│
├── test_manifest.csv
│
├── rnet2_features_new_test/
│   ├── RNA001.npz
│   ├── RNA002.npz
│   └── ...
│
├── evidence_new_test.pkl
├── prediction_top5.csv
├── submission.csv
├── prediction_summary.csv
├── router_audit.csv
└── inference_manifest.json

prediction_top5.csv

Each nucleotide can have five candidate coordinate sets:

ID
resname
resid
x_1 y_1 z_1
x_2 y_2 z_2
x_3 y_3 z_3
x_4 y_4 z_4
x_5 y_5 z_5

prediction_summary.csv

This file may include:

target_id
length
inference_mode
number_of_windows
candidate_score_1
...
candidate_score_5

router_audit.csv

This file can be used to inspect evidence usage:

sequence
RibonanzaNet2
MSA
template
secondary_structure

---

17. RibonanzaNet2 Cache for New Test Data

RibonanzaNet2 cache from an old dataset must not be reused for unrelated new RNA sequences.

The inference pipeline creates:

OUT_DIR/
└── rnet2_features_new_test/

Each RNA receives a file such as:

RNA001.npz

containing:

single
pair

If inference is interrupted, valid feature files can be reused on the next run.

To force complete feature regeneration, enable the corresponding rebuild option, for example:

force_rnet2_rebuild = True

---

18. One Tesla T4 GPU

Inference can be run with:

CUDA_VISIBLE_DEVICES=0

Recommended settings are approximately:

AMP                   ON
NOVA batch            1
Pairformer window     256
RNet2 adaptive chunk  128–192
Conformers            5
Flow steps            48

Very long RNAs may still require aggressive chunking because pairwise representations scale quadratically with sequence length.

---

19. Multiple GPUs / L4

Two T4 GPUs do not behave as one unified 32 GB GPU.

A single RNA target is still limited by the VRAM of the GPU processing that target.

For many independent targets, inference can be sharded:

GPU0 -> targets 0,2,4,...
GPU1 -> targets 1,3,5,...

For 4×L4 training, DistributedDataParallel is preferable to nn.DataParallel.

---

20. Supported Workflows

Fresh Training

data
→ build evidence
→ RibonanzaNet2
→ S1-S5
→ NOVA_best.pt

Training Continuation

NOVA_last.pt
→ restore model/optimizer/scheduler/scaler
→ continue stage/epoch

Inference Only

new test folder
→ new RibonanzaNet2 features
→ NOVA checkpoint
→ 5 conformers

If the goal is only to predict new RNA data, the inference-only script should be used instead of resume-training scripts.

---

21. Major Technical Issues Addressed

Several implementation issues were identified and fixed during development.

AMP BCE

Unsafe pattern:

binary_cross_entropy(sigmoid(logits), target)

Preferred:

binary_cross_entropy_with_logits(logits, target)

---

CUDA SVD in FP16

Unsafe on some Tesla T4 configurations:

torch.linalg.svd(H_half)

Preferred:

torch.linalg.svd(H.float())

---

NaN Coordinates

Multiplying invalid values by zero is not sufficient because:

NaN * 0 = NaN

Coordinates must be sanitized before:

cdist
Kabsch
geometry losses

---

Checkpoint Closure Capture

Avoid checkpointing closures that may capture the wrong block:

checkpoint(lambda ...: blk(...))

inside a changing loop.

---

RibonanzaNet2 OOM

CUDA OOM must not be treated as an API mismatch.

Once the correct RibonanzaNet2 encoder API is identified, OOM should trigger:

re-raise / adaptive chunking

rather than repeatedly trying unrelated call signatures.

---

22. Reproducibility

For reproducible inference, save:

NOVA checkpoint
RibonanzaNet2 checkpoint
inference script
test_sequences.csv
MSA files
random seed
PyTorch version
CUDA version
GPU name
prediction_summary.csv
inference_manifest.json

---

23. Current Limitations

Current limitations include:

1. Pair representations have high memory cost for long RNAs.
2. Missing MSA/template evidence may reduce prediction quality.
3. Window stitching is an approximation and is not equivalent to full global Pairformer inference.
4. Candidate confidence/ranking is not a calibrated probability that a structure is correct.
5. The pipeline is experimental and should be benchmarked carefully before strong scientific conclusions are drawn.
6. If a model was trained with template evidence but inference is performed with no template information, this introduces distribution shift.
7. RibonanzaNet2 preprocessing and feature dimensions must remain compatible with the NOVA checkpoint.

---

24. Recommended Inference Workflow

1. Prepare test_sequences.csv
          ↓
2. Add MSA if available
          ↓
3. Load RibonanzaNet2
          ↓
4. Extract NEW RibonanzaNet2 features
          ↓
5. Build SS/BP evidence
          ↓
6. Load NOVA checkpoint
          ↓
7. Run direct or windowed inference
          ↓
8. Generate 5 conformers
          ↓
9. Rank/diversify conformers
          ↓
10. Save prediction_top5.csv

---

25. Summary

EviRNA-NOVA can be summarized as:

[ \boxed{ \text{Sequence} + \text{RibonanzaNet2} + \text{MSA} + \text{SS/BP} + \text{Template (optional)} \rightarrow \text{Dynamic Evidence Fusion} \rightarrow \text{Pairformer} \rightarrow \text{Geometry} \rightarrow \text{Flow} \rightarrow \text{RNA 3D} } ]

The core design goal is to combine multiple evidence sources rather than relying on a single modality, while producing multiple candidate RNA structures for top-k structural prediction.

---

License / Citation

The current pipeline is an internal / experimental research implementation.

Before public release, the repository should include:

• a project license,
• citation information for the Stanford RNA 3D Folding dataset,
• citation and license information for RibonanzaNet2,
• citations for any external checkpoints or components used,
• a clear description of which parts are original EviRNA-NOVA contributions.