STRIDE: Post-Training LLMs to Reason and Refine Bio-Sequences via Edit Trajectories

Motivation

• Known: discrete diffusion model은 이산 토큰 공간에서 iterative refinement를 지원하며 단백질 서열 생성(EvoDiff, DPLM)에서 강력한 성능을 보였고, autoregressive LLM은 SFT와 RLHF, GRPO 계열의 post-training을 통해 reasoning 및 정렬 능력이 향상되어 왔다.
• Gap: diffusion model은 명시적이고 controllable한 discrete-edit 인터페이스를 노출하지 않아 insertion/deletion 하에서 domain-specific validity를 매 중간 단계마다 강제하기 어렵고, 반면 autoregressive LLM은 다단계 constrained edit을 계획할 때 myopic하게 행동하여 tight edit budget 하의 rugged fitness landscape를 잘 탐색하지 못한다는 한계가 있다.
• Why: 단백질·분자 최적화는 치료제 개발, 효소 공학 등 실제 응용에서 de novo generation이 아닌 기존 서열에 대한 goal-directed, parser-valid한 소수의 edit을 요구하는 경우가 많아, 해석 가능하고 controllable한 discrete edit 인터페이스를 갖춘 방법이 실용적으로 중요하다.
• Approach: discrete bio-sequence 최적화를 edit space에서의 trajectory planning으로 재정식화하여, LLM이 Levenshtein alignment 기반 shortest-edit demonstration으로 SFT된 후 GRPO/GSPO 계열의 group-based policy optimization으로 task reward에 정렬되도록 two-stage로 post-training한다.

Achievement

1. Oracle 기반 protein stress test 성능 향상: full-action protein 벤치마크에서 STRIDE가 Vanilla SFT 대비 success를 42%에서 89%로, unique improvement 중 novelty를 47%에서 97%로 크게 향상시켰다.
2. 분자 편집에서의 정렬 개선: instruction-conditioned molecular editing에서 GSPO-aligned STRIDE 변형이 SFT-only STRIDE 대비 strict success, controllability, SMILES validity를 모두 개선하였다.
3. variable-length editing에서의 강점: GFP/AAV 등 protein 벤치마크와 분자 편집 벤치마크 전반에서 특히 variable-length, index-consistent editing 상황에서 이득이 가장 두드러짐을 확인하였다.

How

Figure 1. Overview of the STRIDE workflow. (A) STRIDE optimizes biological sequences through executable edit scripts rat

• Levenshtein alignment 기반 dynamic programming backtracking으로 source-target 서열 쌍으로부터 minimum-cost INSERT/DELETE/REPLACE edit script(Shortest Edit Path, SEP)를 구성하여 index-grounded trajectory supervision 데이터를 생성
• amino-acid token 및 Regex SMILES tokenizer를 사용하여 protein과 molecule 각각에 대해 executable trajectory와 indexing 체계를 정의
• Stage I: 이렇게 구축된 SEP demonstration으로 supervised fine-tuning을 수행하여 edit trajectory와 최종 서열을 동시에 모델링하고 validity 및 minimal-edit bias를 유도
• Stage II: GRPO 및 GSPO, CISPO 등 group-based policy optimization 변형을 사용해 parse-and-execute consistency 하에서 trajectory를 task reward에 맞게 정렬하며, KL regularization으로 coherent editing 행동을 유지
• GFP fluorescence(TAPE), AAV transfer 평가, instruction-conditioned molecular editing 벤치마크에서 성능을 평가

Originality

• diffusion model의 iterative refinement 장점과 autoregressive LLM의 discrete token 생성 능력을 결합하여, 별도의 stochastic transition process 대신 명시적 edit trajectory를 생성하도록 LLM을 post-training하는 새로운 관점을 제시
• Levenshtein alignment의 DP backtracking을 활용해 index-grounded edit script를 자동 생성하는 deterministic supervision 파이프라인을 고안
• GRPO/GSPO 등 reasoning-oriented post-training 기법을 bio-sequence editing이라는 새로운 도메인에 적용하고 parse-and-execute consistency 기반 reward 정렬을 결합

Limitation & Further Study

• 본문 발췌만으로는 reward 함수 설계와 oracle 모델의 세부 사항, GRPO/GSPO/CISPO 간 구체적 비교 실험 결과가 충분히 제시되지 않아 어떤 policy optimization variant가 왜 우수한지에 대한 심층 분석이 부족해 보임
• shortest edit path prior가 minimal-edit bias를 유도하는데, 이는 국소적 변형에 치우쳐 global reformulation이 필요한 경우(대규모 구조 변화가 필요한 최적화)에는 성능이 제한될 수 있음
• 평가가 GFP/AAV 및 특정 분자 편집 벤치마크에 한정되어 있어 다른 생물학적 서열 유형(RNA, 항체 등)이나 실제 wet-lab 검증으로의 일반화 가능성은 추가 검증이 필요함
• 후속 연구로 다양한 reward 설계, 더 큰 edit budget, 실제 실험적 검증(wet-lab feedback loop)과의 결합 등을 고려할 수 있음

Evaluation

총평: discrete diffusion과 autoregressive LLM 각각의 한계를 명확히 짚고 이를 explicit edit trajectory 생성이라는 아이디어로 연결한 참신하고 실용적인 post-training 프레임워크이며, protein 및 분자 편집 벤치마크에서의 실험 결과가 이를 뒷받침하지만 reward 설계와 policy optimization 세부 비교에 대한 추가 분석이 보완되면 더 완성도가 높아질 것이다.