NVIDIA BioNeMo Recipes로 LoRA를 활용한 생물학적 파운데이션 모델 파인튜닝하기

파운데이션 모델은 계산 생물학의 패러다임을 바꾸고 있습니다. 단백질이나 게놈 서열의 방대한 데이터로 사전 학습된 ESM2(단백질 언어 모델)와 Evo 2(DNA 언어 모델) 같은 모델은 생물학적 서열의 통계적 규칙성을 학습합니다. 이러한 특성은 구조 예측, 변이 효과 분석, 기능 주석을 비롯한 다양한 다운스트림 태스크에서 뛰어난 성능으로 이어집니다.

그러나 이러한 모델을 특정 태스크에 적용하는 것은 간단하지 않습니다. 수십억 개의 파라미터를 가진 모델을 전체 파인튜닝하는 것은 컴퓨팅 자원과 옵티마이저 상태 및 체크포인트 저장 측면에서 금세 한계에 부딪힙니다.

저랭크 적응(LoRA)은 이 문제를 직접 해결합니다. 사전 학습된 백본을 고정한 채 소규모의 저랭크 어댑터 행렬만 학습함으로써, LoRA는 파라미터의 약 1%만 학습하면서도 많은 태스크에서 전체 파인튜닝과 동등한 품질을 달성하고, 단일 워크스테이션 GPU에서 수십억 규모의 모델과 어댑터 상태를 모두 수용할 수 있습니다.

이러한 워크플로우 구축의 어려움을 줄이기 위해, NVIDIA BioNeMo Recipes는 친숙한 PyTorch, Hugging Face, Megatron-Bridge 패턴을 기반으로 한 단계별 학습 레시피를 제공합니다. NVIDIA Transformer Engine(TE)과 같은 성능 지향 컴포넌트 및 스케일아웃 전략이 적재적소에 통합되어 있으면서도, 레시피 자체는 읽기 쉽게 유지됩니다.

이 글에서는 단일 NVIDIA RTX 6000 Blackwell Workstation Edition GPU에서 동일한 파라미터 효율적 레시피가 두 가지 생물학적 모달리티에 어떻게 적용되는지 보여주는 두 가지 사례 연구를 소개합니다.

• ESM2-3B + LoRA를 활용한 단백질 2차 구조 예측(PSSP)
• Evo2-1B + LoRA를 활용한 DNA 스플라이스 사이트 분류

이 결과를 커스터마이즈하거나 재현하는 데 필요한 모든 소스 코드는 NVIDIA BioNeMo Recipes에서 확인할 수 있습니다.

LoRA가 대규모 파인튜닝을 가능하게 하는 방법

사례 연구에 앞서 해당 방법론을 간단히 짚고 넘어가겠습니다. 전체 파인튜닝은 모든 모델 파라미터와 옵티마이저 상태를 저장하고 업데이트해야 하기 때문에 자원이 많이 소모되며, 모델 규모가 커질수록 금세 한계에 부딪힙니다.

LoRA는 모든 모델 파라미터의 옵티마이저 상태를 업데이트하거나 저장하지 않고도 대규모 사전 학습 트랜스포머를 파인튜닝하는 실용적인 방법입니다. LoRA의 핵심 아이디어는 밀집 모델의 가중치 행렬 \(W\)를 직접 업데이트하는 대신, 병렬로 새로운 학습 가능한 저랭크 행렬 \(W=BA\)를 추가하고 \(W\)는 고정한 상태로 유지하는 것입니다. 이를 통해 학습 가능한 파라미터 수와 옵티마이저 및 메모리 사용량을 크게 줄일 수 있습니다.

LoRA는 용량, 안정성, 비용 간의 균형을 조절하는 소규모 하이퍼파라미터 집합으로 파라미터화됩니다. 랭크 \(r\)은 추가된 저랭크 행렬의 크기, 즉 학습 가능한 파라미터 수를 제어합니다. 타깃 모듈은 어느 레이어에 어댑터를 적용할지를 지정하며, 어텐션 및 MLP 프로젝션이 일반적인 선택입니다. 소규모 데이터셋의 경우, 추가적인 정규화 방법으로 LoRA 드롭아웃을 활성화할 수 있습니다.

아래 두 사례 연구는 모달리티(단백질 대 DNA), 태스크 유형(토큰 분류 대 시퀀스 분류), 기반 아키텍처(트랜스포머 대 striped Hyena)가 다르지만, 모두 동일한 LoRA 레시피 패턴을 사용합니다.

단백질 2차 구조 예측을 위한 ESM2-3B

PSSP는 단백질 서열의 각 아미노산에 구조적 레이블을 할당하는 태스크입니다. 2차 구조 레이블은 완전한 3D 구조 예측 없이도 나선(helix)과 가닥(strand) 등 국소 백본 형태를 기술합니다. 많은 단백질에서 이러한 국소 패턴은 기능적 모티프 및 전체 접힘 구성과 상관관계가 있습니다.

PSSP는 생물학의 다양한 다운스트림 응용을 위한 핵심 구성 요소입니다. 국소 구조는 단백질 기능과 강하게 연관되어 있으므로, PSSP는 유용한 기능적 맥락을 제공할 수 있습니다. 또한 이 예측은 3차 구조 예측, 용매 접근성 예측, 단백질-단백질 상호작용 예측, 구조적 클래스 또는 도메인 관련 예측에도 활용될 수 있습니다.

모델링 측면에서 PSSP는 토큰 분류 문제입니다. 입력값은 아미노산 서열이며, 출력값은 각 잔기에 대한 구조적 레이블입니다.

일반적으로 사용되는 두 가지 평가 방식은 레이블 공간만 다릅니다.

• Q3 (3-state): H (Helix), E (Strand/Sheet), C (Coil/Loop)
• Q8 (8-state): H (α-helix), B (β-bridge), E (β-strand),  G (3₁₀ helix), I (π-helix), T (turn), S (bend), C (coil/other)

ESM2-3B는 30억 파라미터의 단백질 언어 모델로, 전체 파인튜닝에는 상당한 컴퓨팅 자원과 메모리가 필요합니다. LoRA는 소수의 추가 파라미터만 학습하면서도 PSSP에서 강력한 성능을 달성함으로써 적응을 실용적으로 만들어 줍니다.

BioNeMo Recipes의 ESM2 + PEFT (TE 가속)

팀은 가벼운 잔기별 분류 헤드(Q3/Q8 레이블용)를 추가하고 PEFT 라이브러리를 통해 LoRA 어댑터를 학습하면서 사전 학습된 백본 가중치는 고정한 채 PSSP를 위한 ESM2-3B를 파인튜닝했습니다. 데이터는 Porter 6 모델 저자들이 공개한 정제된 분할을 사용했으며, 제공된 테스트 세트에서 결과를 보고했습니다. 처리량을 극대화하기 위해 TE와 시퀀스 패킹을 활성화하고, 단일 NVIDIA RTX 6000 Blackwell Workstation Edition GPU에서 1시간 이내에 전체 학습 워크플로를 실행했습니다.

다음 코드는 BioNeMo Recipes ESM2 + PEFT 예제에서 발췌한 것으로, TE 호환 ESM2 모델을 로드하고 융합된 쿼리/키/값(QKV) 프로젝션에 LoRA 어댑터를 연결하는 방법을 보여줍니다.

import peft
import torch
from transformers import AutoConfig, AutoModelForTokenClassification

# Load config and token-classification model (use a local checkpoint path or HF model ID, e.g. nvidia/esm2_t36_3B_UR50D).
config = AutoConfig.from_pretrained("nvidia/esm2_t36_3B_UR50D", trust_remote_code=True)
model = AutoModelForTokenClassification.from_pretrained(
    "nvidia/esm2_t36_3B_UR50D", config=config, trust_remote_code=True, dtype="bfloat16"
)

peft_config = peft.LoraConfig(
    task_type=peft.TaskType.TOKEN_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=16,
    target_modules=["layernorm_qkv"],
    bias="none",
)

peft_model = peft.get_peft_model(model, peft_config)
peft_model.to("cuda", dtype=torch.bfloat16)

이 PEFT 모델을 학습 루프에 연결하면 됩니다. 전체 레시피에는 데이터로더, 손실 함수, 옵티마이저 설정이 포함되어 있습니다.

표 1은 Porter 6 논문에 보고된 강력한 공개 베이스라인과 함께 ESM2-3B + LoRA 모델의 Q3/Q8 테스트 정확도를 요약합니다. 표 1은 ESM2-3B의 상위 5개 검증 체크포인트 평균 점수를 보고합니다.

전반적으로 LoRA 파인튜닝은 최신 PSSP 접근 방식과 경쟁력 있는 정확도를 달성합니다. 그림 2는 파인튜닝 단계별 검증 손실 및 정확도를 보여줍니다.

시퀀스 패킹은 어떻게 GPU 활용률과 처리량을 높이는가?

단백질 데이터셋은 일반적으로 길이가 다양한 서열을 포함합니다. 단순하게 배치로 묶으면(패딩된 BSHD 형식) 배치 내 최대 길이에 맞춰 패딩되어 많은 토큰이 패딩 토큰이 됩니다. 이는 어텐션 및 MLP 레이어 내의 컴퓨팅과 메모리 대역폭을 낭비합니다.

시퀀스 패킹(패킹된/평탄화된 THD 형식)은 패딩이 아닌 토큰만 연결하고 누적 길이 메타데이터로 서열 경계를 추적함으로써 이러한 낭비를 줄입니다. 그 결과 어텐션/MLP 커널이 패딩 토큰 대신 실제 토큰에 대해 동작합니다. 패킹이 실제로 어떻게 작동하는지(그리고 TE 패킹 형식과 어떻게 상호작용하는지)에 대한 자세한 설명은 Scale Biology Transformer Models with PyTorch and NVIDIA BioNeMo Recipes를 참고하세요. 그림 3은 THD 대 BSHD로 파인튜닝할 때의 처리량(토큰/초)을 보여줍니다.

이 설정에서 BSHD에서 THD로 전환하면 토큰/초가 약 5.5배 향상되었으며, 이는 주로 패딩 오버헤드 제거 덕분입니다. 실제 속도 향상은 서열 길이 분포, 마이크로배치 크기, GPU에 따라 달라집니다.

처리량 외에도 THD 패킹은 메모리 효율을 개선합니다. 패딩 토큰에 소비되는 활성화 또는 어텐션 작업량을 줄여, GPU 메모리 트래픽과 컴퓨팅의 더 많은 비중이 유효한(비패딩) 토큰에 사용됩니다.

동일한 입력 서열과 배치 크기에서 THD는 패딩 토큰 구체화를 피하기 때문에 일반적으로 BSHD보다 메모리를 덜 사용합니다. 실제로 이렇게 확보된 여유 공간은 단계별로 처리하는 실제 토큰 수를 늘리는 데 활용됩니다.

DNA 스플라이스 사이트 분류를 위한 Evo2-1B

Evo 2는 모든 생명 도메인에 걸친 게놈 서열로 학습된 생성형 DNA 파운데이션 모델입니다. 아키텍처적으로는 상태 공간 방식의 장거리 컨볼루션 연산자와 소수의 어텐션 레이어가 혼합된 striped Hyena 블록을 기반으로 합니다. 이를 통해 긴 DNA 컨텍스트를 효율적으로 처리할 수 있습니다. ESM2가 아미노산 서열에서 단백질 “문법”을 학습하는 것처럼, Evo2는 뉴클레오타이드 서열에서 직접 게놈의 규칙성을 학습하며, 이는 변이 효과 예측, 조절 요소 분류, (이 글의 주제인) 스플라이스 사이트 식별 등 다양한 다운스트림 태스크에 전이됩니다.

스플라이스 사이트 분류란 무엇인가?

스플라이싱(splicing)은 pre-mRNA에서 인트론(intron)을 제거하고 엑손(exon)을 연결하는 세포 내 과정입니다. 경계는 두 가지 짧은 서열 모티프로 정의됩니다. 공여 사이트(인트론 시작, 일반적으로 GT)는 인트론의 5′ 말단에, 수용 사이트(인트론 끝, 일반적으로 AG)는 인트론의 3′ 말단에 위치합니다.

원시 DNA에서 이러한 사이트를 식별하는 것은 단순히 다이뉴클레오타이드 모티프를 매칭하는 것보다 훨씬 어렵습니다. 동일한 GT/AG 패턴은 게놈 전체에 걸쳐 나타나며, 그 중 극히 일부만이 기능적 스플라이스 사이트입니다. 유용한 예측기는 후보 위치 주변의 더 넓은 컨텍스트를 학습해야 합니다.

Nucleotide Transformer 다운스트림 태스크 데이터셋의 splice_sites_all 태스크를 사용했습니다. 각 예제는 길이가 고정된 600 bp DNA 윈도우이며, 레이블은 중앙 위치를 설명하는 세 가지 클래스(비스플라이스, 수용, 공여) 중 하나입니다. 이 벤치마크는 약 3만 건의 학습 / 3천 건의 테스트 예제로 구성되며, 클래스 균형이 대체로 맞춰져 있습니다.

모델링 수준에서 이는 시퀀스 분류 문제입니다. PSSP의 토큰별 레이블과 달리, 입력 서열당 단일 레이블을 출력합니다.

BioNeMo Recipes의 Evo2 + LoRA

팀은 Megatron Hyena 모델을 서브클래싱해 평균 풀링된 은닉 상태 위에 소규모 시퀀스 분류 헤드를 추가하고, 백본 어텐션, MLP, Hyena 믹서 프로젝션에 LoRA 어댑터를 학습해 스플라이스 사이트 분류를 위한 Evo2-1B를 파인튜닝했습니다. 사전 학습된 백본 가중치는 고정 상태로 유지되었으며, LoRA 어댑터와 분류 헤드만 학습되었습니다.

LoRA의 기여를 맥락에서 파악하기 위해, 동일한 데이터로 두 가지 구성을 학습하고 비교했습니다.

• 헤드만 학습 베이스라인: 백본 고정; 어댑터 없이 분류 헤드만 학습 가능. 학습 가능 파라미터 합계: 약 370만 개 (모델의 0.33%)
• LoRA + 헤드: 백본 고정; 지정된 타깃 모듈에 LoRA 어댑터 적용, 분류 헤드 학습 가능. 학습 가능 파라미터 합계: 약 1,600만 개 (모델의 1.42%)

표 2는 홀드아웃 3천 건의 테스트 정확도를 보여줍니다.

차이는 매우 큽니다. 파라미터의 약 1%만 학습 가능한 상태에서 LoRA는 사전 학습된 Evo2 백본이 스플라이싱에 대해 보유한 신호의 거의 전부를 회복하는 반면, 풀링만으로는 충분하지 않습니다. LoRA 모델의 잔여 오류 대부분은 공여↔수용 방향에서 발생합니다. 두 모티프 모두 GT/AG 다이뉴클레오타이드 구조를 공유하고 있어 구별하려면 더 넓은 컨텍스트가 필요하기 때문에 예상된 결과입니다.

전체 워크플로는 단일 RTX 6000 Workstation Edition에서 약 1시간 내에 엔드 투 엔드로 실행됩니다.

다음 코드는 ESM2 예제와 스타일적으로 유사하게 Evo2 백본을 로드하고, Hyena Model 서브클래스를 통해 분류 헤드를 연결하며, 어텐션, MLP, Hyena 믹서 프로젝션에 LoRA 어댑터를 구성합니다.

from bionemo.evo2.models.evo2_lora import Evo2LoRA
from evo2_classifier import (
    Hyena1bClassifierProvider,
    HyenaForSequenceClassification,
)

# Backbone provider: a HyenaModel subclass with a small classification head
# (LayerNorm → Linear → GELU → Dropout → Linear) on top of mean-pooled hidden states.
model_provider = Hyena1bClassifierProvider(
    num_classes=3,            # no-splice / acceptor / donor
    classifier_dropout=0.1,
    pool="mean",
)

# LoRA adapters on attention (linear_qkv, linear_proj), MLP (linear_fc1, linear_fc2),
# and the Hyena mixer (dense_projection, dense). The classification head is kept
# trainable via the skip_freeze_modules pattern; everything else is frozen.
peft = Evo2LoRA(
    target_modules=[
        "linear_qkv", "linear_proj",
        "linear_fc1", "linear_fc2",
        "dense_projection", "dense",
    ],
    dim=16,
    alpha=32,
    dropout=0.1,
    skip_freeze_modules=["*classification_head*"],
)

Megatron-Bridge 사전 학습 진입점은 분산 학습, 옵티마이저, 스케줄러, 체크포인팅, 데이터 로딩, 로깅을 처리합니다.

엔드 투 엔드 파인튜닝 실행을 시작하기 위해 레시피는 CLI를 제공합니다.

torchrun --nproc_per_node=1 evo2_classifier.py \
    --train-jsonl splice_train.jsonl \
    --val-jsonl   splice_val.jsonl \
    --test-jsonl  splice_test.jsonl \
    --base-ckpt-dir evo2_1b_bf16_mbridge \
    --result-dir   splice_run \
    --experiment-name lora_finetune \
    --num-classes 3 \
    --seq-length-tokens 600 \
    --train-iters 1000 \
    --global-batch-size 32 --micro-batch-size 32 \
    --lr 5e-4 --min-lr 5e-5 --warmup-iters 30 \
    --lora-finetune --lora-dim 16 --lora-alpha 32 --lora-dropout 0.1

--lora-finetune을 제거하고 배치 크기를 늘리면 헤드만 학습하는 베이스라인을 재현할 수 있습니다. 데이터, 옵티마이저, 스케줄러, 평가는 동일하게 유지됩니다.

전체 학습 루프, 데이터셋 코드, 파라미터 계산, 평가 유틸리티는 Evo2 LoRA 파인튜닝 노트북을 참고하세요.

생물학적 파운데이션 모델 파인튜닝 시작하기

단백질(ESM2)과 DNA(Evo2)라는 전혀 다른 두 가지 생물학적 모달리티에서 동일한 파라미터 효율적 레시피를 사용할 수 있습니다. 사전 학습된 백본을 고정하고, 소규모 LoRA 어댑터와 태스크별 헤드를 학습하면, 단일 워크스테이션 GPU에서 전체 파인튜닝이나 특화 모델에 버금가는 정확도를 달성할 수 있습니다.

ESM2-3B의 경우, LoRA는 PSSP 성능을 Porter 6, SPOT-1D-LM과 같은 강력한 공개 베이스라인 수준으로 끌어올리며, TE와 THD 시퀀스 패킹을 통해 단일 NVIDIA RTX 6000 Blackwell Workstation Edition GPU에서의 학습을 실용적으로 만들어 줍니다. Evo2-1B의 경우, 동일한 접근 방식이 파라미터의 약 1.4%만 학습하면서 스플라이스 사이트 분류 정확도를 고정 백본 베이스라인의 약 52%에서 약 97%로 높입니다.

수십억 파라미터 규모의 생물학적 파운데이션 모델을 적절한 학습 스택(TE, Megatron-Bridge, 패킹된 시퀀스, PEFT)이 잘 통합된다면 일반 하드웨어에서도 이제 적응시킬 수 있습니다. NVIDIA BioNeMo Recipes는 이러한 통합을 예외가 아닌 기본값으로 만들기 위해 설계되었습니다.

LoRA, TE, 확장 가능한 PyTorch 워크플로로 생물학적 파운데이션 모델 파인튜닝을 시작하려면 NVIDIA BioNeMo Recipes를 확인하세요.