Model Immunization from a Condition Number Perspective

TL;DR

This work addresses the risk of harmful fine-tuning of open-source models by introducing a Hessian condition-number perspective on model immunization. It develops a principled framework with two differentiable regularizers, $\mathcal{R}_{\text{ill}}$ and $\mathcal{R}_{\text{well}}$, and a gradient-based algorithm that increases the harmful-task Hessian conditioning while stabilizing the pre-training-task conditioning, analyzed for linear probes and extended to deep nets. Empirically, the method achieves strong immunization on linear models (high $\text{RIR}$) and shows robust performance on non-linear models like ResNet18 and ViT, indicating practical safety benefits for releasing pre-trained models. The approach provides a clear, theory-grounded pathway to preemptively curb relearning of harmful concepts while preserving utility, with broad implications for the safe deployment of open-source AI systems.

Abstract

Model immunization aims to pre-train models that are difficult to fine-tune on harmful tasks while retaining their utility on other non-harmful tasks. Though prior work has shown empirical evidence for immunizing text-to-image models, the key understanding of when immunization is possible and a precise definition of an immunized model remain unclear. In this work, we propose a framework, based on the condition number of a Hessian matrix, to analyze model immunization for linear models. Building on this framework, we design an algorithm with regularization terms to control the resulting condition numbers after pre-training. Empirical results on linear models and non-linear deep-nets demonstrate the effectiveness of the proposed algorithm on model immunization. The code is available at https://github.com/amberyzheng/model-immunization-cond-num.

Figures (4)

• Figure 1: Norm ratio Eq. \ref{['eq:norm_ratio']}vs. Epochs. We visualize the convergence of linear probing of different immunized models using gradient descent with an exact line search. Here, Identity corresponds to not using a feature extractor, i.e., $\theta_{\tt I} = {\bm{I}}$. Observe that Ours made the convergence faster on ${\mathcal{D}}_{\tt P}$ while slower in ${\mathcal{D}}_{\tt H}$ when compared to the other baselines; consistent with the results in Tab. \ref{['tab:tabular']}.
• Figure 2: Visualization of $\log(\texttt{RIR})$ of binary classification tasks created from MNIST. Each element in the figure corresponds to the $\log(\texttt{RIR})$ of a model immunized against ${\mathcal{D}}_{\tt H}$ from the pre-training task of ${\mathcal{D}}_{\tt P}$. We color the block blue if $\texttt{RIR} \gg 1$, and red otherwise. Our method succeeds in immunizing the model across all digit pairs, while the baselines failed in most pairs.
• Figure 3: Test accuracy vs. Fine-tuning Epochs on ${\mathcal{D}}_{\tt H}$. We visualize the test accuracy of linear probing on ImageNet of different immunized models using gradient descent. Here ${\mathcal{D}}_{\tt H}$ is the Stanford Cars dataset.
• Figure 4: Dummy layer with selective inverse feature covariance matrix in backward function.

Theorems & Definitions (22)

• Definition 3.0
• Proposition 3.0
• Theorem 4.1: Properties of $\kappa$-maximizing regularizer ${\mathcal{R}}_{\tt ill}({\bm{S}})$
• Theorem 4.2
• Theorem 4.3
• Theorem 1.1: Properties of $\kappa$-minimizing regularizer ${\mathcal{R}}_{\tt well}({\bm{S}})$, Theorem 2.1, 2.2, 3.1, 3.2 in Nenov2024smoothsailing
• Proposition 2.0
• proof
• proof
• Lemma 2.1