Skip to main content
The GradientOptimizer designs a protein by gradient descent on a differentiable representation of the sequence rather than by discrete mutation. This program hallucinates a 20-residue binder in two stages that mirror the Germinal pipeline: a logit phase, in which a soft, continuous sequence evolves, and a softmax phase, in which the temperature is annealed to sharpen it into a discrete sequence. To keep the example runnable on CPU, it uses a mock differentiable constraint in place of a real structure model; the structure of the program is identical when real constraints are substituted. Open as a runnable notebook View as a Python script
Demonstration: this example uses a mock constraint that pushes every position toward alanine, so it runs on CPU without model weights. The all-alanine result is the expected output of that mock; substitute real structure constraints (for example an AF2-backed confidence constraint) to perform actual design.

​
A differentiable constraint

Gradient optimization requires constraints that return a gradient, not just a score. Instead of the discrete function used elsewhere, a gradient constraint supplies a backward callable that reads each sequence’s continuous logits (an (L, |vocab|) matrix) and returns a GradientConstraintOutput carrying two things: gradient, the gradient of the objective with respect to those logits, and loss, the scalar objective value the optimizer minimizes. The mock below builds a target distribution that places all mass on alanine (column 0 of the vocabulary), sets the gradient to logits - target, and reports the mean squared difference as the loss; a real constraint would backpropagate through a structure predictor instead.
python
import numpy as np
from pydantic import BaseModel

from proto_language import GradientConstraintOutput
from proto_language.core import Constraint, Construct, Program, Segment
from proto_language.generator import PositionWeightGenerator, PositionWeightGeneratorConfig
from proto_language.optimizer import GradientOptimizer, GradientOptimizerConfig


class EmptyConfig(BaseModel):
    """This mock constraint takes no parameters."""


def mock_structure_backward(input_sequences, *, config, temperature=1.0, soft=1.0, **kwargs):
    """Push the logits toward a target distribution (here, all alanine)."""
    results = []
    for (seq,) in input_sequences:
        logits = seq.logits
        target = np.zeros_like(logits)
        target[:, 0] = 1.0  # prefer alanine at every position
        grad = logits - target
        results.append(
            GradientConstraintOutput(gradient=(grad,), loss=float(np.mean(grad**2)), metrics={})
        )
    return results

​
Stage 1: logit phase

The design starts from a VHH seed sequence carried on the Segment. The GradientOptimizer does not mutate that string directly; its optimization variable is each proposal’s per-position logits, which it updates by gradient descent over num_steps. PositionWeightGenerator never proposes here; it only decodes those logits back into a discrete sequence (by default the most likely token at each position) at tracked steps. GradientOptimizerConfig.germinal_logit_preset() configures this stage: 65 SGD steps with the soft blend ramping from 0 to 1 while the softmax temperature stays fixed, so the continuous logits relax freely. Passing the same target_segment to both stages is what lets stage two continue from this stage’s logits. The custom_logging callback track records the decoded sequence at each tracked step.
python
segment = Segment(sequence="EVQLVESGGGLVQPGGSLRL", sequence_type="protein", label="binder")
construct = Construct([segment])

# Record the decoded sequence at each tracked step, across both stages.
trajectory = []


def track(step, segments):
    trajectory.append(str(segments[0].proposal_sequences[0].sequence))


gen1 = PositionWeightGenerator(PositionWeightGeneratorConfig())
gen1.assign(segment)

con1 = Constraint(
    inputs=[segment],
    backward=mock_structure_backward,
    backward_config=EmptyConfig(),
    label="structure_s1",
)

stage1 = GradientOptimizer(
    target_segment=segment,
    constructs=[construct],
    generators=[gen1],
    constraints=[con1],
    config=GradientOptimizerConfig.germinal_logit_preset(),
    custom_logging=track,
)

​
Stage 2: softmax phase

The second stage applies GradientOptimizerConfig.germinal_softmax_preset(): 35 SGD steps that hold the soft blend at 1 and anneal the softmax temperature from 1 toward 0.01 on a quadratic schedule, sharpening the soft distribution into a near-discrete one. Lowering the temperature concentrates each position’s probability mass on a single token. Because it reuses the same segment object by identity, it picks up the logits stage one left behind rather than reinitializing. A fresh PositionWeightGenerator and a gradient Constraint are bound for this stage, and the same track callback keeps appending to the shared trajectory list.
python
gen2 = PositionWeightGenerator(PositionWeightGeneratorConfig())
gen2.assign(segment)

con2 = Constraint(
    inputs=[segment],
    backward=mock_structure_backward,
    backward_config=EmptyConfig(),
    label="structure_s2",
)

stage2 = GradientOptimizer(
    target_segment=segment,
    constructs=[construct],
    generators=[gen2],
    constraints=[con2],
    config=GradientOptimizerConfig.germinal_softmax_preset(),
    custom_logging=track,
)

​
Run

The Program holds both optimizers and runs them in order, so the logit phase completes before the softmax phase begins on the same segment. num_results=1 runs a single trajectory.
python
program = Program(optimizers=[stage1, stage2], num_results=1)
program.run()

​
Inspect the result

segment.result_sequences[0] is the decoded design; segment.original_sequence holds the VHH seed it started from. The printout compares the two and walks the trajectory recorded by the track callback, sampling a handful of evenly spaced snapshots to show the decoded sequence sharpening from the seed toward the mock target. With this mock constraint the optimizer drives every position to alanine, so the designed sequence is all A; substituting real structure constraints would yield a real design while the program structure stays the same.
python
result = segment.result_sequences[0]

def representative(traj, n=4):
    if len(traj) <= n:
        return traj
    idx = sorted({round(i * (len(traj) - 1) / (n - 1)) for i in range(n)})
    return [traj[i] for i in idx]

print(f"input:     {segment.original_sequence.sequence}")
print("trajectory (the soft sequence sharpens toward the mock target, alanine):")
for seq in representative(trajectory, n=5):
    print(f"  {seq}")
print(f"designed:  {result.sequence}")

input:     EVQLVESGGGLVQPGGSLRL
trajectory (the soft sequence sharpens toward the mock target, alanine):
  DMFWPQQDMPMHEAFCDEIK
  AAAAAAAAAAAAAAAAAAAA
  AAAAAAAAAAAAAAAAAAAA
  AAAAAAAAAAAAAAAAAAAA
  AAAAAAAAAAAAAAAAAAAA
designed:  AAAAAAAAAAAAAAAAAAAA

​
Next Steps

Using Generators

The gradient generator family.

Using Optimizers

The optimizer that drives this design.