Upload README.md with huggingface_hub

README.md

@@ -3,125 +3,101 @@ license: apache-2.0
 tags:
 - dna
 - genomics
-- mutation-prediction
-- sad
+- mutation
 - omni-dna
-- hyenadna
-- causal-lm
-base_model: zehui127/Omni-DNA-20M
+- sad
+base_model: Nhoodie/omni-dna-ici-dc
 pipeline_tag: text-generation
 ---
 
-# Omni-DNA SAD Mutation Model
-
-Fine-tuned Omni-DNA-20M for cross-domain HGT (Horizontal Gene Transfer) mutation prediction using **SAD (Sequential Attenuation Denoising)**.
-
-## Architecture
-
-- **Base**: [Omni-DNA-20M](https://huggingface.co/zehui127/Omni-DNA-20M) (OLMo-based, 20M params, BPE tokenizer, 250-token context)
-- **Purpose**: Given a parent DNA sequence, predict the mutated child sequence
-- **Format**: `"mutate: {parent} -> {child}"` — instruction-tuned as a causal LM
+# Omni-DNA SAD Checkpoint
 
-## SAD Pipeline
+**Sequential Attenuation Denoising (SAD)** — ICI-DC checkpoint fine-tuned on real mutation data.
 
-```
-Stage 1 (ICI — Interleaved Codon Interference):
-  Dual-model consensus generation (Omni-DNA-20M + HyenaDNA tiny-1k)
-  Every 3 codons, a 1-codon gap is introduced; both models predict the gap
-  Agreement = consensus (high quality), disagreement = contested (kept, weighted lower)
-  → Produces 8,112 synthetic mutation pairs
-
-Stage 2 (SAD — Sequential Attenuation Denoising):
-  Fine-tune on synthetic pairs first (builds broad mutation prior)
-  Then fine-tune on real data at 10x lower LR
-  Uncontradicted synthetic patterns persist, contradicted ones get attenuated
-```
+Starting from the [ICI-DC checkpoint](https://huggingface.co/Nhoodie/omni-dna-ici-dc), this model
+undergoes attenuation: real mutation pairs at 10x lower learning rate. Weights that don't contradict
+real data persist from ICI-DC; contradicted patterns get corrected.
 
-### Key Parameters
+## SAD Training Details
 
 | Parameter | Value |
-|-----------|-------|
-| **ICI Codon Gap Distance** | 3 (every 3rd codon = 9 bp spacing) |
-| **SAD Coefficient** | 4.89 (81,120 synthetic exposures / 16,585 real exposures) |
-| Stage 1 LR | 5e-5, 10 epochs |
-| Stage 2 LR | 1e-5, 5 epochs |
-| Batch Size | 16 × 2 (effective 32) |
+|---|---|
+| Base model | `Nhoodie/omni-dna-ici-dc` (ICI-DC pre-trained) |
+| Training data | 3,317 real mutation pairs |
+| Epochs | 5 |
+| Learning rate | 1e-5 (10x lower than ICI-DC) |
+| **SAD coefficient** | **4.89** (81,120 synthetic exposures / 16,585 real exposures) |
+| Batch size | 32 effective |
 | Precision | fp32 |
 
-## Checkpoints
+> **Note**: SAD coefficient of 4.89 is considered too high. A coefficient of ~1.5 was also tested
+> (16 real epochs) with similar results. See benchmarks below.
+
+## Multi-Axis Benchmark (100 test pairs, 4 models)
 
-| Path | Description |
-|------|-------------|
-| `stage1_ici/` | Weights after synthetic pre-training (ICI stage) |
-| `stage2_sad/` | Weights after real-data attenuation (final SAD model) |
+### Axis 6: Discriminative (most important)
 
-## Benchmarks (200 test pairs, corrected methodology)
+| Model | AUC | Best F1 | Score Gap |
+|-------|:---:|:---:|:---:|
+| Base Omni (no fine-tune) | 0.588 | 0.688 | 25.6 |
+| [ICI-DC](https://huggingface.co/Nhoodie/omni-dna-ici-dc) | 0.887 | 0.858 | 370.2 |
+| **SAD coeff=4.89 (this)** | **0.904** | **0.862** | **411.4** |
+| SAD coeff=1.5 | 0.908 | 0.862 | 413.7 |
 
-| Model | Levenshtein ↓ | Similarity ↑ | Mutation Recall ↑ | SAD Coeff |
-|-------|:---:|:---:|:---:|:---:|
-| Base Omni (no fine-tune) | 360.4 | 43.9% | **30.1%** | — |
-| SAD coeff=4.89 (5 real ep) | **264.3** | 43.7% | 26.3% | 4.89 |
-| SAD coeff=1.5 (16 real ep) | 271.7 | 43.1% | 29.4% | 1.50 |
+### Axis 5: Mutation Surprise
 
-Earlier runs reported inflated recall (74.3%, 67.8%) due to a bug in mutation_recall(). 
-The corrected numbers above reflect actual per-position mutation accuracy.
+| Model | Surprise | p-value | Interpretation |
+|-------|:---:|:---:|---|
+| Base Omni | +0.777 | <0.0001 | Expects parent (mutations are surprising) |
+| ICI-DC | -0.216 | 0.0001 | Expects mutations everywhere |
+| **SAD coeff=4.89** | -0.203 | 0.0003 | Partially attenuated |
+| SAD coeff=1.5 | -0.169 | 0.0026 | More attenuation |
 
-**Key finding**: Fine-tuning dramatically improves sequence length prediction 
-(637→445 bp, matching ground truth 443) and Levenshtein distance (360→264), 
-but mutation recall remains flat. The 20M param model learns sequence structure 
-but struggles to localize specific mutations.
+### Full 6-Axis Comparison
 
-### Known Issues
+| Axis | Metric | Base | ICI-DC | SAD 4.89 | SAD 1.5 |
+|---|---|:---:|:---:|:---:|:---:|
+| A1: Detection | Recall | 0.511 | 0.518 | 0.517 | 0.497 |
+| A2: Logits | Top-3 acc | 0.846 | 0.846 | 0.846 | 0.846 |
+| A3: Ti/Tv | Predicted ratio | 0.20 | 0.41 | 0.42 | 0.43 |
+| A4: Embeddings | Seq AUC | 0.41 | 0.35 | 0.40 | 0.40 |
+| A5: Surprise | ΔLL | +0.78 | -0.22 | -0.20 | -0.17 |
+| **A6: Discrim.** | **AUC** | **0.588** | **0.887** | **0.904** | **0.908** |
 
-1. **SAD coefficient too high**: At 4.89, synthetic exposures overwhelmed real data attenuation. Mutation recall dropped from base 74.3% → 67.8%. A lower coefficient (~1.5) is recommended for future runs.
-2. **Synthetic mutation rate mismatch**: Synthetic data has 17.8% mean mutation rate vs 3.9% in real data — model learns to over-mutate.
-3. **Omni `.generate()` broken after fine-tuning**: Model outputs PAD/EOS immediately. Use manual autoregressive loop with special token suppression instead (see Usage).
+### Interpretation
+
+- **ICI-DC provides the main training signal** (AUC 0.59→0.89). Synthetic data builds a strong
+  internal representation of valid mutation pairs.
+- **SAD fine-tunes that representation** (0.887→0.908). The attenuation step does help, but the
+  marginal gain over ICI-DC is modest.
+- **The model is a judge, not a generator.** 90.8% discriminative AUC but ~30% generative recall.
+  The representations encode mutation structure, but autoregressive decoding can't access it efficiently.
+- **Ti/Tv ratio converges** from 0.20 (base) to 0.43 (SAD 1.5), approaching the biological value of 0.48.
 
 ## Usage
 
 ```python
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load Stage 2 (final SAD model)
-tokenizer = AutoTokenizer.from_pretrained("Nhoodie/omni-dna-sad-mutation", subfolder="stage2_sad", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("Nhoodie/omni-dna-sad-mutation", subfolder="stage2_sad", trust_remote_code=True).to("cuda").eval()
-
-def predict_child(parent, max_new_tokens=200, temperature=0.2):
-    """Manual autoregressive generation (Omni .generate() outputs PAD after fine-tuning)"""
-    prompt = f"mutate: {parent} -> "
-    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)
-    generated = input_ids
-    suppress = {0, 1, 2, 3}  # UNK, CLS, SEP, PAD
-
-    with torch.inference_mode():
-        for _ in range(max_new_tokens):
-            logits = model(input_ids=generated).logits[:, -1, :] / max(temperature, 0.01)
-            for s in suppress:
-                logits[0, s] = float("-inf")
-            probs = torch.softmax(logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            generated = torch.cat([generated, next_token], dim=-1)
-
-    text = tokenizer.decode(generated[0, input_ids.shape[1]:], skip_special_tokens=True)
-    return "".join(c for c in text.upper() if c in "ACGT")
-
-parent = "ATGGCTAGCTGATCGATCGATCG..."
-child = predict_child(parent)
-print(child)
+tokenizer = AutoTokenizer.from_pretrained("Nhoodie/omni-dna-sad-mutation", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("Nhoodie/omni-dna-sad-mutation", trust_remote_code=True)
+
+# Score a mutation pair (discriminative use)
+prompt = "mutate: ATGGCTAGCTGA -> ATAGCTGGCTAA"
+logits = model(**tokenizer(prompt, return_tensors="pt")).logits
 ```
 
-## Dataset
+## Related
 
-Training data available at: [Nhoodie/omni-dna-sad-mutation-dataset](https://huggingface.co/datasets/Nhoodie/omni-dna-sad-mutation-dataset)
+- **[Nhoodie/omni-dna-ici-dc](https://huggingface.co/Nhoodie/omni-dna-ici-dc)** — ICI-DC checkpoint (synthetic-only)
+- **[Nhoodie/omni-dna-sad-mutation-dataset](https://huggingface.co/datasets/Nhoodie/omni-dna-sad-mutation-dataset)** — Training data
 
 ## Source Code
 
-GitHub commit: `83634b7` (branch `dev/sad-omni-hyena`, private repo)
+Git commit: `7be4e73` (branch `dev/sad-omni-hyena`, private repo)
 
-## Training Hardware
+## Citation
 
-- **GPU**: NVIDIA GTX 1080 8GB
-- **Stage 1**: ~8 min (8,112 synthetic pairs, 10 epochs)
-- **Stage 2**: ~9 min (3,317 real pairs, 5 epochs)
-- **Total**: <20 min training time
+- Omni-DNA: Zehui127 et al.
+- HyenaDNA: Nguyen et al., NeurIPS 2023
+- ENBED: Malusare et al., Bioinformatics Advances, 2024 (arXiv: 2311.02333)