MPSelectTune: Prompt-type Selection for Fine-tuning improves Concept Unlearning in LLMs

2026-07-04

AI Summary

This work addresses the challenge that current large language models struggle to robustly unlearn harmful or biased concepts across diverse prompt types, a factor overlooked by existing methods. To tackle this, the authors propose a two-stage adversarial fine-tuning approach that explicitly incorporates prompt-type selection into the concept unlearning process. Their method first identifies the “worst-case” prompt type—i.e., the one yielding the highest concept prediction accuracy—and then employs a novel multi-task loss function to minimize concept retention under this prompt while simultaneously optimizing primary task performance. Experimental results demonstrate that the proposed method improves main-task accuracy by 2–15% across four benchmarks and reduces concept prediction accuracy under the worst-case prompt by up to 17%, substantially outperforming existing baselines.

Abstract

LLMs can be conveniently adapted to a diverse set of tasks, e.g, prediction, question-answering tasks, etc, using appropriate prompts with few-shot examples. Biased or harmful concepts, e.g. gender or bio-weapons, present in pre-trained LLMs can lead to unsafe or unethical responses for many such prompts. Removing such undesirable concepts robustly across different prompt types remains a challenging problem, since existing unlearning methods typically ignore the impact of prompt variation. In this paper, we explore a novel adversarial approach to use a joint prompt for the main task and concept task prediction. We show that fine-tuning using the ``worst prompt type'' for concept prediction (with the highest concept accuracy) improves the average unlearning performance over a fine-tuning method that uses a combination of all prompt types. Our proposed method, MPSelectTune, is a two-stage approach that minimizes the concept accuracy of the highest accuracy-prompt type, after fine-tuning using a novel multi-task loss using multiple prompt types. Experimental results on four benchmarks show $2 - 15\%$ main task accuracy improvements over recent baselines and while reducing the worst-case concept accuracy by up to $17\%$ compared to recent baselines.

Problem

• concept unlearning
• prompt-type variation
• large language models
• harmful concepts
• bias mitigation

Innovation

• prompt-type selection
• concept unlearning
• adversarial fine-tuning
• multi-task loss
• LLM safety