experiment/knowledge_editing/IMPLEMENTATION_PLAN.md

# Implementation Plan: Knowledge Editing Baselines

## Baseline Lineup

| Method | Category | Role |
|--------|----------|------|
| **LoRA v5** | Your method | Optimization-based, handles distributional visual subjects + conditional suppression |
| **OVERTONE (LoRA v3)** | Text KME baseline (adapted) | Token-level adaptive targets, shows text-only KME on VLM |
| **MEMIT** | Classic locate-then-edit | Shows point-subject methods struggle with distributional setting |
| **DualEdit** | VLM-specific editor (COLM 2025) | Best existing VLM editor, gating mechanism provides visual conditioning |

---

## Part 1: MEMIT Baseline

**Effort:** Small (1-2 days)
**Status:** EasyEdit has MEMIT at `EasyEdit/easyeditor/models/memit/`. Just needs wiring up.

### Files to Create

#### 1. `experiment/knowledge_editing/hparams/memit.yaml`

```yaml
alg_name: "MEMIT"
model_name: "llava-hf/llava-1.5-7b-hf"
device: 0
stats_dir: "./data/stats"

# Target mid-range layers for LLaMA backbone
layers: [4, 5, 6, 7, 8]
layer_selection: "all"
fact_token: "subject_last"

v_num_grad_steps: 25
v_lr: 5e-1
v_loss_layer: 31
v_weight_decay: 1e-3
clamp_norm_factor: 4
kl_factor: 0.0625
mom2_adjustment: true
mom2_update_weight: 15000

# Module templates for LLaVA-1.5's LLaMA backbone
# MEMIT operates on model.language_model (a LlamaForCausalLM)
# These paths are relative to the language_model, not the full LlavaForConditionalGeneration
rewrite_module_tmp: "model.layers.{}.mlp.down_proj"
layer_module_tmp: "model.layers.{}"
mlp_module_tmp: "model.layers.{}.mlp"
attn_module_tmp: "model.layers.{}.self_attn"
ln_f_module: "model.norm"
lm_head_module: "lm_head"

mom2_dataset: "wikipedia"
mom2_n_samples: 100000
mom2_dtype: "float32"
model_parallel: false
```

#### 2. `experiment/knowledge_editing/memit_wrapper.py`

Adapts text-only MEMIT for the multimodal pipeline:

- **Extract `model.language_model`** from `LlavaForConditionalGeneration` — MEMIT operates on CausalLM directly
- **Convert edit requests** from image-based format to MEMIT's text-only format:
  ```python
  # Our format (from build_requests):
  {"prompt": "Describe this image.", "target": "A bathroom with a sink and mirror.", "image": <PIL>}

  # MEMIT format:
  {"prompt": "In the bathroom, the {} is", "subject": "toilet", "target_new": {"str": "not visible"}}
  ```
- **Call `apply_memit_to_model()`** from EasyEdit on the language model
- **Return the full LlavaForConditionalGeneration** with edited language model

Function signature must match `run_baselines.py` pattern:
```python
def apply_memit_to_multimodal_model(model, processor, requests, hparams,
                                     copy=False, return_orig_weights=True,
                                     keep_original_weight=False):
```

**Key limitation:** MEMIT is text-only, cannot condition on images. It does an unconditional edit
suppressing the bathroom→toilet association globally. Expected to hurt locality (suppress toilet
even when present). This is the point — demonstrates why vision-conditioned methods are needed.

**Pitfalls:**
- Module path mismatch: paths must be relative to `model.language_model`, not full model
- Covariance computation (`mom2`): runs model on Wikipedia text, needs CausalLM not full VLM
- Tokenizer: use `processor.tokenizer` or `processor._processor.tokenizer`
- `target_new` vs `target`: wrapper must translate field names

### Files to Modify

#### 3. `experiment/knowledge_editing/run_baselines.py`

Add `memit` to three places:

**`load_hparams()` (~line 181):**
```python
elif method == "memit":
    from easyeditor.models.memit import MEMITHyperParams
    return MEMITHyperParams.from_hparams(yaml_path)
```

**`get_apply_algo()` (~line 197):**
```python
elif method == "memit":
    from experiment.knowledge_editing.memit_wrapper import apply_memit_to_multimodal_model
    return apply_memit_to_multimodal_model
```

**CLI args (~line 459):** Add `"memit"` to choices list.

#### 4. `experiment/evaluation/validate.py`

MEMIT edits weights directly (like LoRA merged), so treat `--model_type memit` same as `merged`.
Minimal change — just document it or add alias.

---

## Part 2: DualEdit Baseline

**Effort:** Medium (3-5 days)
**Paper:** "DualEdit: Dual Editing for Knowledge Updating in Vision-Language Models" (COLM 2025)
**Code:** https://github.com/zhiyiscs/DualEdit

### How DualEdit Works

1. Inserts **two learnable cross-attention adapters** at modality-specific key layers:
   - Layer 16: text adapter (modifies textual representations)
   - Layer 19: vision adapter (modifies visual representations)
2. **Gating mechanism**: cosine similarity of last-token representations decides whether to apply edit
   - `Sim = cos(h^e, h^i)` — edit sample vs input sample
   - If `Sim > τ`: route through adapter (apply edit)
   - Else: use original model (preserve behavior)
3. **Loss** = `L_rel + L_gen + L_loc` (reliability + generality + locality)

### Why It Fits Our Problem

- Gating = natural implementation of our `g(image)` visual grounding function
- Dual-modality editing matches our insight (prior in LLM, triggered by visual input)
- Near-perfect locality (99.89% M-Loc in paper) via gating

### Key Adaptation: Per-Sample → Distributional

DualEdit is designed for per-sample edits. We adapt to **batched training**:
- Train ONE pair of adapters across ALL bathroom-no-toilet images
- Store mean last-token representation as gate prototype
- At inference, gate fires for inputs similar to prototype

### Files to Create

```
experiment/knowledge_editing/dualedit/
    __init__.py
    adapter.py
    gating.py
    dualedit_main.py
    dualedit_hparams.py
```

#### 1. `dualedit/adapter.py` — Dual Adapter Modules

Cross-attention adapter (from paper Eq. 5, Section 3.2):

```python
class DualEditAdapter(nn.Module):
    """Learnable adapter using cross-attention.

    Inserted at a specific layer. Uses the edit signal (h_e from the edit sample)
    as Key/Value, and the current hidden state as Query.

    Separate weight matrices W_1, W_2, W_3 for text vs vision modalities.
    """
    def __init__(self, hidden_size, d_a=64):
        # W_1^{t/v}: Query projection  (hidden_size -> d_a)
        # W_2^{t/v}: Key projection    (hidden_size -> d_a)
        # W_3^{t/v}: Value projection   (hidden_size -> hidden_size)
        # Text and vision get separate weight matrices
        ...

    def forward(self, h_k, h_e_k):
        # h_k: current layer's hidden state for text or vision tokens
        # h_e_k: edit signal from edit sample at this layer
        # Returns: edited hidden state
        # h_hat = Softmax(h_k @ W_1 . (h_e_k @ W_2)^T) . h_e_k @ W_3
        ...
```

Two instances needed: one at layer 16 (text), one at layer 19 (vision).

#### 2. `dualedit/gating.py` — Gating Mechanism

```python
class DualEditGate:
    """Cosine similarity gate on last-token representations.

    At training time: stores edit sample representations.
    At inference time: compares input repr to stored prototypes.
    If similarity > threshold: activate adapters.
    Otherwise: bypass (use original model).
    """
    def __init__(self, threshold=0.6):
        self.threshold = threshold
        self.edit_prototypes = []  # mean last-token reprs from edit samples

    def register_edit_repr(self, h_last_token):
        """Store representation during training."""
        ...

    def compute_prototype(self):
        """Average all stored reprs into a single gate key."""
        ...

    def should_edit(self, h_input_last_token) -> bool:
        """Check if input should be routed through adapters."""
        sim = F.cosine_similarity(h_input_last_token, self.prototype, dim=-1)
        return sim > self.threshold
```

Threshold: τ=0.6 for LLaVA-1.5 (from paper Appendix D).

#### 3. `dualedit/dualedit_main.py` — Main Apply Function + Training

```python
def apply_dualedit_to_multimodal_model(model, processor, requests, hparams, **kwargs):
    """
    1. Insert adapter modules at layers 16 and 19
    2. Freeze all params except adapters
    3. Train adapters on edit requests:
       - L_rel: -log P(target | edit_image, prompt) on bathroom_no_toilet
       - L_gen: same on rephrased prompts + different bathroom_no_toilet images
       - L_loc: KL preservation on bathroom_with_toilet + unrelated
    4. Compute and store gate prototype (mean last-token repr)
    5. Return edited model with adapters + gating state
    """
```

Training details (from paper Appendix D):
- Learning rate: 1e-4
- Batch size: 4
- Max iterations: 50,000 (but we can use fewer for our distributional task)
- Checkpoint every 1000 iterations, select best by loss

#### 4. `dualedit/dualedit_hparams.py` — Hyperparameters

```python
@dataclass
class DualEditHyperParams:
    model_name: str = "llava-hf/llava-1.5-7b-hf"
    text_adapter_layer: int = 16
    vision_adapter_layer: int = 19
    adapter_dim: int = 64
    gating_threshold: float = 0.6
    edit_lr: float = 1e-4
    n_iterations: int = 5000
    batch_size: int = 4
    reliability_weight: float = 1.0
    generality_weight: float = 0.5
    locality_weight: float = 1.0
    checkpoint_every: int = 500
```

#### 5. `experiment/knowledge_editing/hparams/dualedit.yaml`

YAML version of the above hparams.

### Files to Modify

#### 6. `experiment/knowledge_editing/run_baselines.py`

Add `dualedit` to `load_hparams()`, `get_apply_algo()`, and CLI args.

#### 7. `experiment/evaluation/validate.py`

Add `model_type="dualedit"` support:
1. Load base LlavaForConditionalGeneration
2. Insert adapter modules at layers 16 and 19
3. Load adapter weights + gating keys from `.pt` file
4. Register forward hooks for gating

Follow existing GRACE/WISE save/restore pattern in `_extract_adapter_states()`.

### Pitfalls

1. **Layer numbering**: Verify 0-indexed layers 16 and 19 match paper's intent for LLaVA-1.5
2. **Gating prototype noise**: Mean of diverse bathroom images may be noisy — may need multiple prototypes or tuned threshold
3. **Memory**: Training requires backprop through frozen model. ~16-20GB VRAM with batch_size=4 and fp16. May need gradient checkpointing.
4. **Multimodal loss**: Locality loss must handle both text-only (unrelated) and image+text (bathroom-with-toilet) inputs

---

## Implementation Order

```
Week 1:
  Day 1-2: MEMIT
    ├── Create memit.yaml hparams
    ├── Create memit_wrapper.py
    ├── Wire into run_baselines.py
    └── Test: run MEMIT on small edit set, verify weights change, run eval

  Day 3-4: DualEdit (architecture)
    ├── Clone DualEdit repo, study source code
    ├── Create dualedit/ directory structure
    ├── Implement adapter.py (cross-attention modules)
    └── Implement gating.py (cosine similarity gate)

  Day 5: DualEdit (training + integration)
    ├── Implement dualedit_main.py (training loop)
    ├── Create hparams, wire into run_baselines.py
    └── Add save/restore + validate.py support

Week 2:
  Day 1: Testing & debugging
    ├── Run DualEdit training on bathroom dataset
    ├── Debug gating behavior (check threshold sensitivity)
    └── Verify locality on bathroom-with-toilet images

  Day 2: Joint evaluation
    ├── Run all methods through validate.py
    ├── Compare: LoRA v5, OVERTONE, MEMIT, DualEdit
    └── Generate comparison table
```

---

## Reference: Existing Infrastructure

| Component | Location | Notes |
|-----------|----------|-------|
| LoRA v5 (your method) | `experiment/training/finetune_lora_v5.py` | Already implemented |
| OVERTONE (LoRA v3) | `experiment/training/finetune_lora_v3.py` | Already implemented |
| EasyEdit baselines | `experiment/knowledge_editing/run_baselines.py` | WISE, GRACE, LoRA, IKE |
| EasyEdit MEMIT | `EasyEdit/easyeditor/models/memit/memit_main.py` | Needs wrapper |
| Evaluation | `experiment/evaluation/validate.py` | Supports lora, merged, delta_w, grace, wise |
| Edit set builder | `experiment/knowledge_editing/build_edit_set.py` | Generates edit_set.json |
| Dataset | `experiment/data/datasets.py` | CC3M bathroom categories |
| LLaVA compat | `experiment/knowledge_editing/llava15_compat.py` | Processor wrapper |
| Config | `experiment/config/train_config.py` | Comprehensive config system |
| DualEdit paper code | https://github.com/zhiyiscs/DualEdit | To be cloned |

## Reference: Key Papers

- **MEMIT**: Meng et al. (2023). "Mass-Editing Memory in a Transformer." ICLR.
- **DualEdit**: Shi et al. (2025). "DualEdit: Dual Editing for Knowledge Updating in Vision-Language Models." COLM.
- **VisEdit**: Chen et al. (2024). "Attribution Analysis Meets Model Editing." AAAI. (Not implementing, but VisEdit/ code available for reference)
- **OVERTONE**: Liu et al. (2025). "Mitigating Heterogeneous Token Overfitting in LLM Knowledge Editing." ICML.
- **AlphaEdit**: Fang et al. (2025). "AlphaEdit: Null-Space Constrained Knowledge Editing." ICLR. (Optional post-hoc enhancement)