File size: 3,872 Bytes
56cdeab
 
f946378
 
 
 
 
 
 
 
 
 
56cdeab
9c0f6fc
f946378
9c0f6fc
 
f946378
9c0f6fc
 
 
f946378
9c0f6fc
 
f946378
9c0f6fc
f946378
9c0f6fc
f946378
9c0f6fc
 
 
 
f946378
 
9c0f6fc
f946378
 
 
44a0ca3
f946378
 
9c0f6fc
f946378
 
9c0f6fc
f946378
44a0ca3
f946378
 
9c0f6fc
 
f946378
 
 
9c0f6fc
 
 
f946378
 
 
9c0f6fc
f946378
 
 
9c0f6fc
f946378
 
 
9c0f6fc
f946378
 
 
 
9c0f6fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f946378
 
 
9c0f6fc
f946378
 
9c0f6fc
 
 
 
 
 
 
 
f946378
 
 
 
9c0f6fc
f946378
9c0f6fc
f946378
 
9c0f6fc
f946378
9c0f6fc
f946378
 
 
 
 
9c0f6fc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
---
license: mit
base_model: facebook/esm2_t12_35M_UR50D
base_model_relation: finetune
tags:
  - biology
  - protein-language-model
  - antibody
  - sequence-editing
  - flow-matching
  - esm-2
  - reproduction
---
![EditJumps architecture](hero-light.svg)

One antibody sequence in, near neighbours out: insertions, deletions and substitutions placed
where related natural sequences actually differ.

- Discrete **edit-flow** generative model over amino-acid sequences.
- ESM-2 encoder trunk + per-position edit-rate and token heads.
- Trained on aligned antibody homolog pairs to carry one natural sequence to another.

> **This repo holds weights only.** Please check-out the official repo to learn more about the code, docs, training pipeline and evaluation artefacts:
> **https://github.com/VisiumCH/editjumps**

## 🚀 Quickstart

### 1. Install

```bash
git clone git@github.com:VisiumCH/editjumps.git
cd editjumps
uv sync --all-groups
```

### 2. Download the weights

```bash
pip install -U "huggingface_hub[cli]"
hf download VisiumSA/EditJumps --include "model/*" --local-dir ./weights
```

From Python:

```python
from huggingface_hub import snapshot_download

folder = snapshot_download("VisiumSA/EditJumps", allow_patterns="model/*")
```

### 3. Run inference
From the terminal directly
```bash
uv run --group train editjumps edit \
  --sequence QVQLVESGGGLVQPGGSLRLSCAAS \
  --model ./weights/model \
  --rate-head mlp \
  --q-head esm_lm_head \
  --edits 5
```

Or from Python:

```python
from pathlib import Path

from editjumps.pipeline.train.evoflows import EvoFlowsModel

model = EvoFlowsModel.load_trained(
    Path("weights/model"),
    rate_head="mlp",
    q_head="esm_lm_head",
)
```

## Model details

| | |
|---|---|
| Model type | Discrete flow matching over insert / delete / substitute |
| Trunk | Stock `facebook/esm2_t12_35M_UR50D` — 12 layers, hidden 480, 20 heads, vocab 33 |
| Heads | 3 rate heads (`mlp`) + 2 token heads, FiLM-conditioned on a time MLP <!-- TODO: confirm head count/names --> |
| Parameters | 35,213,590 |
| Precision | fp32 |
| Input form | Joined `VH.VL` antibody chains — `.` is a real token in the ESM-2 vocabulary |
| Training data | 1,660,105 unordered homolog pairs from the OAS paired corpus (2,586,927 `VH.VL` sequences) |
| Objective | Bregman rate-matching path (eq. 6), linear `kappa` schedule |
| Budget | 20,000 steps × batch 16 = 320,000 pairs ≈ **19% of one epoch** — not converged, and not claimed to be |
| Optimizer | Adam, lr 1e-4, constant; no warm-up, decay or weight decay; grad-clip 1.0, seed 0 |
| Best val loss | 130.48 at step 18,500 <!-- TODO: confirm step --> |
| Training cost | ~3 h on one A100 (~0.54 s/step); ~5.8 h on an L4 |

The trunk is deliberately **stock** ESM-2 rather than the project's OAS-adapted trunk: EvoFlows §3.4
specifies no domain adaptation, only "the encoder trunk of a pre-trained ESM-2 model".

## Citation

Cite this reproduction via `CITATION.cff` in the source repository. The papers reproduced:

```bibtex
@article{editJumps2026,
  title  = {When Edit Flows are Edit Jumps: replicating Edit Flows and {EvoFlows}},
  author = {Bénédict, Gabriel and Buechler, Melanie and Riera-Solà, Gerard and
            Ancos, Chloé de and Teukam, Yves Gaetan Nana and Freidank, Moritz},
  year   = {2026},
  note   = {arXiv:2609.18745}
}

@article{evoflows2026,
  title  = {EvoFlows: Evolutionary Edit-Based Flow-Matching for Protein Engineering},
  author = {Deutschmann and Ferragu and Bixby},
  year   = {2026},
  note   = {arXiv:2603.11703. ICLR 2026 "AI for Science" workshop}
}

@article{editflows2025,
  title  = {Edit Flows: Flow Matching with Edit Operations},
  author = {Havasi and Karrer and Gat},
  year   = {2025},
  note   = {arXiv:2506.09018. NeurIPS 2025}
}
```

## License

MIT, for both the code and these weights. See the source repository.