Model

Build and Train a Deep Learning Model

Model:

To turn the "mutation → simulate → select" loop into a knowledge-preserving generator, we trained a multimodal latent autoencoder on RFdiffusion+ProteinMPNN variants labeled in silico by Brownian dynamics (BD). The encoder has two synchronized inputs:

• Sequence stream: Amino-acid tokens are embedded with ESM-2 (650M) and passed through a lightweight Transformer encoder to produce a context-aware sequence embedding.
• Structure stream: The same variant is represented as a residue-level contact graph; node features encode residue identity/secondary-structure hints and edge features are RBFs of inter-residue distances. A GINEConv stack with Set2Set readout yields a structural embedding. During pretraining, we add three self-supervised signals on the graph branch—contrastive learning with a memory bank of hard negatives, prototype consistency, and masked-node recovery—so the structural representation becomes geometry-aware yet robust to small contact perturbations.

Both streams are L2-normalized and fused into a single latent vector z. A small performance head maps z to the normalized kon, shaping the manifold around what the BD labeler deems favorable. A compact Transformer decoder then reconstructs only the mutated positions (edit-only reconstruction): we supply a mutation mask derived from the RFdiffusion/MPNN delta to (i) copy untouched residues straight through via a skip connection, and (ii) apply the reconstruction loss exclusively on mutated sites. Non-mutated tokens incur no loss and are passed through unchanged. This design has three payoffs crucial for Seed&Seek:

• Protects conserved motifs. Catalytic/structural residues are preserved by construction, preventing the decoder from "improving" what must not change.
• Focuses capacity on edits. The latent is forced to encode how and why successful edits work (electrostatic tuning, loop stiffening), not to memorize the entire sequence.
• Enables constrained generation. At proposal time we keep the mask fixed (or tighten it), letting the model suggest targeted edits while guaranteeing global fold-stabilizing regions remain intact.

We tie input/output embeddings for stability, keep the decoder shallow to avoid rote memorization, and regularize with label smoothing, token/edge dropout, and occasional encoder-drop (decode from z alone). The result is a fitness-aware, geometry-aware latent that captures the semantics of beneficial edits and can be searched directly in the recursive seek phase.

Methods:

# ====== Data ======
# batch = (seq_tokens, structure, mut_mask, y_target)  # y_target = normalized k_on (BD)

# ====== Modules ======
# --- Sequence side (ESM + adapters + small encoder) ---
ESM = ESM2_650M()                         # frozen weights by default
ESM_Adapters = LoRA_Adapters(ESM, rank=r) # trainable; attach to last L layers
SeqEncoder   = TransformerEncoder_small() # light encoder on top of ESM embeddings

# --- Graph side (residue-level contact graph) ---
GraphEncoder = GINE_Stack_With_Set2Set(   # GINEConv x L, Set2Set readout
    node_in_dim=20 + ss_hints, rbf_k=32, hidden=256, set2set_steps=3
)

# --- Fusion / heads / decoder ---
Fuse       = LinearFuse()                 # concat → linear → L2
PerfAdapter= MLP()                        # Z → Z_v (performance-adapted)
PerfHead   = Linear()                     # predicts performance scalar (normalized k_on)
Decoder    = TransformerDecoder_tiedIO()  # compact; edit-only reconstruction

# --- Graph SSL heads (for pretraining the GNN) ---
ContrastHead = MoCoQueue(temp=0.07, queue=2048)   # hard negatives
ProtoHead    = PrototypeConsistency(K=100, temp=0.05)
MaskNodeHead = LinearClassifier(20)               # AA category at masked nodes

# ====== Utilities ======
function BuildResidueGraph(structure):
    nodes = residues(structure)
    edges = kNN_or_distance_edges(nodes)          # e.g., k=16 or d<8Å
    E_attr = RBF(distance(edges), k=32, dmin=0, dmax=10)
    X_node = onehotAA(nodes) + secStructHint(nodes)
    return Graph(X_node, edges, E_attr)

function EditOnlyCE(logits, targets, mut_mask):
    idx = [i | mut_mask[i] == 1]
    return CrossEntropy(logits[idx], targets[idx])

L2(x) = x / max(ε, ||x||)

# ====== Forward (joint AE) ======
function FORWARD(seq_tokens, graph, mut_mask):
    # SEQUENCE: ESM embeddings (with optional adapter fine-tuning)
    esm_tok, esm_pooled = ESM.embed(seq_tokens)        # gradients allowed only through adapters
    h_seq  = SeqEncoder(esm_tok)                       # small Transformer on top
    Z_seq  = L2(pool(h_seq))

    # GRAPH: GINE + Set2Set
    Z_graph = L2( GraphEncoder(graph) )                # Set2Set readout → projection

    # FUSE → Z → Z_v → Z_latent
    Z       = L2( Fuse(CONCAT(Z_seq, Z_graph)) )
    Z_v     = L2( PerfAdapter(Z) )
    Z_lat   = L2( CONCAT(Z, Z_v) )

    # HEADS
    y_hat   = PerfHead(Z_lat)
    logits  = Decoder(Z_lat, context=seq_tokens)       # logits per position

    # EDIT-ONLY reconstruction
    rec_seq = seq_tokens
    for i in range(len(seq_tokens)):
        if mut_mask[i] == 1:
            rec_seq[i] = ARGMAX(logits[i])

    return Z_seq, Z_graph, Z, Z_v, Z_lat, y_hat, logits, rec_seq

# ====== Sequence-side adaptation (ESM + SeqEncoder) ======
# Masked-LM on mutated or high-entropy positions; keeps ESM stable via adapters.
function SEQ_PRETRAIN_STEP(seq_tokens, mut_mask):
    mask_idx = sample_positions(mut_mask, rate=0.15)          # prefer mutated sites
    seq_in, seq_tgt = apply_MLM_mask(seq_tokens, mask_idx)

    esm_tok, _ = ESM.embed(seq_in)                            # adapters get gradients
    h_seq  = SeqEncoder(esm_tok)
    logits = LM_head(h_seq)                                   # tied with decoder embeddings

    L_mlm  = CrossEntropy(logits[mask_idx], seq_tgt[mask_idx])
    OPT_SEQ.step(L_mlm)                                       # optimize ESM_Adapters + SeqEncoder

# ====== Graph-side pretraining (GINE + Set2Set) ======
function GRAPH_SSL_STEP(graph_batch):
    g1 = augment_graph(graph_batch)                           # edge-drop, rbf noise, node-drop
    g2 = augment_graph(graph_batch)

    z1 = L2(GraphEncoder(g1))
    z2 = L2(GraphEncoder(g2))

    L_contrast         = ContrastHead(z1, z2)                 # InfoNCE + memory bank
    L_proto, L_entreg  = ProtoHead(z1, z2)                    # prototype consistency + entropy reg
    L_masknode         = CE(MaskNodeHead(mask_nodes(g1)), true_labels_of_masked_nodes(g1))

    L_ssl = L_contrast + λp*L_proto + λe*L_entreg + λm*L_masknode
    OPT_GNN.step(L_ssl)                                       # optimize GraphEncoder + SSL heads

# ====== Joint AE training step (supervised by BD labels) ======
function JOINT_AE_STEP(seq_tokens, structure, mut_mask, y_target):
    graph = BuildResidueGraph(structure)
    _, _, Z, Z_v, Z_lat, y_hat, logits, _ = FORWARD(seq_tokens, graph, mut_mask)

    L_pred = MSE(y_hat, y_target)                             # performance shaping (k_on)
    L_rec  = EditOnlyCE(logits, seq_tokens, mut_mask)         # only mutated positions
    L_total = L_pred + α*L_rec + REG()                        # label smoothing, token/edge dropout

    OPT_AE.step(L_total)                                      # optimize SeqEncoder, GraphEncoder,
                                                              # Fuse, PerfAdapter, PerfHead, Decoder
# NOTE: when fine-tuning, freeze ESM backbone; unfreeze only LoRA adapters.

# ====== Optimizers / parameter groups ======
OPT_SEQ = Adam(params={ESM_Adapters, SeqEncoder, LM_head},        lr=1e-4, wd=1e-5)
OPT_GNN = Adam(params={GraphEncoder, ContrastHead, ProtoHead,
                       MaskNodeHead},                              lr=1e-3, wd=1e-5)
OPT_AE  = Adam(params={SeqEncoder, ESM_Adapters, GraphEncoder,    # small LR for adapters
                       Fuse, PerfAdapter, PerfHead, Decoder},      lr=5e-4, wd=1e-5)

# ====== Training loop ======
for epoch in 1..EPOCHS:
    for batch in DataLoader(TRAIN, B):
        seq_tokens, structure, mut_mask, y = batch

        # 1) sequence-side MLM adaptation (few steps per epoch)
        if rand() < p_seq_ssl:
            SEQ_PRETRAIN_STEP(seq_tokens, mut_mask)

        # 2) graph-side SSL (few steps per epoch)
        if rand() < p_graph_ssl:
            GRAPH_SSL_STEP(BuildResidueGraph(structure))

        # 3) supervised joint AE step (always)
        JOINT_AE_STEP(seq_tokens, structure, mut_mask, y)

    # validation (token accuracy on mutated sites + R^2 on k_on)
    val_tok_acc = evaluate_edit_only_accuracy(VALID)
    val_R2      = evaluate_R2(VALID)
    early_stop_if_needed(val_tok_acc, val_R2)
                                    

Data regime. Our training signals come from synthetic variants plus noisy BD labels, so we lean on large, frozen priors (ESM) and a small trainable head to avoid overfitting. The graph stream encodes what sequence cannot tell us—spatial neighborhood chemistry, orientation, and long-range couplings that are local in 3D—so mutations that preserve the fold but shift approach-face electrostatics can be separated in latent space.

Noise and regularization. BD labels have counting noise (react vs escape paths), so we use label-smoothing for the performance head, stochastic token/edge drop, and encoder-drop (occasionally decoded from z only) to prevent the model from attributing noise to specific residues. Temperatures in the contrastive/prototype heads are clamped to keep gradients well-scaled; a small EMA of weights stabilizes training.

Active-learning readiness. We deliberately keep z low-dimensional and normalized. In later "seeking" rounds, an acquisition function can move z toward better regions under compute budgets, then the decoder proposes concrete sequences—exactly the operation our mutation-only experiment lacked.

Results:

Reconstruction and prediction trade-off.

Training on RFdiffusion+MPNN variants labeled by BD produced a smooth, fitness-aware latent. Quantitatively, increasing the latent dimensionality from 96 → 128 raises token reconstruction accuracy from 0.942 → 0.987, confirming that the decoder can nearly perfectly reproduce mutants. However, the predictive alignment between latent and kinetics shows the opposite trend: the performance head's R² on held-out kon drops from 0.855 → 0.793. This is the classic capacity trade-off—larger latents allocate more degrees of freedom to memorize sequence detail, while slightly weakening the pressure to organize z along fitness-relevant directions. For downstream seeking, we therefore favor d=96 as the default (better R² while maintaining high reconstruction), and reserve d=128 when near-lossless reconstruction is essential (e.g., strict motif preservation).

Latent–fitness geometry.

Projections of z by PCA and UMAP retain clear fitness gradients: warmer points (higher normalized kon) form contiguous neighborhoods, indicating that small moves in latent space correspond to meaningful, chemistry-preserving edits. The t-SNE fitness surface exhibits broad ascents rather than spiky ridges, and "weight-climbing" trajectories converge toward basins—both signals that the latent is navigable for active learning.

Interpretation: for Seed&Seek's recursive search, d=96 offers the best balance—accurate reconstruction (accuracy > 0.9 shows highly efficient sequence reconstruction) with stronger fitness shaping—while d=128 is a high-fidelity option when exact sequence recovery is prioritized.

Explore on Latent and Generate

Model:

At inference we search directly inside the autoencoder's latent space. You give the system a target performance (for example, a high normalized k_on). A small latent optimizer then moves to a point in the latent space where the model's performance head predicts that target. From that optimized point, we create a small cloud of nearby latent vectors to promote diversity. Finally, each latent vector is decoded by the Transformer decoder to produce a sequence candidate. Because decoding is driven only by the latent vector, there is no need to rebuild graphs or run encoders at this stage; it is fast and easy to batch. Two simple controls shape exploration: the regularization that keeps the latent near regions seen during training, and the "spread" rule that determines how widely we sample around the optimized point (either isotropic Gaussian noise or covariance‐aware noise estimated from nearest neighbors in the training latents).

Methods:

############################################################
# EXPLORE ON LATENT AND GENERATE — PSEUDOCODE
# Inputs:
#   - target_perf: desired normalized performance (e.g., k_on)
#   - N: number of sequences to generate
#   - cfg: knobs for optimization and diversity (below)
# Outputs:
#   - candidates: list of {sequence, pred_perf, latent}

PREP:
  model = load_trained_AE()            # exposes perf_head(z) and z-only decoder
  Z_train = load_training_latents()    # (N, D) latent database

function KNN_NEIGHBORS(z):
  d = pairwise_distance(z, Z_train)
  idx = argsort(d)[:knn_k]
  return Z_train[idx]                  # (k, D)

function KNN_MEAN(z):
  return mean(KNN_NEIGHBORS(z), axis=0, keepdims=True)

function OPTIMIZE_LATENT(target_perf):
  best = None
  repeat restarts times:
    z = normal_init(D, std=0.5)
    for t in 1..steps:
      y_hat = model.perf_head(z)
      loss = (y_hat - target_perf)^2 + l2_weight*norm2(z)
      z_knn = KNN_MEAN(z)
      loss += knn_weight * norm2(z - z_knn)
      z = z - lr * grad(loss, z)
    score = -abs(model.perf_head(z) - target_perf)
    if best is None or score > best.score:
      best = {z: copy(z), pred: model.perf_head(z), score: score}
  return best.z, best.pred

function SAMPLE_KNN_COV(z_opt, m):
  nbr = KNN_NEIGHBORS(z_opt)             # (k, D)
  X = nbr - mean(nbr, axis=0)
  V, S = top_svd(X)                      # low-rank basis and scales
  list = []
  for i in 1..m:
    eps = normal_r_like(S)               # draw in r-dim principal space
    jitter = (eps * S) @ V^T
    list.append(z_opt + spread_scale * jitter)
  return list

function DECODE_LATENTS(z_list):
  out = []
  for z in z_list:
    seq = decode_with_z(model, z,
                        temperature=decode_temperature,
                        top_p=decode_top_p,
                        no_repeat_ngram=no_repeat_ngram)
    pred = model.perf_head(z)
    out.append({sequence: seq, pred_perf: pred, latent: z})
  return out

# main
z_opt, pred_at_opt = OPTIMIZE_LATENT(target_perf)
z_pool = SAMPLE_KNN_COV(z_opt, num_candidates)
candidates = DECODE_LATENTS(z_pool)
save_PCA_background_and_candidates(Z_train, z_pool, candidates.pred_perf)
return candidates
                                    

We search the autoencoder's latent space for a point whose predicted performance matches a target while staying on the training manifold. Let f(z) be the performance head (normalized k_on). The loss we minimize is:

where z_knn(z) is the mean of the K nearest training latents to z. After optimization yields z_opt, we generate diversity by sampling around it using the local KNN covariance: compute the SVD of centered KNN latents to get principal directions v_r and scales σ_r, then get:

Each z_i is decoded autoregressively (temperature and top-p) to a sequence.

Results:

We apply the following defaults: target_k_norm=2.4, num=600, lambda_knn=0.05, per_candidate_latent=True, spread_mode="knn_cov", knn_k=128, spread_scale=1.5, temperature=1.0, top_p=0.9.

Observations from the PCA plots (KNN / spread = 2.5 vs 4.5).

Both panels place candidates along warm (high-performance) bands of the latent manifold—evidence that the KNN prior keeps samples on realistic directions.

• Spread = 2.5: candidates form a compact cloud around the optimized point; the median predicted score (white tag) is slightly above the surrounding background and variance is modest—good exploitation.
• Spread = 4.5: candidates occupy a wider region following manifold ridges; the median predicted score remains comparable, but the tail covers more diverse zones—better exploration with a reasonable quality floor because the local covariance steers away from empty space.

How to tune during active learning

• Early cycles (explore): choose a larger spread (e.g., 3.5–4.5) and keep knn_k moderately large (64–128). This yields diverse proposals without falling off-manifold. Keep temperature near 1.0 and top_p≈0.9 to let decoding express the latent diversity.
• Mid cycles (balance): reduce spread toward 2.0–3.0 while keeping lambda_knn=0.05 to maintain realism. If simulation budget is limited, drop num but increase restarts in optimization to refine z_opt.
• Late cycles (exploit): tighten spread to 1.0–2.0 and optionally lower temperature (0.8–0.9) for higher per-candidate quality. Keep per_candidate_latent=True so proposals are not identical.

Safety rails: if candidates begin to drift into cooler regions, increase lambda_knn or knn_k; if all candidates cluster too tightly, raise spread_scale or slightly relax lambda_knn.

Practical recipe. Start with the defaults above; visualize each batch on PCA colored by predicted performance. If the candidate cloud is too concentrated, raise spread_scale toward 4.0 - 4.5; if quality drops noticeably, nudge lambda_knn to 0.075 – 0.1 or increase knn_k to 160 to pull proposals back onto the warm manifold. This schedule gives a smooth transition from exploration to exploitation while keeping generation fast and well-behaved.

Experiment: AI-guided Directed Evolution

Design: Active training

The full loop (diagram) runs as follows. We start from the pre-trained Seed&Seek autoencoder and its latent optimizer. Each active learning round proceeds in five steps:

1. Seek in latent. We optimize a latent point toward a target performance (normalized kon), then sample a KNN-covariance cloud around the optimum to get 600 candidates.
2. Decode & filter. We decode each latent to a sequence (edit-aware decoder).
3. Structure → physics. New sequences are sent to AlphaFold for structural hypotheses and then to Brownian dynamics (Browndye) to label kon.
4. Curate for fine-tuning. To avoid forgetting and collapsing, we carry over the top 300 from all previous rounds (diversity-aware sampling) and merge them with 300 newly ranked candidates → 600 total for that round's training set.
5. Fine-tune the model. We fine-tune the autoencoder on this 600-sample set:

• The performance head regresses the new kon.
• Decoder is trained in edit-only reconstruction so conserved positions pass through untouched; this keeps scaffolds stable.
• Encoders receive a small learning rate with a KNN-prior and contrastive push to reshape the latent around new chemistry without drifting off-manifold.

Preventing "recursive collapse"

Two mechanisms make the loop robust:

• Latent repulsion during re-encode. When new sequences are encoded back to z, we add a repulsive penalty that discourages assigning them to the exact same mode occupied by the previous round's proposals. Concretely, we compute KNNs between "old" and "new" latents and add a margin loss that pushes the new latents away from the old centroid while keeping them on-manifold (via the KNN prior). This forces both the decoder and the performance head to adapt to genuinely new directions rather than over-fitting one pocket.
• Replay + diversity quota. We always replay the top-300 historical points and enforce a quota for diverse structural neighborhoods (based on AlphaFold contact maps). The fine-tune therefore sees both the frontier and the backbone, preventing drift and catastrophic forgetting.

Each new round re-runs this cycle, updating the model parameters and the latent optimizer's landscape.

Results: Seeking for optimized solution

Latent PCA

Left panel (colored by dataset). Pre-train points (gray) occupy the right-hand cluster. As we go through loop1→loop4 (brown→pink), the cloud shifts left and spreads, indicating the model is learning new latent directions rather than collapsing back to the original mode.

Right panel (colored by kon). Warmer colors (higher kon) concentrate in the shifted region, showing that the discovered latent territory corresponds to better association rates.

Kinetic distributions across rounds.

The density plot shows clear rightward shifts: pre-train ≈20k, loop1 ≈25–33k, loop2 ≈30–49k, loop3 ≈37–67k, loop4 ≈42–66k. The bulk improves and the high-end tail extends each round.

Mean and best kon.

The summary line plot quantifies the trend: mean increases nearly linearly across loops, while the best candidate improves even faster, consistent with KNN-covariance exploration plus targeted fine-tuning

Head-to-head efficiency.

With ~2,900 BD evaluations, Seed&Seek active training reaches a best simulated kon = 67,313 M⁻¹ s⁻¹ (+237% vs SpyCatcher002). A baseline "mutate-simulate-select" (RFdiffusion+MPNN only) needs ~3,500 evaluations to reach 41,022 M⁻¹ s⁻¹ (+105%). Active training achieves higher best performance with fewer simulations than the baseline, indicating better sample efficiency.

Analysis: why does the loop work and how do the knobs interact

Sample-efficiency and effect size.

Across ~2.9k BD evaluations, Seed&Seek's best simulated kon reaches 67.3k M⁻¹s⁻¹, a +237% lift over SpyCatcher002 and +64% over the baseline's best (41.0k) despite 17% fewer simulations (2,900 vs 3,500). The density curves shift right each loop and broaden asymmetrically—evidence that we're not just nudging the mean but expanding the high-performance tail. The mean/best lines rise in tandem, which usually indicates the model is re-shaping the latent rather than cherry-picking rare outliers.

Exploration–exploitation schedule.

Early rounds benefit from a larger KNN spread (2.5–3.5) and higher decoding temperature; this seeds multiple pockets that later rounds can exploit. The steady rise of the mean shows exploitation is working; the sustained growth of the best shows exploration hasn't shut down. Practically, we switched to a tighter spread (~1.5–2.0) once the density mode crossed ~35–40k to consolidate gains, which is reflected by the shrinking variance but rising mean.

Limits and next steps.

Plateauing will eventually occur as the frontier meets the model's prior or physics constraints. When the mean curve's slope begins to flatten, two actions usually help: (i) increase repulsion margin or decrease knn_k slightly to discover new pockets; (ii) inject small wet-lab batches (if available) to recalibrate BD biases. For deployment, we'd also add uncertainty-aware selection (e.g., ensemble perf heads) so each round mixes high-score and high-uncertainty candidates.