github.com/escalante-bio/mosaic @main

Functional, multi-objective protein design using continuous relaxation.

WARNING: Unlike BindCraft (which is a well-tested and well-tuned method for generic binder design), mosaic may require substantial hand-holding (tuning learning rates, etc), often produces proteins that fail simple in-silico tests, should be combined with standard filtering methods, etc. This is not for the faint of heart: the intent is to provide a framework in which to implement custom objective functions and optimization algorithms for your application. You can read about some applications in our blog.

WARNING: We rely heavily on just-in-time compilation via JAX; the first call to a JAX-compiled function will be slow. After that things should be pretty fast. If you're tuning loss weights or optimizer parameters, you should use an interactive session or a notebook!

Why?

mosaic is an attempt to reimplement a zoo of protein property models with a common interface to make it easier to run them together without dealing with containers, horrible dependencies, etc. Because they're all implemented using the same backend (JAX), they can be efficiently and conveniently connected with code rather than bash scripts. We use this for protein binder design, but there are many applications.

Protein design tasks almost always involve multiple constraints or properties that must be satisfied or optimized. For instance, in binder design one may want to simultaneously ensure: - the chance of binding the intended target is high
- the chance of binding to a similar off-target protein is low - the binder expresses well in bacteria - the binder is highly soluble.

There has been a recent explosion in the application of machine learning to protein property prediction, resulting in fairly accurate predictors for each of these properties. What is currently lacking is an efficient and flexible method for combining these different predictors into one design/filtering/ranking framework.

Models

Applications & case studies

• Minibinder design is just not that hard — end-to-end walkthrough with experimentally validated binders against PD-L1 (8/10) and IL7Ra (7/10). Optimizing a complex objective leads to higher hitrates and affinities.
• ~180 lines of code to win the in silico portion of the Adaptyv Nipah binding competition + Winning the de novo portion of the Adaptyv Nipah binder competition — a minimal, self-contained design script (Boltz-2 + soluble ProteinMPNN + simplex_APGM) wins a design competition against a novel target with single digit nanomolar binders and a 90% wet-lab success rate.
• Teaching generative models to hallucinate — using a mosaic loss functional as the reward to finetune and RL-align a generative model (BoltzGen).
• Mosaic multispecifics — mosaic-designed designing bispecific IL7Ra and PD-L1 minibinders
• Adaptyv × Muni "AI agents vs humans" (TREM2) — a mosaic design produced the top-affinity binder overall (1.11 nM), ahead of six autonomous LLM design agents (results on Proteinbase). Note: Claude is very good at writing mosaic scripts and HPO.
• GEM × Adaptyv RBX1 binder design competition — won by ORBIT, a hallucination-based method built on mosaic.

Citing mosaic

If you like mosaic or build on it please cite us:

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

Installation

We recommend using uv, e.g. run uv sync --group jax-cuda after cloning the repo to install dependencies.

To run the example notebook try uv run marimo edit examples/example_notebook.py.

You may need to add various uv overrides for specific packages and your machine, take a look at pyproject.toml

You'll need a GPU or TPU-compatible version of JAX for structure prediction. You might need to install this manually, i.e. uv add jax[cuda12].

Introduction

This project combines two simple components to make a powerful protein design framework:

• A functional, modular interface to easily combine multiple learned or hand-crafted loss terms and optimization algorithms (as in A high-level programming language for generative protein design etc)
• Gradient-based optimization over a continuous, relaxed sequence space (as in ColabDesign, RSO, BindCraft, etc)

The key observation is that it's possible to use this continuous relaxation simultaneously with multiple learned objective terms [^1].

This allows us to easily construct objective functions that are combinations of multiple learned potentials and optimize them efficiently, like so:

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

boltz1 = Boltz1()
mpnn = ProteinMPNN.from_pretrained()

target_sequence = "DYSFSCYSQLEVNGSQHSLTCAFE..."
binder_length = 80

# Generate features for binder-target complex
boltz_features, _ = boltz1.binder_features(
    binder_length=binder_length,
    chains=[TargetChain(sequence=target_sequence)],
)

# Generate features for binder alone (monomer)
mono_features, _ = boltz1.binder_features(
    binder_length=binder_length,
    chains=[]
)

combined_loss = (
    boltz1.build_loss(
        loss=4 * sp.BinderTargetContact()
        + sp.DistogramRadiusOfGyration(target_radius=15.0)
        + sp.WithinBinderContact()
        + 0.3 * sp.HelixLoss()
        + 5.0 * InverseFoldingSequenceRecovery(mpnn, temp=jax.numpy.array(0.01)),
        features=boltz_features,
        recycling_steps=1,
    )
    + 0.5 * esm_loss
    + trigram_ll
    + 0.1 * stability_loss
    + 0.5
    * boltz1.build_loss(
        loss=0.2 * sp.PLDDTLoss()
        + sp.DistogramRadiusOfGyration(target_radius=15.0)
        + 0.3 * sp.HelixLoss(),
        features=mono_features,
        recycling_steps=1,
    )
)

_, PSSM = simplex_APGM(
    loss_function=combined_loss,
    n_steps=150,
    x=jax.nn.softmax(
        0.5 * jax.random.gumbel(
            key=jax.random.key(np.random.randint(100000)),
            shape=(binder_length, 20),
        )
    ),
    stepsize=0.1,
    momentum=0.9,
)

Here we're using ~5 different models to construct a loss function: the Boltz-1 structure prediction model (which is used twice: once to predict the binder-target complex and once to predict the binder as a monomer), ESM2, ProteinMPNN, an n-gram model, and a stability model trained on the mega-scale dataset.

It's super easy to define additional loss terms, which are JIT-compatible callable pytrees, e.g.

class LogPCysteine(LossTerm):
    def __call__(self, soft_sequence: Float[Array, "N 20"], key = None):
        mean_log_p = jnp.log(soft_sequence[:, IDX_CYS] + 1E-8).mean()
        return mean_log_p, {"log_p_cys": mean_log_p}

Though a better way to exclude cysteine is to wrap an existing loss, as in NoCys.

WARNING: Optimization is hard: it's quite easy to create an objective function that's difficult to minimize using our standard optimizers. For many design problems you may have to come up with your own heuristic optimization algorithms (for instance by guiding generative models or combining multiple designs).

There's no reason custom loss terms can't involve more expensive (differentiable) operations, e.g. an EVOLVEpro-style fitness predictor.

The marimo notebooks give a few examples of how this can work.

It's very easy to swap in different optimizers. For instance, let's say we really wanted to try projected gradient descent on the hypercube $[0,1]^N$. We can implement that in a few lines of code:

from mosaic.optimizers import _print_iter, _eval_loss_and_grad
def RSO_box(
    *,
    loss_function,
    x: Float[Array, "N 20"],
    n_steps: int,
    stepsize: float,
    max_grad_norm: float,
    key=None,
):
    if key is None:
        key = jax.random.PRNGKey(np.random.randint(0, 10000))

    for _iter in range(n_steps):
        (v, aux), g = _eval_loss_and_grad(
            x=x,
            loss_function=loss_function,
            key=key
        )
        g_norm = np.linalg.norm(g)
        if g_norm > max_grad_norm:
            g = g * (max_grad_norm / g_norm)
        x = (x - stepsize * g).clip(0,1)
        key = jax.random.fold_in(key, 0)
        _print_iter(_iter, aux, v)

    return x

Take a look at optimizers.py for examples.

Structure Prediction

We provide a simple interface in mosaic.structure_prediction and mosaic.models.* to nine structure prediction models: OpenFold3, Boltz1, Boltz2, AF2, ProtenixMini, ProtenixTiny, ProtenixBase, Protenix2025, and ESMFold2.

To make a prediction or design a binder, you'll need to make a list of mosaic.structure_prediction.TargetChain objects. This is a simple dataclass that contains a protein (or DNA or RNA) sequence, a flag to tell the model if it should use MSAs (use_msa), and potentially a template structure (as a gemmi.Chain).

For example, we can make a prediction with Protenix for IL7Ra like so:

```python

import jax from mosaic.structure_prediction import TargetChain from mosaic.models.protenix import Protenix2025

model = Protenix2025()

target_sequence = "DYSFSCYSQLEVNGSQHSLTCAFEDPDVNTTNLEFEICGALVEVKCLNFRKLQEIYFIETKKFLLIGKSNICVKVGEKSLTCKKIDLTTIVKPEAPFDLSVVYREGANDFVVTFNTSHLQKKYVKVLMHDVAYRQEKDENKWTHVNLSSTKLTLLQRKLQPAAMYEIKVRSIPDHYFKGFWSE