Order-Agnostic Decoding for Sample-Efficient RNA Inverse Folding

Motivation

• Known: RNA 2차 구조 역설계는 RNA 합성생물학의 핵심 문제로, 현재 최고 성능(Struct2SeQ 등)은 L→R autoregressive 디코더와 대규모 best-of-N 샘플링(예: K=98,000, 256 A100 GPU)에 의존해 개별 샘플의 낮은 성공률을 보완하고 있다. 단백질 역설계(ProteinMPNN, ESM-IF)와 달리 RNA는 4-letter 알파벳과 평평한 열역학 지형으로 인해 서열-구조 매핑이 더 퇴화되어 있어 이러한 레시피가 그대로 전이되지 않는다.
• Gap: L→R autoregressive 생성은 downstream pairing partner를 관측하기 전에 upstream nucleotide를 고정시켜 버리므로 pseudoknot과 같은 장거리 페어링을 다루는 데 근본적 한계가 있고, 부분 서열 조건화(in-painting, 모티프 보존/재설계)를 원천적으로 지원할 수 없다는 구조적 갭이 존재한다.
• Why: Best-of-N에 의존하는 현재 패러다임은 계산 비용이 막대하고(수십만 샘플, 수백 GPU) 제약이 추가될수록 성공률이 곱셈적으로 감소해 긴 서열이나 제약이 많은 설계 문제로 확장하기 어렵다. 샘플 효율성을 근본적으로 높이고 in-painting 기반의 압타머·리보스위치 설계를 가능하게 하는 것은 RNA 합성생물학 응용에 실질적으로 중요하다.
• Approach: Struct2SeQ의 L→R autoregressive 디코더를 decoding-order permutation에 대해 학습된 순수 Transformer 기반 order-agnostic 디코더로 교체하고, 학습된 relative-position bias를 permutation σ에 keyed하여 위치와 순환성을 분리함으로써 동일 체크포인트로 전체 설계, 모티프 보존, 모티프 재설계 세 가지 태스크를 모두 지원한다.

Achievement

Figure 1. One order-agnostic decoder, three inverse-design modalities. Because the decoder is trained to generate nucleo

1. 샘플 효율성 대폭 향상: OpenKnot Round 7b 240-mer pseudoknot 벤치마크에서 매칭된 K에서 최고 AR 베이스라인 대비 5배 많은 perfect-structure 설계를 샘플당 생성함.
2. 토큰 샘플링 없는 다양성 확보: argmax-only 추론에서도 random order decoding을 통해 다양성을 복원, best-of-N에서 쓰이는 토큰 수준 샘플링 전략 없이도 다양한 설계 생성.
3. 분포 외(out-of-distribution) 전이: 재학습 없이 Eterna100 벤치마크로 전이하여 64/100 퍼즐 해결(ViennaRNA-2 기준), 매칭 예산에서 base model 샘플링 대비 우수.
4. 네이티브 in-painting: 단일 체크포인트로 모티프 보존과 모티프 재설계를 추론 시점 연산으로 지원하며 후자는 AR 디코딩에서는 원천적으로 불가능함을 입증.
5. 메커니즘적 설명 제공: post-hoc rescue가 AR 디코더(국소적으로 수정 가능한 over-pair 오류)에는 도움이 되지만 random-permutation 디코더(문맥 의존적 under-pair 오류)에는 도움이 되지 않는 비대칭성을 규명함.

How

Figure 2. Order-agnostic decoder architecture. The structure-conditioned decoder predicts nucleotides in a randomly perm

• Struct2SeQ의 encoder-decoder 구조, RibonanzaNet-SS 기반 dense per-position reward, twice-shifted deep-Q-learning 학습 신호, Hungarian-matched Jaccard scoring을 그대로 채택
• 기존 LSTM-PE + causal-conv 디코더를 permutation σ에 keyed된 learned relative-position bias를 갖는 순수 Transformer 디코더로 교체
• 학습 시 매 예제마다 σ ∼ Uniform(S_L)을 샘플링하고 causal mask를 RNA 위치 순서가 아닌 step 순서로 적용
• 공개된 Struct2SeQ 체크포인트에서 이어서 4×A100 80GB GPU, bf16, FlashAttention/SDPA, torch.compile로 5 episode 추가 학습(약 7배 처리량 향상)
• 추론 시 Algorithm 1(random-permutation argmax decoding)로 매 샘플마다 새로운 permutation을 뽑아 디코딩
• in-painting은 Algorithm 2와 같이 고정할 nucleotide를 디코딩 시작 전 출력 버퍼에 pre-fill하여 partner-mask를 통해 위치와 무관하게 조건화
• OpenKnot Round 7b 240-mer 20개 퍼즐과 Eterna100에서 K~1000 샘플, RibonanzaNet-SS 예측 구조를 target과 Hungarian-matched Jaccard로 채점, perfect-Jaccard rate와 solved puzzle 수를 보고
• 동일 체크포인트로 L→R+argmax vs random-perm+argmax를 비교하는 ablation으로 decoding order와 architecture 기여를 분리

Originality

• Best-of-N 스크리닝이 지배적인 RNA inverse folding 분야에서 sample volume이 아닌 sample efficiency를 핵심 축으로 재정의한 문제 설정 전환
• 단백질 역설계(ProteinMPNN 등)에서 사용되던 order-agnostic decoding 아이디어를 pseudoknot을 포함한 RNA 2차 구조 역설계에 처음 적용하고, deep-Q-learning 학습 신호와 결합
• 단일 체크포인트로 전체 설계·모티프 보존·모티프 재설계 세 가지를 모두 지원하는 통합 in-painting 프레임워크 제시, 특히 모티프 재설계는 L→R AR 디코딩에서는 구조적으로 불가능함을 명확히 함
• Post-hoc rescue의 효과가 AR과 random-permutation 디코더 간에 비대칭적으로 나타나는 이유(over-pair vs under-pair 오류의 국소 수정 가능성 차이)를 메커니즘적으로 설명

Limitation & Further Study

• 평가가 RibonanzaNet-SS 예측 구조와 ViennaRNA-2 기반 스코어링에 의존하며, 실제 실험적(예: SHAPE 등) 검증 없이 in-silico 지표에만 기반함
• 벤치마크 규모가 OpenKnot 20개 퍼즐과 Eterna100 100개 퍼즐로 제한적이며, 더 다양한 길이·위상의 RNA 구조에 대한 일반화는 추가 검증이 필요
• AR 베이스라인의 rescue 프로토콜 대비 coverage(puzzle 수)에서는 여전히 뒤처지는 경우가 있어(예: 매칭 프로토콜에서 AR이 2 puzzle 우위), rescue와 결합 가능한 하이브리드 전략에 대한 후속 연구 필요
• 실제 압타머/리보스위치 기능 검증(결합 친화도, 조절 기능 등) 실험 데이터가 제시되지 않아 실질적 합성생물학 응용 가능성은 향후 wet-lab 검증이 필요함

Evaluation

총평: Best-of-N 패러다임에 갇혀 있던 RNA inverse folding 분야에 order-agnostic decoding을 도입해 샘플 효율성과 in-painting 능력을 동시에 확보한 실용적이고 독창적인 연구로, 향후 압타머/리보스위치 설계로 이어질 잠재력이 크다.