README.md

---
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.