FairViT: Fair Vision Transformer via Adaptive Masking

The transformer architecture was initially designed for natural language processing (NLP) [29] tasks. Unlike convolutional neural networks, the transformer network relies on the attention mechanism to process sequences of input tokens in parallel.

Recently, the transformer architecture has been adapted to computer vision tasks [7], utilizing the self-attention mechanism to model relationships between different parts of an image. ViT’s advantages include flexibility in handling various resolutions, capturing global information, parameter efficiency, and potential for better generalization. In many scenarios, ViT outperforms CNN and achieves considerable robustness [23].

In ViT, sample $\mathbf{x}$ contains $p$ input image patches. ViT first employs an embedding layer to each patch to convert it into a embedding vector. Subsequently, ViT applies a series of transformer encoder layers to the embeddings, and each encoder layer consists of two parts: a Multi-Head Attention mechanism (MHA) and a position-wise FeedForward Network (FFN). The MHA layer models the interactions between the patch embeddings using self-attention, while the FFN layer implements a non-linear transformation on each patch embedding individually. The self-attention mechanism [7] can be illustrated as:

where $Q$, $K$, and $V$ are the query, key, and value matrices, respectively, $\text{S}(\cdot)$ is the softmax function, $d$ is the dimension of the key vectors. Self-attention is an important building block for transformers and raises a huge amount of interest in the CV domain, since it is shown that the reliance on CNNs is not necessary and a pure transformer applied directly to sequences of image patches can perform very well on image classification tasks [7]. There are abundant works that aim to explore the transformers’ attention mechanism, such as Gradient Attention Rollout [1] in explainability and Swin Transformer [16] in efficiency.

Most of the existing debiasing methods for image classification tasks are specified for CNN or deep neural network (DNN) models [31, 21], and can not be directly applied to ViTs. However, several studies show that CV models make predictions by mixing sensitive features with input features [35, 21], and the sensitive features may capture biased relationships between the input features and the target labels. For example, the sensitive feature “gender” usually influences the accuracy of a face recognition task. In this case, it may lead to discriminatory results towards underrepresented groups, which causes serious social and ethical problems.

Fairness of ViT has been investigated in several recent works [27, 24, 22]. A targeted alignment technique TADeT was proposed in [27], which seeks to identify and eliminate bias from the query matrix in ViT. However, their directly manipulating $Q$ sacrifices accuracy for fairness. Dr-Fairness [24] proposes a bilevel optimization that finds the optimal data sampling ratios between real and generated data, and it leads to an improving tradeoff between fairness and accuracy, yet they require relatively high computing power. Debiased Self-Attention (DSA) [22] is a fairness-target approach that enforces ViT to eliminate spurious features correlated with the sensitive label, and DSA uses adversarial machine learning to enhance fairness-accuracy balance.

In this paper, we present a novel fair and accurate training framework designed for vision transformers. FairViT outperforms existing works by demonstrating superior accuracy and appreciable fairness. Furthermore, our time cost experiment and multi-task testing show that FairViT is applicable in real deployments, and maintains reasonable computational efficiency.

Cross-entropy loss uses the sign function that activates the output from the target label, deactivates the outputs from other labels. Solely minimizing cross-entropy loss actually omits the information from other labels. Therefore, we design the distance loss, considering not only maximizing the target label’s score but also minimizing other labels’ scores. Following some regularization techniques [34, 8], we formulate a regularizer that improves accuracy. During the validation phase, we use logistic regression, a binary classifier to extract a hyperplane that underlines data distribution. Subsequently, in the training stage, we leverage this hyperplane to guide the processing of each sample. This innovative strategy allows us to align the training process more closely with the actual distribution of the data, enhancing adaptability and performance.

In details, we define $\hat{y}=f_{{y}}(\mathbf{x};\boldsymbol{\theta})$ as the predicted score corresponding to the target label $y$. Additionally, we denote $\hat{y}_{k}=\sum_{i\in\{topk\}/\{y\}}f_{i}(\mathbf{x};\boldsymbol{\theta})$ as the cumulative score derived from labels in the top $k$ set, excluding the target label. In the validation stage, we train a linear classifier utilizing logistic regression, i.e. $\zeta(\hat{y},\hat{y}_{k})=\mathcal{S}(\hat{y}+\omega\hat{y}_{k}+\beta)$, where $\mathcal{S}(x)=\frac{1}{1+e^{-x}}$ is the sigmoid function, $\omega$ and $\beta$ are trainable parameters. The samples are labeled by $z=\mathbbm{1}(\hat{y}={\max}_{i}(f_{i}(\mathbf{x};\boldsymbol{\theta})))\in\{0,1\}$, indicating if the sample is classified correctly. The decision boundary of this linear classifier can be shown as

Since $\omega$ and $\beta$ are updated during the validation stage, they remain constant in the training stage. Then we introduce the distance term as [2] for the training stage

which measures the distance between point $(\hat{y},\hat{y}_{k})$ and the hyperplane in Equation (9). We train $\zeta$ in the validation stage to obtain $\omega$ and $\beta$, which remain constant in the next training stage. These fixed values facilitate the computation of our distance loss. To elaborate, the distance loss is as follows:

where $\gamma$ is a non-negative hyperparameter. Minimizing $L_{dist}$ makes the points $(\hat{y},\hat{y}_{k})$ that $\hat{y}+\omega\hat{y}_{k}+\beta\geq 0$ retain, and the points that $\hat{y}+\omega\hat{y}_{k}+\beta<0$ are moved to the decision boundary, as our objective is to shift all the points to $\hat{y}+\omega\hat{y}_{k}+\beta\geq 0$. The overall loss function consists of two parts, shown as

where $L_{ce}=\text{CE}(f(\textbf{x};\boldsymbol{\theta}),y)$, and it guides the initial training phase when the hyperplane lacks meaningful definition. During the first epoch, $\boldsymbol{\theta}$ undergoes training solely based on $L_{ce}$. In each subsequent of the epochs, we update and refine the hyperplane, training the model with $L=L_{ce}+\alpha L_{dist}$.

Introducing the distance loss contributes to accuracy enhancement. Recall the definitions of $\hat{y}$ and $\hat{y}_{k}$, since the model selects the highest score’s label as the predicted label, a higher $\hat{y}$ and a lower $\hat{y}_{k}$ suggest a greater likelihood of correct classification. Because we determine $\omega$ by the logistic regression on $z$, $\omega$ should assume a negative value, resulting in a positive slope for the hyperplane in Equation (9). Equation (10) is the distance between $(\hat{y},\hat{y}_{k})$ and the hyperplane. In Equation (11), $L_{dist}$ encourages both points that satisfy $\hat{y}+\omega\hat{y}_{k}+\beta<0$ and $\hat{y}+\omega\hat{y}_{k}+\beta\geq 0$ to move towards $\hat{y}+\omega\hat{y}_{k}+\beta\geq 0$. In both scenarios, the distance loss encourages an increase in $\hat{y}$ and a decrease in $\hat{y}_{k}$, ultimately improving accuracy. Furthermore, we can extend the distance loss directly to other models, such as deep neural networks (DNNs) and convolutional neural networks (CNNs), as it solely requires the information from the model’s output.

We conduct experiments in three distinct scenarios on the CelebA dataset [17], a large-scale dataset of facial attributes, containing over 200,000 images of celebrities’ faces. We attached our code in the supplementary materials and we will open-source it upon publication. In standard deployments, we initalize $\varsigma$ as $\mathbf{2}$ and clamp the weight to the range $(\epsilon,4-\epsilon)$, where $\epsilon$ is a small value, set as $1e-8$ in our experiment. We initialize $\mathbf{M}_{l,h}$ with $\mathbf{1/2G}$ and clamp it to the range $(-1,1)$, therefore $\widetilde{\mathbf{M}}_{l,h}$ is initialized with $\mathbf{1}$. We set the training-validation split ratio as $0.9:0.1$, $\gamma=0.5$, $G=10$ and $\alpha=0.01$ in our experiments. For measuring fairness, the evaluation metric is as follows:
Balanced Accuracy (BA) [20] measures the performance of a classification model, particularly when dealing with imbalanced datasets. Specifically, The formulation is shown as

It takes into account the imbalance in the dataset by calculating the average accuracy of each sensitive group and the target label.

We select five SOTA fairness-aware baselines to compare with our work, i.e., Vanilla [7], TADeT-MMD [27], TADeT [27], FSCL [21] and FSCL+ [21]. As the source code of DSA [22] is not publicly available, we did not include DSA in our experiment. We implement Vanilla, FSCL and FSCL+ based on their published source codes, TADeT-MMD, and TADeT using illustrations in the paper. Table 1 demonstrates that FairViT exhibits superior fairness performance alongside appreciable accuracy. In comparison to FSCL+, which is our main competitor, FairViT achieves a significantly higher accuracy of at least 4.5%. In terms of fairness metrics, FairViT showcases excellent unbiased effects.

Method ablation study: We conducted an ablation study to evaluate the effectiveness of adaptive masking and the distance loss. The results are presented in Table 2. Here, $\Theta$ denotes the adaptive masking method without updating masks and weights, while $\Delta\Theta$ signifies the deployment of adaptive masking with updating masks and weights. The experiment results show that both $L_{dist}$ and $\Delta\Theta$ contribute to accuracy improvement, with $\Delta\Theta$ playing a much more crucial role in enhancing accuracy and fairness. Continuously updating adaptive masks performs much better in both accuracy and fairness than static masks.

Impact of $\alpha$ in the loss function: In Table 3, we observe that as $\alpha$ increases, the accuracy gradually improves while EO and BA are not hurt at all. we can empirically observe that $\alpha$ is even beneficial to fairness at around 0.1. When $\alpha$ reaches a certain threshold (e.g. between 0.1 and 1 in the second case from Table 3), the model experiences a decrease in both accuracy and fairness. This phenomenon could be attributed to the model’s need to strike a balance between the distance loss and the cross-entropy loss. An inappropriate $\alpha$ might disrupt the optimization process, keeping the model from focusing on the optimization objective. For subsequent experiments, we set $\alpha=0.01$.

Impact of $\gamma$ in Distance Loss: In Table 4, we note that as the values of $\gamma$ increase, both EO and BA typically exhibit an initial rise until reaching 0.5, after which they decline. Meanwhile, accuracy follows an increasing trend till $0.5$, then it begins to decrease. We can observe $\gamma$ positively influences the impact of the distance loss, benefiting accuracy. However, excessively high values of $\gamma$ might create an imbalance between the two losses, potentially leading to decreased accuracy. Given these observations, we use $\gamma=0.5$ as it represents an optimal balance between fairness and accuracy based on our analysis.

Impact of $G$ in the adaptive masking: In Figure 4, we demonstrate the impact of $G$, the number of distinct parts, in the adaptive masking. When $G$ is small, the accuracy and fairness matrices do not bring huge benefits, however, as $G$ reaches a threshold, the accuracy and fairness tend to be stabilized and perform well. A possible explanation is that $G$ has a positive correlation with the adjustment capabilities of the model, as the model can consider more different parts at the same time and judge them in a more personalized way. However, an exceeding $G$ may result in too few images in one part, where each part will not have enough training data to obtain adequate performance, resulting in a slight decrease in performance. Therefore, the best choice of $G$ may vary across different problem scenarios and datasets.

We conduct an interpretability study on FairViT to elucidate the reasons behind its superior performance across diverse scenarios, shown in Figure 5. We use Gradient Attention Rollout (GAR) [1] to generate heat maps that accentuate crucial decision-making zones in ViT, and details about GAR are in Appendix A. We observe there are noteworthy distinctions of focused areas between Vanilla and FairViT. The Vanilla method appears to capture information relevant to sensitive attributes like Gender in the first scenario and Hair color in the second scenario. In contrast, FairViT, leveraging adaptive masking, demonstrates a tendency to extract information relevant to the target attributes, such as Expression in the first scenario and Attraction in the second scenario. This phenomenon illustrates the effectiveness of FairViT in fairness and accuracy. Furthermore, FairViT generates heat maps that are distributed more distinctly and densely in space, potentially indicating enhanced model learning.

To evaluate the efficiency of FairViT , we conduct a comparative analysis of computational costs between FairViT and the baselines, as illustrated in Table 5. Our findings reveal that FairViT exhibits comparable computational costs to the baselines. Moreover, FairViT attains a superior balance between accuracy and fairness while maintaining reasonable computational efficiency. Compared with FSCL+, FairViT is 6 times faster while achieving better accuracy and competitive fairness results. The core incremental consumption in FairViT is the adaptive masking, requiring $O(p*d)$, where $p$ is the patch number and $d$ is the dimension of the key. The time complexity is better than FSCL, which requires cubic time complexity.