Skip to content

Machine Unlearning via Task Simplex Arithmetic

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Task vector arithmetic for VLM unlearning is unstable because individual task vectors vary across fine-tuning runs, and this prediction-level variance correlates directly with poor forget-set erasure. This paper models the space of all possible task vectors as a (Q−1)-dimensional simplex and derives a closed-form ensemble over an infinite number of uniformly sampled interpolated task vectors using a second-order Taylor expansion and Dirichlet distribution moments. The result is a retraining-free unlearning method that outperforms prior task-vector and model-merging baselines while matching the cost of single task vector arithmetic.

Problem

Standard task vector unlearning (subtract fine-tuned−pretrained parameter delta from θ₀) is sensitive to fine-tuning configuration: different augmentations yield different task vectors with high prediction variance, which weakens forget-set erasure. Naively ensembling Q independently fine-tuned functions reduces variance (per the Bienaymé formula) but requires Q full fine-tuning runs and Q-times the inference cost, making it computationally prohibitive.

Method

Given Q task vectors τ₁,…,τ_Q (each from fine-tuning CLIP on the forget set under distinct RandAugment configs), the method defines a (Q−1)-simplex ∆_θ whose vertices are these vectors. The unlearning prediction is the expectation of f(x; θ₀ − λτ) over τ drawn uniformly from ∆_θ.

A second-order Taylor expansion of f around θ₀ reduces this expectation to three terms: (0th) f(x; θ₀), (1st) −λ⟨J_f(θ₀), µ_τ⟩, and (2nd) a quadratic term ω requiring E[τ^⊤Hτ] over the simplex. Theorem 1 derives this in closed form using the known moments of the Dirichlet(1_Q) distribution: E[αᵢ²] = 2/Q(Q+1), E[αᵢαⱼ] = 1/Q(Q+1) for i≠j. This yields an exact closed-form ensemble of an infinite number of interpolated task vectors with no sampling overhead.

Additional components: - Advanced aggregator (Theorem 2): replaces average pooling with probability-of-at-least-one-success over infinite Bernoulli trials, computed via exp(|∆_θ|·µ_g) where g = log(1 − f). - Vertex importance weighting (Corollary 1): extends Theorem 1 to weighted vertices; weights are optimized to minimize ensemble variance. - Vertex perturbation: small ℓ₂-ball perturbations ∆τᵢ of simplex vertices, also optimized for variance reduction. - Distillation: optionally distills the ensemble back to a single task vector τ⁻ via ℓ₂ alignment to µ_f with a variance penalty β·σ²_f, enabling deployment without an ensemble.

Key Contributions

  • Empirical and theoretical analysis linking prediction-level variance (Bienaymé formula) to unlearning effectiveness across 8 forget datasets.
  • (Q−1)-task simplex formulation capturing the convex hull of Q task vectors and enabling infinite interpolation.
  • Closed-form ensemble (Theorem 1 + Corollary 1) via Dirichlet distribution moments, avoiding any sampling from the simplex.
  • Advanced aggregator (Theorem 2) extending Poisson Binomial PMF to infinite task-vector trials.
  • Interpolation bound (Theorem 4) characterizing deviation of any simplex sample from the convex combination of vertex outputs.
  • Incremental unlearning protocol for sequential multi-dataset removal.
  • Multimodal extension validated on CLEAR VQA with LLaVA-1.5-7B.

Results

Standard unlearning (forget accuracy ↓, retain ~95% of pretrained): - ViT-B/32: Ours 15.20% vs. Function Ensemble 22.75%, EMR-Merging 21.83% (best prior) - ViT-B/16: Ours 12.17% vs. Function Ensemble 19.89%, EMR-Merging 19.10% - ViT-L/14: Ours 9.98% vs. Function Ensemble 16.32%, EMR-Merging 15.67%

Linear task vector unlearning: - ViT-B/16: Ours 6.46% vs. EMR-Merging 7.09% (best prior) - ViT-L/14: Ours 5.70% vs. EMR-Merging 6.51%

Incremental unlearning (CLIP ViT-B/32, avg forget after all 8 steps): - Ours 20.1% vs. EMR-Merging 26.8% (best prior)

Multimodal (CLEAR VQA, LLaVA-1.5-7B): - Forget: Ours 31.5% vs. EMR-Merging 34.8%; Retain: 50.2% vs. 49.3%

Efficiency: Closed-form solution achieves ~15.20% forget accuracy in 0.8 hours vs. 6.3 hours for 300-sample explicit simplex sampling at comparable performance.

Limitations

  • Requires Q independent fine-tuning runs on the forget set to obtain simplex vertices (Q=30 in experiments), which is still non-trivial for very large VLMs.
  • Second-order Taylor approximation is exact only for small λ; accuracy degrades if λ is large or task vectors have large norm.
  • Theorem 1's closed form uses Dirichlet(1_Q) (uniform); non-uniform α requires recomputing moments via Corollary 1.
  • Hessian computation (m×m for m=86M parameters) is approximated; full Hessians are infeasible and the paper relies on the Taylor truncation.
  • Retain-set evaluation uses only ImageNet; generalization of retained knowledge to other out-of-distribution tasks is not assessed.
  • Incremental unlearning is demonstrated only on ViT-B/32; interference effects at scale are unexplored.

Relevance to Vision-Language Models

This work directly targets CLIP-based VLMs (ViT-B/32, ViT-B/16, ViT-L/14) and extends to LLaVA-1.5-7B, addressing the growing need for privacy-compliant deployment of large vision-language models under GDPR and similar regulations. It advances the task vector arithmetic paradigm—a lightweight, fine-tuning-free editing tool increasingly used in VLM customization—by grounding it in a principled bias-variance framework. The closed-form simplex aggregation is particularly relevant because it scales to the parameter dimensionality of modern VLMs (86M+) without ensemble inference overhead, making it practical alongside PEFT methods like LoRA. For VLM researchers, the variance-unlearning correlation finding (Figure 1) provides a new diagnostic lens for evaluating the reliability of any task-vector-based model editing scheme.