adaptyvbio/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

| Requirement | Minimum | Recommended | |-------------|---------|-------------| | Python | 3.11+ | 3.11 | | Framework | JAX with CUDA or TPU | JAX CUDA 12 | | GPU VRAM | 24GB | 48GB+ (depends on predictors used) |

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

| Category | Options | |----------|---------| | Structure predictors | AF2, Boltz-1, Boltz-2, Protenix, OpenFold3, ESMFold2 | | Generative / design | BoltzGen, Proteina-Complexa | | Inverse folding | ProteinMPNN, SolubleMPNN, AbMPNN | | Language models | ESM-2, ESM-C, AbLang, trigram | | Property heads | Stability (megascale-trained) |

Optimizers

| Optimizer | Use | |-----------|-----| | simplex_APGM | Default; proximal gradient / mirror descent on the probability simplex | | batched_simplex_APGM | The same, vmapped over many designs | | gradient_MCMC | Discrete moves for fine-tuning a sequence |

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

| Issue | Cause | Fix | |-------|-------|-----| | Designs fail simple in-silico checks | Under-constrained objective | Add inverse-folding and confidence terms; filter with protein-qc | | Optimization unstable | Step size too large | Lower the simplex_APGM step size | | First call very slow | JIT compilation | Expected; reuse the compiled loss across designs | | OOM with large predictors | Several models in one loss | Use smaller predictors or a larger GPU |

---

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