adaptyvbio/protein-design-skills/mosaic

Mosaic Multi-Objective Design

Mosaic (Escalante Bio) is a JAX framework for "functional, multi-objective protein design using continuous relaxation." It optimizes a soft sequence by gradient descent over a continuous relaxation of sequence space, in the lineage of ColabDesign, RSO, and BindCraft, with one key difference: it composes multiple learned objectives from different models in a single differentiable loss.

When Mosaic fits

Mosaic is a framework for custom objectives, not a one-click method. The README is explicit: it "may require substantial hand-holding (tuning learning rates, etc), often produces proteins that fail simple in-silico tests, [and] should be combined with standard filtering methods." Reach for it when a fixed pipeline cannot express the objective you need. For a turnkey binder run, use bindcraft instead.

Prerequisites

JIT compilation makes the first call to any loss slow; later calls are fast.

Install

Mosaic runs locally on a JAX GPU or TPU build. It has no CLI and no Modal integration; you drive it through the marimo notebooks or the Python API.

git clone https://github.com/escalante-bio/mosaic && cd mosaic
uv sync --group jax-cuda      # or --group jax-tpu / --group jax-cpu
uv add jax[cuda12]            # may be needed for a GPU build
uv run marimo edit examples/example_notebook.py

Ready-made examples include esmfold_minibinder.py, esmfold_vhh.py, boltzgen_pipeline.py, and batched_protenix.py.

Core idea

A design objective is built from LossTerm objects that you add and scale with plain Python arithmetic, then hand to an optimizer.

import mosaic.losses.structure_prediction as sp
# Compose a loss from interface, confidence, and inverse-folding terms
design_loss = (
    sp.BinderTargetContact()
    + sp.WithinBinderContact()
    + 0.05 * sp.TargetBinderPAE()
    + 0.05 * sp.BinderTargetPAE()
    + 0.025 * sp.IPTMLoss()
    + 0.1 * sp.PLDDTLoss()
)

Loss terms can wrap one model used several ways (for example a structure predictor scoring both the binder-target complex and the binder as a monomer). Composing different architectures also lowers the chance of finding adversarial sequences that fool a single predictor.

What you can compose

Optimizers

A reasonable simplex_APGM step size is about 0.1 * sqrt(binder_length).

Worked example: ranking with ipSAE

The published Nipah competition recipe optimizes a design loss on Boltz-2, then ranks candidates with a separate multi-sample loss built from ipTM and ipSAE. The multi-sample loss is a method on the Boltz2 model, not a free function:

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,
)

On the Adaptyv Nipah de novo target, this recipe produced 8 binders out of 9 tested designs at nanomolar affinity, the highest hit-rate of any method on that target in the public results. That is a small, expert-tuned sample on one hard target, not a guarantee across targets, so treat Mosaic as a high-ceiling option that rewards careful objective design rather than a turnkey default.

Two practices from that work are worth carrying over:

• Let the optimizer choose the epitope. Asking for a binder, without fixing

hotspots, can find a better interface than a manually chosen one.

• Match filter stringency to assay throughput. With high-throughput testing,

filter lightly to keep diversity rather than applying heavy consensus filters that can reject good binders.

Decision tree

Should I use Mosaic?
│
├─ Need a custom objective across multiple models? → Mosaic
├─ Want one-click binders with default filters?    → bindcraft
├─ Want all-atom diffusion design?                  → boltzgen
└─ Want backbone-only diversity?                    → rfdiffusion + proteinmpnn

Cost

Adaptyv's own tests of these models showed Mosaic costing about $0.55 per accepted design, averaged across 7 targets, among the cheapest per design of the methods tested. That is compute only; the setup and tuning effort is the real cost of using Mosaic.

Troubleshooting

Next: Validate designs with boltz or chai, rank with ipsae, then filter with protein-qc.