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

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:

  # 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:

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):

elif method == "memit":
    from easyeditor.models.memit import MEMITHyperParams
    return MEMITHyperParams.from_hparams(yaml_path)

get_apply_algo() (~line 197):

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):

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

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

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

@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)