mosaic

用途

Mosaic（Escalante Bio）是一个基于 JAX 的框架，用于“使用连续松弛进行功能性、多目标蛋白质设计”。它通过对序列空间连续松弛的软序列进行梯度下降来优化，属于 ColabDesign、RSO 和 BindCraft 一脉，但有一个关键区别：它能在单个可微损失中组合来自不同模型的多个学习目标。

何时适合使用 Mosaic

Mosaic 是用于自定义目标的框架，而不是一键式方法。README 明确指出：它“可能需要大量手动调参（调整学习率等），经常产生无法通过简单 in-silico 测试的蛋白质，[并且] 应与标准过滤方法结合使用。”当固定流程无法表达你所需的目标时，才应使用它。对于开箱即用的 binder 运行，请改用 bindcraft。

前置条件

要求	最低	推荐
Python	3.11+	3.11
框架	带 CUDA 或 TPU 的 JAX	JAX CUDA 12
GPU 显存	24GB	48GB+（取决于所用预测器）

JIT 编译会使任何损失的首次调用变慢；后续调用很快。

安装

Mosaic 在本地 JAX GPU 或 TPU 构建上运行。它没有 CLI，也没有 Modal 集成；你通过 marimo notebook 或 Python API 驱动它。

git clone https://github.com/escalante-bio/mosaic && cd mosaic
uv sync --group jax-cuda      # 或 --group jax-tpu / --group jax-cpu
uv add jax[cuda12]            # GPU 构建可能需要
uv run marimo edit examples/example_notebook.py

现成示例包括 esmfold_minibinder.py、esmfold_vhh.py、boltzgen_pipeline.py 和 batched_protenix.py。

核心思想

设计目标由 LossTerm 对象构建，你可以用普通 Python 算术添加和缩放它们，然后交给优化器。

import mosaic.losses.structure_prediction as sp

# 从界面、置信度和逆折叠项组合损失
design_loss = (
    sp.BinderTargetContact()
    + sp.WithinBinderContact()
    + 0.05 * sp.TargetBinderPAE()
    + 0.05 * sp.BinderTargetPAE()
    + 0.025 * sp.IPTMLoss()
    + 0.1 * sp.PLDDTLoss()
)

损失项可以包装一个模型并以多种方式使用（例如结构预测器同时对 binder-target 复合物和作为单体的 binder 打分）。组合不同架构还能降低找到欺骗单一预测器的对抗序列的可能性。

可以组合的内容

类别	选项
结构预测器	AF2、Boltz-1、Boltz-2、Protenix、OpenFold3、ESMFold2
生成 / 设计	BoltzGen、Proteina-Complexa
逆折叠	ProteinMPNN、SolubleMPNN、AbMPNN
语言模型	ESM-2、ESM-C、AbLang、trigram
属性头	稳定性（megascale 训练）

优化器

优化器	用途
simplex_APGM	默认；概率单纯形上的近端梯度 / 镜像下降
batched_simplex_APGM	相同，但对许多设计进行 vmapped
gradient_MCMC	用于微调序列的离散移动

合理的 simplex_APGM 步长约为 0.1 * sqrt(binder_length)。

实例：使用 ipSAE 排序

已发表的 Nipah 竞赛配方在 Boltz-2 上优化设计损失，然后用由 ipTM 和 ipSAE 构建的独立多样本损失对候选进行排序。多样本损失是 Boltz2 模型上的方法，而不是自由函数：

from mosaic.models.boltz2 import Boltz2

boltz2 = Boltz2()
ranking_loss = boltz2.build_multisample_loss(
    loss=1.00 * sp.IPTMLoss()
    + 0.5 * sp.TargetBinderIPSAE()
    + 0.5 * sp.BinderTargetIPSAE(),
    features=design_features,
    num_samples=6,
    recycling_steps=3,
)

在 Adaptyv Nipah de novo 靶点上，该配方在 9 个测试设计中产生 8 个 binder，亲和力为纳摩尔级，是公开结果中该靶点上任何方法中最高的命中率。这是一个在单一困难靶点上的小型、专家调优样本，并非跨靶点的保证，因此应将 Mosaic 视为高上限选项，奖励仔细的目标设计，而不是开箱即用的默认选择。

该工作中有两个值得借鉴的实践：

让优化器选择表位。要求一个 binder 而不固定热点，可能找到比手动选择更好的界面。
将过滤严格度与测定通量匹配。在高通量测试中，轻度过滤以保持多样性，而不是应用可能拒绝好 binder 的严格共识过滤器。

决策树

我应该使用 Mosaic 吗？
│
├─ 需要跨多个模型的自定义目标？ → Mosaic
├─ 想要带默认过滤器的一键 binder？    → bindcraft
├─ 想要全原子扩散设计？                  → boltzgen
└─ 只想要骨架多样性？                    → rfdiffusion + proteinmpnn

成本

Adaptyv 对这些模型自己的测试显示，Mosaic 每个被接受设计的成本约为 0.55 美元，在 7 个靶点上平均，是所测试方法中每个设计成本最低的之一。这只是计算成本；设置和调优工作才是使用 Mosaic 的真正成本。

故障排除

问题	原因	修复
设计未通过简单 in-silico 检查	目标约束不足	添加逆折叠和置信度项；用 protein-qc 过滤
优化不稳定	步长过大	降低 simplex_APGM 步长
首次调用非常慢	JIT 编译	预期行为；跨设计重用编译后的损失
大型预测器 OOM	一个损失中有多个模型	使用更小的预测器或更大的 GPU

下一步：用 boltz 或 chai 验证设计，用 ipsae 排序，然后用 protein-qc 过滤。