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.