Upload README.md with huggingface_hub

README.md

@@ -12,36 +12,40 @@ The GitHub with the implementation and requirements.txt can be found [here](http
 
 # FastESMFold
 
-FastESMFold is a self-contained, HuggingFace-compatible reimplementation of ESMFold with **built-in Test-Time Training (TTT)** and multi-backend attention (SDPA, Flash, Flex).
+FastESMFold is a self-contained, HuggingFace-compatible reimplementation of ESMFold with optional **Test-Time Training (TTT)** and multi-backend attention (SDPA, Flash, Flex).
 
 No dependency on `fair-esm`, `proteinttt`, or `openfold`. Just `transformers`, `torch`, and `einops`.
 
-## Key Features
+## Why Test-Time Training?
 
-- **Always-on TTT**: Runs 10 steps of masked language model adaptation via LoRA before folding. Returns the structure with the highest pLDDT across all steps.
-- **Best structure selection**: Folds after each TTT step, tracks per-step pLDDT, returns the best.
-- **FastESM2 attention**: SDPA/Flash/Flex backends for the 3B ESM2 backbone.
-- **Self-contained LoRA**: lora_diffusion-compatible implementation (no peft dependency). `Normal(0, 1/r)` initialization, `scale=alpha`.
-- **3.5B parameters**: Full ESMFold v1 architecture (ESM2-3B backbone + folding trunk).
+Protein language models like ESM2 are trained on millions of sequences, but at inference time they process each new protein in a single forward pass with no adaptation. This is a missed opportunity: the input sequence itself contains structural signal that the model could learn from.
 
-## Benchmark
+**Test-Time Training (TTT)** adapts the model to each individual protein before predicting its structure. The idea is simple: before folding, we briefly train the ESM2 backbone on the input sequence using masked language modeling (the same objective it was pretrained with). This forces the model to "study" the specific sequence, strengthening its internal representation of that protein's structural features.
 
-Tested on 10 difficult sequences on A10G GPU:
+The adaptation uses **LoRA** (Low-Rank Adaptation) for efficiency: only small adapter weights are trained (~4.4M parameters out of 3.5B), and the base model is restored after each prediction. This takes 20-45 seconds per sequence on an A10G GPU but can dramatically improve structure prediction quality, especially on difficult targets where standard ESMFold produces low-confidence predictions.
 
-| Metric | Value |
-|--------|-------|
-| Mean baseline pLDDT | 0.549 |
-| Mean best TTT pLDDT | 0.637 |
-| Mean improvement | +0.088 |
-| Sequences improved >5pt | 5/10 |
-| Time per sequence | ~20-45s |
-| GPU memory peak | 18.3 GB |
+**When is TTT most useful?**
+- Sequences with low baseline pLDDT (< 0.5): TTT can improve pLDDT by 10-30+ points
+- Novel proteins with limited homology in training data
+- Disordered or multi-domain proteins where ESMFold struggles
+
+**When is TTT unnecessary?**
+- Sequences that already fold well (baseline pLDDT > 0.7): TTT rarely helps and may slightly degrade predictions
+- High-throughput screening where speed matters more than accuracy
+
+## Key Features
 
-On the hardest sequence (baseline pLDDT 0.38), TTT achieves 0.72 (+34 points).
+- **Standard ESMFold**: Full ESMFold v1 structure prediction, loadable via `AutoModel`
+- **Optional TTT**: Enable test-time training for improved structure prediction on difficult sequences
+- **Best structure selection**: When TTT is enabled, folds after each step and returns the structure with the highest pLDDT
+- **FastESM2 attention**: SDPA/Flash/Flex backends for the 3B ESM2 backbone
+- **Self-contained LoRA**: lora_diffusion-compatible implementation (no peft dependency)
+- **3.5B parameters**: Full ESMFold v1 architecture (ESM2-3B backbone + folding trunk)
 
 ## Use with transformers
 
-### Basic usage
+### Standard structure prediction (no TTT)
+
 ```python
 import torch
 from transformers import AutoModel
@@ -52,16 +56,38 @@ model = AutoModel.from_pretrained(
     torch_dtype=torch.float32,
 ).cuda().eval()
 
+# Standard fold (no TTT)
+with torch.no_grad():
+    output = model.infer("MKTLLILAVVAAALA...")
+pdb_strings = model.output_to_pdb(output)
+plddt = output["plddt"].mean().item()
+print(f"pLDDT: {plddt:.3f}")
+```
+
+### Structure prediction with TTT
+
+TTT adapts the ESM2 backbone to a specific input sequence via masked language modeling before folding. This can dramatically improve pLDDT on difficult sequences (e.g., 0.38 to 0.72).
+
+```python
+# Configure TTT
+model._ttt_cfg.steps = 10      # 10 optimizer steps (default)
+model._ttt_cfg.lora_rank = 8   # LoRA rank (default)
+model._ttt_cfg.lora_alpha = 32 # LoRA scale (default)
+
+# fold_protein() runs TTT, folds after each step, returns best structure
 result = model.fold_protein("MKTLLILAVVAAALA...")
 print(f"pLDDT: {result['plddt']:.3f}")
-print(f"Best step: {result['best_step']}")
-print(f"Step pLDDTs: {result['step_plddts']}")
-print(f"PDB length: {len(result['pdb_string'])} chars")
+print(f"Best step: {result['best_step']} (0=baseline, 1-10=TTT steps)")
+print(f"Step pLDDTs: {[f'{p:.2f}' for p in result['step_plddts']]}")
+
+# Save PDB
+with open("structure.pdb", "w") as f:
+    f.write(result["pdb_string"])
 ```
 
 ### Return values
 
-`fold_protein(sequence)` returns a dict with:
+`fold_protein(sequence)` returns a dict:
 
 | Key | Type | Description |
 |-----|------|-------------|
@@ -71,28 +97,38 @@ print(f"PDB length: {len(result['pdb_string'])} chars")
 | `step_plddts` | list[float] | pLDDT at each step [baseline, s1, ..., s10] |
 | `best_step` | int | Which step produced the best structure (0=baseline) |
 
-### Loading from Synthyra/ESMFold-v1 with custom config
+### Disabling TTT
+
+To use FastESMFold as a standard ESMFold (no TTT), set `steps=0` or call `infer()` directly:
+
 ```python
-from esmfold.modeling_fast_esmfold import FastEsmFoldConfig, FastEsmForProteinFolding
-
-config = FastEsmFoldConfig.from_pretrained("Synthyra/ESMFold-v1")
-config.attn_backend = "sdpa"
-config.ttt_config = {
-    "lr": 4e-4,
-    "steps": 10,
-    "ags": 4,
-    "batch_size": 4,
-    "lora_rank": 8,
-    "lora_alpha": 32.0,
-    "seed": 0,
-}
-model = FastEsmForProteinFolding.from_pretrained(
-    "Synthyra/ESMFold-v1",
-    config=config,
-    torch_dtype=torch.float32,
-).cuda().eval()
+# Option 1: Set TTT steps to 0
+config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True)
+config.ttt_config = {"steps": 0}
+model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True)
+result = model.fold_protein("MKTLLILAVVAAALA...")  # No TTT, just baseline fold
+
+# Option 2: Call infer() directly (inherited from EsmForProteinFolding)
+with torch.no_grad():
+    output = model.infer("MKTLLILAVVAAALA...")
+pdb_strings = model.output_to_pdb(output)
 ```
 
+## TTT Benchmark
+
+Tested on 10 difficult sequences on A10G GPU:
+
+| Metric | Value |
+|--------|-------|
+| Mean baseline pLDDT | 0.549 |
+| Mean best TTT pLDDT | 0.637 |
+| Mean improvement | +0.088 |
+| Sequences improved >5pt | 5/10 |
+| Time per sequence | ~20-45s |
+| GPU memory peak | 18.3 GB |
+
+On the hardest sequence (baseline pLDDT 0.38), TTT improves to 0.72 (+34 points).
+
 ## Attention backends
 
 The ESM2 backbone supports multiple attention backends via `config.attn_backend`:
@@ -104,41 +140,60 @@ The ESM2 backbone supports multiple attention backends via `config.attn_backend`
 | Flex Attention | `"flex"` | Skips padding tokens via block mask. First use compiles a Triton kernel. |
 | Auto | `"auto"` | Picks best available: `kernels_flash` > `flex` > `sdpa`. |
 
+```python
+from transformers import AutoConfig, AutoModel
+
+config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True)
+config.attn_backend = "kernels_flash"
+model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True)
+```
+
 ## TTT Configuration
 
-TTT parameters can be customized via `config.ttt_config`:
+TTT parameters are set via `config.ttt_config` (a dict) or by modifying `model._ttt_cfg` after loading:
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
 | `lr` | 4e-4 | Learning rate for SGD optimizer |
-| `steps` | 10 | Number of optimizer steps |
+| `steps` | 10 | Number of optimizer steps (0 to disable TTT) |
 | `ags` | 4 | Gradient accumulation steps per optimizer step |
 | `batch_size` | 4 | Batch size for masked language model training |
 | `mask_ratio` | 0.15 | Fraction of tokens to mask |
-| `lora_rank` | 8 | LoRA rank (0 for full fine-tuning) |
-| `lora_alpha` | 32.0 | LoRA scaling factor |
-| `seed` | 0 | Random seed for reproducibility |
+| `lora_rank` | 8 | LoRA rank (0 for full backbone fine-tuning) |
+| `lora_alpha` | 32.0 | LoRA scaling factor (applied as `scale=alpha`, matching lora_diffusion) |
+| `seed` | 0 | Random seed for reproducible LoRA initialization and masking |
+| `lora_target_class` | `"EsmSelfAttention"` | Which module class to inject LoRA into |
 
 ## How TTT Works
 
 1. **Baseline fold** (step 0): Standard ESMFold prediction
-2. **LoRA injection**: Rank-8 LoRA adapters on ESM2 attention Q/K/V projections
-3. **Masked LM training**: 10 optimizer steps of BERT-style masked language modeling on the input sequence
+2. **LoRA injection**: Rank-8 LoRA adapters on all `nn.Linear` layers inside ESM2 attention modules
+3. **Masked LM training**: 10 optimizer steps (each with 4 gradient accumulation sub-steps) of BERT-style masked language modeling on the input sequence
 4. **Per-step folding**: After each optimizer step, fold the sequence and record pLDDT
 5. **Best selection**: Return the structure with the highest pLDDT
 6. **Reset**: Restore LoRA weights to initial state for the next sequence
 
-This is based on the ProteinTTT paper (test-time compute for protein structure prediction).
+## Citations
 
-### Citation
-If you use this implementation please cite it:
-```
-@misc {FastPLMs,
-    author       = { Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
-    title        = { FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
+If you use this implementation, please cite FastPLMs and the original ProteinTTT paper:
+
+```bibtex
+@misc{FastPLMs,
+    author       = {Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
+    title        = {FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
     year         = {2024},
-    url          = { https://huggingface.co/Synthyra/ESMplusplus_small },
-    DOI          = { 10.57967/hf/3726 },
-    publisher    = { Hugging Face }
+    url          = {https://huggingface.co/Synthyra/ESMplusplus_small},
+    DOI          = {10.57967/hf/3726},
+    publisher    = {Hugging Face}
+}
+
+@misc{bushuiev2026proteinneed,
+    title        = {One protein is all you need},
+    author       = {Anton Bushuiev and Roman Bushuiev and Olga Pimenova and Nikola Zadorozhny and Raman Samusevich and Elisabet Manaskova and Rachel Seongeun Kim and Hannes St\"ark and Jiri Sedlar and Martin Steinegger and Tom\'a\v{s} Pluskal and Josef Sivic},
+    year         = {2026},
+    eprint       = {2411.02109},
+    archivePrefix= {arXiv},
+    primaryClass = {cs.LG},
+    url          = {https://arxiv.org/abs/2411.02109},
 }
 ```