SWiFT: Soft-Mask Weight Fine-Tuning for Fairness

• SWiFT is a post-hoc debiasing framework that uses soft-mask generation to identify and modulate parameters driving bias versus prediction.
• It employs Fisher Information Matrix-based measures to compute continuous mask values, enabling selective gradient scaling during fine-tuning.
• The two-step fine-tuning process preserves core predictive features while reducing bias, improving both in-distribution and out-of-distribution performance.

Soft-Mask Weight Fine-Tuning (SWiFT) is a post-hoc debiasing framework for trained machine learning models that seeks to improve fairness while preserving discriminative performance through parameter-selective fine-tuning. It is defined by two coupled components: a soft-mask generation procedure that estimates each parameter’s relative contribution to bias and predictive performance, and a two-step fine-tuning process that uses those mask values to modulate gradient flow and partially reinitialize the classification head. In the reported formulation, SWiFT requires only a small external dataset and only a few epochs of model fine-tuning, and it is evaluated in bias-sensitive healthcare settings using Statistical Parity Difference (SPD), Equalized Odds (EOdds), and AUC, with particular emphasis on out-of-distribution (OOD) generalization (Yan et al., 26 Aug 2025).

1. Motivation and problem setting

SWiFT is motivated by the observation that trained models in ethically sensitive domains such as healthcare can exhibit bias with respect to sensitive attributes such as gender, skin tone, and age. The framework is introduced against two specific limitations of existing debiasing approaches: pre-/in-processing methods demand the original training data and substantial retraining, and fairness improvements often come with a reduction in discriminative accuracy. The central premise is that biased models often still capture useful core features, so complete retraining is unnecessary; instead, post-hoc interventions should distinguish, at the parameter level, how much each parameter contributes to bias versus predictive performance (Yan et al., 26 Aug 2025).

The operational setting is explicit. Given a biased trained model $f_\theta(\cdot)$, decomposed into a feature extractor $E$ and a classification head $C$, and a small, group-balanced external dataset $\mathcal{D}_e$, the goal is to fine-tune the model for improved fairness while preserving or improving AUC. The fairness objectives are expressed through common group-fairness metrics, especially SPD and EOdds, and the framework is presented as applicable to skin tone, gender, and age, with an extension to multi-attribute settings.

This positioning differentiates SWiFT from indiscriminate full-model retraining. Rather than treating all parameters as equally suspect or equally useful, SWiFT assumes that parameter contributions are heterogeneous and that debiasing should exploit this heterogeneity. This suggests a parameter-modular view of bias mitigation in which some weights are closely tied to core predictive structure, whereas others are more strongly implicated in spurious correlations.

2. Soft-mask construction from bias and prediction importance

The first major component of SWiFT is soft-mask generation. For each parameter $\theta_i$, the framework quantifies importance to prediction and importance to bias using the diagonal of the Fisher Information Matrix (FIM). Prediction importance is defined through the gradient of weighted binary cross-entropy (WBCE), while bias importance is defined through the gradient of a differentiable proxy for the Equalized Odds function (Yan et al., 26 Aug 2025).

For parameter $\theta_i$, the reported quantities are

$$\mathcal{I}_l(\theta_i, \mathcal{D}_e) = \mathbb{E}\left[\left(\frac{\partial \mathcal{L}_{\mathrm{WBCE}}}{\partial \theta_i}\right)^2\right]$$

and

$$\mathcal{I}_b(\theta_i, \mathcal{D}_e) = \mathbb{E}\left[\left(\frac{\partial \mathcal{B}}{\partial \theta_i}\right)^2\right].$$

These per-parameter importance estimates are then converted into a continuous mask value:

$$\mathcal{M}_i = \left| \tanh \left( \frac{\mathrm{Norm}(\mathcal{I}_{i,b})}{\mathrm{Norm}(\mathcal{I}_{i,l})} \right) \right|$$

where $\mathrm{Norm}$ denotes min-max normalization per layer. The resulting $E$0 is interpreted as the relative importance of a parameter to bias over prediction. High mask values are assigned to parameters that strongly drive bias but weakly contribute to correct prediction; low values correspond to parameters that are crucial to prediction and should be updated less.

This construction is a defining feature of SWiFT. The mask is neither a binary keep/drop operator nor an unstructured regularizer. It is a continuous control signal over gradient magnitude, derived from relative sensitivity rather than absolute magnitude. In the reported interpretation, the soft mask disentangles core feature parameters from those responsible for bias, enabling targeted updates that remove bias features while preserving, or even enhancing, class-discriminative capability.

3. Two-step fine-tuning dynamics

SWiFT’s second major component is a two-step fine-tuning strategy that uses the soft mask to impose differential update dynamics across the parameter space. The first step fine-tunes the feature extractor while freezing the classification head; the second step partially reinitializes the classification head and fine-tunes it while freezing the feature extractor (Yan et al., 26 Aug 2025).

In the first step, the optimization objective is

$E$1

where $E$2 is the bias function and $E$3 balances accuracy versus fairness. At this stage, $E$4 is set small. During backpropagation, each parameter’s gradient is scaled by its mask value:

$E$5

Only the feature extractor is updated, and the classification head is frozen. This design biases the optimization toward modifying parameters that are more strongly implicated in bias, while reducing perturbation of parameters that are important for predictive performance.

In the second step, parameters in the classification head with a mask value above a threshold are reinitialized to zero. The threshold is specified as $E$6 mean mask value, and the reported rule is

$E$7

After this partial reinitialization, the head is fine-tuned with the feature extractor frozen, again using the combined loss, but now with a higher $E$8 that emphasizes predictive performance. The intended effect is to relearn the classifier on top of debiased features without discarding all previously useful head parameters.

The two-step structure is important because SWiFT does not treat debiasing as a single homogeneous update. The feature extractor stage targets the internal representation, and the head stage re-aligns the decision layer after those representational changes. The reported ablations state that removing either step substantially worsens both AUC and bias metrics, and that partial re-initialization of the head is more effective than no or full re-initialization.

4. Experimental regime, metrics, and reported findings

The reported evaluation spans dermatological and chest X-ray settings. For skin lesion classification, models are pre-trained and fine-tuned on ISIC, with OOD evaluation on Fitzpatrick-17k, DDI, Atlas, and PAD-UFES-20 for skin tone and gender. For chest X-ray classification, models are pre-trained and fine-tuned on MIMIC-CXR, with OOD evaluation on CheXpert and NIH Chest Xray8 for gender and age. The architectures listed are ResNet-50, EfficientNet-B3, and DenseNet-121. SWiFT is applied post hoc using an external dataset of at most $E$9 of the original size, constructed from the validation split and balanced across groups; only a few fine-tuning epochs are used, typically $C$0 of standard training epochs (Yan et al., 26 Aug 2025).

Performance is assessed by AUC for accuracy and by SPD and EOdds for fairness. The reported baseline comparison includes 9 SOTA debiasing methods, including data balancing, fine-tuning, regularization, pruning methods such as FairPrune and DiffPrune, adversarial methods such as DiffGda, and the hard-mask method BMFT. The principal empirical claim is that SWiFT consistently reduces model bias while achieving competitive or even superior diagnostic accuracy under these metrics, and that the improved models show superior performance on several OOD datasets.

The paper provides concrete examples. For skin tone debiasing with ResNet-50 on Fitzpatrick-17k, the baseline is reported as AUC $C$1, SPD $C$2, and EOdds $C$3, whereas SWiFT is reported as AUC $C$4, SPD $C$5, and EOdds $C$6, corresponding to an SPD reduction of approximately $C$7, an EOdds reduction of approximately $C$8, and an AUC improvement of approximately $C$9. For chest X-ray age debiasing with ResNet-50 on CheXpert, EOdds is reported to decrease from $\mathcal{D}_e$0 for the baseline to $\mathcal{D}_e$1 for SWiFT, with no AUC drop and, in some cases, slight improvement (Yan et al., 26 Aug 2025).

The ablation findings are similarly specific. The soft mask is reported to provide a consistently better trade-off than the hard-mask BMFT variant and a random mask. The method is stated to be robust to the mask normalization strategy, with min-max normalization reported as best, and to the partial head re-initialization threshold, with mean mask value reported as effective. It is also reported to operate with very small external datasets, even with only $\mathcal{D}_e$2 of the usual size. Qualitatively, Class Activation Map visualizations are said to show that SWiFT shifts model attention from spurious features such as skin tone toward disease-relevant regions, for both minority and majority groups.

5. Relation to BMFT, MFT, and earlier masking methods

SWiFT belongs to a broader family of masking-based adaptation methods, but its specific mechanism differs from both hard-mask debiasing and mask learning aimed at other objectives. The most direct comparison in the provided literature is with Bias-based Weight Masking Fine-Tuning (BMFT), which is also a post-processing debiasing technique for pre-trained neural network classifiers and also uses a small, group-balanced external dataset without access to the original training data. BMFT derives a binary mask analytically from FIM-based sensitivity analysis, selecting the top $\mathcal{D}_e$3 most bias-contributing, prediction-irrelevant weights through a threshold on the ratio $\mathcal{D}_e$4, and then applies an impair-repair two-step fine-tuning process in which the classification layer is reinitialized and fine-tuned with a composite loss. By contrast, SWiFT uses a continuous mask in $\mathcal{D}_e$5 to scale gradient flow rather than a binary selection rule (Xue et al., 2024).

The relation to Mask Fine-Tuning (MFT) in LLMs is more conceptual than application-specific. MFT challenges the assumption that model integrity must be maintained during fine-tuning by freezing the post-FFT model weights and learning binary masks under the standard auto-regressive language modeling loss. Its goal is performance boost rather than efficiency, and it emphasizes hard binary masking with straight-through estimators. The MFT paper explicitly contrasts this with soft-masking approaches such as SWiFT, noting that SWiFT assigns soft, real-valued mask values to each weight and allows partial contribution and smoother gradient flow, whereas MFT learns hard binary masks in $\mathcal{D}_e$6 (Zhang et al., 27 Mar 2025).

A related computer-vision line of work studies differentiable weight masks for domain transfer. In that setting, binary masks are used to partition parameters into frozen and adaptable sets so as to mitigate forgetting on a source task while enabling target-task adaptation. The paper distinguishes naive masking, editor-network masking, and direct binary mask learning via Gumbel-Softmax, and describes the direct binary mask learning method as essentially implementing SWiFT with hard masks learned via differentiable relaxation. The reported trade-off is that stricter masking better preserves source performance but can reduce target adaptation (Khanna et al., 2023).

Historically, binary masks also appear in incremental and multi-task learning through weight transformations. A 2018 method adapts a pre-trained network to new tasks using learned binary masks and affine transformations of the base weights, achieving high adaptation performance with roughly one bit per parameter per additional task. That work is not a fairness method, but it situates SWiFT within a longer trajectory in which masks evolve from storage-efficient task adaptation mechanisms to parameter-level control structures for transfer, continual learning, pruning, and debiasing (Mancini et al., 2018).

6. Interpretation, significance, and points of caution

Within the reported evidence, SWiFT is presented as a response to the common fairness–accuracy trade-off. Other debiasing methods are described as often achieving fairness at substantial cost to accuracy, sometimes characterized as “leveling down,” whereas SWiFT is reported to lower bias and maintain or improve AUC across the tested datasets, architectures, and sensitive attributes (Yan et al., 26 Aug 2025). In that sense, the framework is positioned not merely as a fairness regularizer, but as a method for preserving or enhancing discriminative performance under post-hoc debiasing constraints.

The emphasis on OOD performance is also central. The paper states that SWiFT-improved models show superior performance on OOD datasets for both fairness and accuracy. This suggests that the method is not only suppressing an in-distribution fairness metric, but also modifying features that are implicated in spurious demographic correlations across dataset shifts. The CAM visualizations are presented as qualitative support for that interpretation by showing attention shifts toward disease-relevant regions.

At the same time, the framework is explicitly tied to a particular design space: a small, group-balanced external dataset; FIM-based per-parameter importance estimates; a differentiable proxy for Equalized Odds; and a two-step training schedule with partial head reinitialization. A plausible implication is that the effectiveness of SWiFT depends on the representativeness of the external dataset and on the match between the selected fairness proxy and the operational fairness requirement. A further plausible implication is that the continuous mask is useful precisely because it avoids the rigidity of binary inclusion/exclusion rules, allowing finer control of parameter updates than hard-mask alternatives.

More broadly, SWiFT illustrates a shift in how masking is used in modern model adaptation. In earlier work, masks were often associated with pruning, modular transfer, or continual learning. In BMFT they become a fairness-targeting mechanism based on explicit bias-versus-prediction ratios; in MFT they become a means of improving downstream performance by breaking structural integrity after full fine-tuning. SWiFT occupies a distinct point in this landscape by using soft, parameter-wise masks to modulate debiasing gradients rather than to impose hard structural sparsity or permanent binary partitioning.