escalante-bio/mosaic

mosaic – 다목적 단백질 설계를 위한 JAX 기반 프레임워크

무엇인가요

• github.com/escalante-bio/mosaic 에 호스팅된 오픈소스 Python 라이브러리로, 연구자들이 다양한 미분 가능한 단백질 특성 예측기와 구조 예측 모델을 하나의 그래디언트 기반 최적화 루프에 통합할 수 있게 해줍니다.
• JAX 기반으로 구축되어 있어 모든 모델이 공통의 JIT 컴파일 백엔드를 공유하며, 무거운 컨테이너 오케스트레이션 없이 체인 가능합니다.

핵심 아이디어

1. 통합된 손실 인터페이스 – 각 예측기 (예: Boltz‑1, AlphaFold‑Multimer, ESM‑2, ProteinMPNN, 안정성 모델 등)는 JAX 호환 손실 항목을 반환하는 build_loss 메서드를 제공합니다.
2. 연속적인 시퀀스 리라크세이션 – 디자인은 20개의 아미노산에 대한 소프트맥스 확률 행렬로 표현되어 시퀀스 공간에서 그래디언트 하강이 가능합니다.
3. 플러그인 가능한 최적화기 – 라이브러리는 간단한 프로젝션 그래디언트 최적화기 (simplex_APGM)를 제공하며, 박스 제약 RSO와 같은 사용자 정의 최적화기 예제도 포함되어 있습니다.

포함된 내용

• 모델: Boltz‑1/2, BoltzGen, AlphaFold‑2/Multimer, OpenFold‑3, OpenDDE, ESMFold‑2, Protenix (mini/tiny/base/2025), ProteinMPNN (standard/soluble/AbMPNN), ESM‑2/C, 안정성 예측기, AbLang/AbLang2, 트리그램 언어 모델, 몇 가지 데모 모델.
• 손실 라이브러리: 구조 예측 손실 (바인더-타겟 접촉, 반지름의 관성 모멘트, PLDDT 등), ProteinMPNN를 사용한 역폴딩 손실, 안정성, n-그램 가능성, 사용자 정의 손실.
• 예제 및 노트북: 미니바인더 설계를 위한 엔드투엔드 노트북, 경쟁 우승 스크립트, 새로운 손실 항목 추가 튜토리얼.
• 유틸리티: 시퀀스, MSA, 선택적 템플릿을 지정하기 위한 dataclass TargetChain; 특징 추출, 모델 독립 예측, PSSM 처리를 위한 헬퍼 함수.

일반적인 워크플로우

from mosaic.models.boltz1 import Boltz1
from mosaic.proteinmpnn.mpnn import ProteinMPNN
from mosaic.losses import structure_prediction as sp
from mosaic.optimizers import simplex_APGM
import jax, numpy as np

# 모델 로드
boltz = Boltz1()
mpnn  = ProteinMPNN.from_pretrained()

# 타겟과 바인더 길이 정의
target_seq = "..."
binder_len = 80

# 복합체와 모노머의 특징 추출
cplx_feat, _ = boltz.binder_features(binder_len, [TargetChain(sequence=target_seq)])
mono_feat, _ = boltz.binder_features(binder_len, [])

# 복합 손실 정의 (결합성, 기하학, 용해성, 역폴딩 등)
loss = (
    boltz.build_loss(sp.BinderTargetContact() + sp.DistogramRadiusOfGyration(15.0), cplx_feat)
    + 0.3 * sp.HelixLoss()
    + 5.0 * InverseFoldingSequenceRecovery(mpnn, temp=jax.numpy.array(0.01))
    + 0.5 * sp.StabilityLoss()
)

# 소프트맥스 시퀀스 행렬 초기화
x0 = jax.nn.softmax(
    0.5 * jax.random.gumbel(jax.random.key(0), (binder_len, 20))
)

# 그래디언트 기반 설계 실행
_, pssm = simplex_APGM(loss_function=loss, x=x0, n_steps=150, stepsize=0.1, momentum=0.9)

결과로 얻은 pssm 은 샘플링하여 실험 검증용 구체적인 아미노산 시퀀스를 얻을 수 있습니다.

왜 중요한가요

• 단백질 공학은 종종 결합성, 특이성, 발현성, 용해성, 안정성 등 여러 경쟁적인 목표를 균형 잡는 것이 필요합니다. mosaic 는 이를 모두 미분 가능한 항으로 표현하고 함께 최적화할 수 있게 해주며, 별도의 스크립트나 블랙박스 파이프라인에서는 어렵습니다.
• JAX에서 모든 것이 실행되므로, CPU, GPU, TPU에서 거의 변경 없이 동일한 코드가 작동하며, JIT 컴파일로 첫 호출 이후 내부 루프가 매우 빠릅니다.

제한 사항 / 주의사항

• 이 라이브러리는 연구용 프레임워크이며, 즉시 사용 가능한 설계 도구는 아닙니다. 사용자는 학습률, 최적화기 하이퍼파라미터, 손실 가중치를 조정해야 하며, 잘못된 설정은 시뮬레이션에서 기본 검사를 통과하지 못하는 설계로 이어질 수 있습니다.
• 처음 JAX 컴파일은 느립니다. 빠른 반복을 위해서는 인터랙티브 노트북을 추천합니다.
• 일부 구조 예측 모델 (예: AF3 스타일)은 설계 시 측쇄 정보를 생략할 수 있어, 후속 필터링에 영향을 줄 수 있습니다.
• 무거운 모델 (AlphaFold, OpenFold, Protenix 등)을 사용하려면 GPU/TPU 호환 JAX 빌드가 필요합니다.

시작하기

1. 리포지토리 복제 후 uv (권장) 또는 pip 로 설치:

   git clone https://github.com/escalante-bio/mosaic.git
   cd mosaic
   uv sync --group jax-cuda   # 또는 수동으로 jax[cuda12] 추가
   

2. 예제 노트북 실행:

   uv run marimo edit examples/example_notebook.py
   

3. README에 링크된 미니바인더 설계 블로그 포스트를 따라가며 PD‑L1 및 IL7Ra에 대한 실험적으로 검증된 바인더를 재현하세요.

인용 mosaic 을 발표된 연구에 사용할 경우 다음을 인용하세요:

Boyd, N., Guns, S. & Escalante Bio. Mosaic. https://github.com/escalante-bio/mosaic (2025).