A Strategy for Label Alignment in Deep Neural Networks

Imani et al. (2022) in their work deduced the following label alignment objective for linear regression settings in general:

where $\Phi\in\mathbb{R}^{n\times d}$ is the representation matrix with each row being the features for the linear regression, $w\in\mathbb{R}^{d\times 1}$ is the weights of the linear regression, and $y\in\mathbb{R}^{n\times 1}$ is the label vector. And $\Phi=U\Sigma V^{\top}$ is the singular value decomposition (SVD) of $\Phi$:

And the label alignment property: $y^{U}_{i}=0,\forall i\in\{k+1,...,d\}$, were used from (2) to (4), assuming that the label alignment property holds for the first $k$ singular vectors.

In the original literature, the first term in (4) was interpreted as linear regression on a smaller subspace of $\Phi$. While the second term in (4) was called the label alignment regularization and interpreted as minimizing $\sigma_{i}w^{V}_{i}=y^{U}_{i},\forall i\in\{k+1,...,d\}$, which has the effect of reducing the influence on the model’s output from those singular vectors that are not the top principal components.

Based on the above interpretation, the following objective for unsupervised domain adaptation was further developed to adapt the linear regression model from a labeled source dataset $(\Phi,y)$ to an unlabeled target dataset $(\tilde{\Phi},\tilde{y})$, with $\tilde{y}$ being unknown:

where the first term is the typical linear regression loss on the source domain, the second term removes the label alignment included in $\min_{w}{\|\Phi w-y\|^{2}}$ as shown in (4), and the third term enforces the label alignment on the target domain with rank $\tilde{k}$.

Directly applying the same rigorous deduction above onto the case of deep neural networks (DNNs) is challenging due to both the diversity and non-linearity properties of DNNs. Instead, we start by reinterpreting the objective given by (5) from a different perspective.

We start by combining (5) with (4) to form the following explicit objective equivalent to (5):

We then make the assumption that the label alignment of the source and the target dataset have approximately the same rank, such that $\tilde{k}\approx k$. This assumption makes the two terms in objective (8) independent, since then $\tilde{k}+1>k$ and all $w_{i}$ in the first term and all $w_{i}$ in the second term become mutually exclusive set. Under this assumption, (8) can be decomposed into the following two respective objectives:

We can rewrite objective (2.2) and (10) back into the following matrix forms respectively:

where $U,V$ and $\tilde{U},\tilde{V}$ follow from the SVD of $\Phi$ and $\tilde{\Phi}$ respectively, but $\Sigma^{+}$ is the reduced-upper singular value matrix $\Sigma$ of $\Phi$ containing $\{\sigma_{i}|i\in\{1,...,k\}\}$, and $\tilde{\Sigma}^{-}$ is the reduced-lower singular value matrix $\tilde{\Sigma}$ of $\tilde{\Phi}$ containing $\{\tilde{\sigma_{i}}|i\in\{k+1,...,d\}\}$. And a zero vector $y^{o}$ that doesn’t affect the optimization is introduced at (14). More formally:

Thus:

Objective (2.2) can be interpreted as performing dimensionality reduction on $\Phi$ onto the top $k$ principal components, feeding the reduced $\Phi$ into the model, and minimzing the model’s prediction loss on the reduced version of $\Phi$. Whereas objective (16), originally the label alignment regularization term, can be interpreted as performing dimensionality reduction on $\tilde{\Phi}$ onto the last $d-k$ principal components, feeding the reduced $\tilde{\Phi}$ into the model, and minimizing the model’s output on the reduced version of $\tilde{\Phi}$.

Following the intuition above, we can further deduce a general strategy for performing label alignment in DNNs. For demonstration, we start by discussing this part in the context of an example image classification task. Nevertheless, the same general strategy can be applied to other tasks as well.

Let’s define $f:\hat{X}\rightarrow\hat{\Phi}$ to be a convolutional feature extractor (convolutional neural network), $g:\hat{\Phi}\rightarrow\hat{y}$ be a feedforward neural network, where $\hat{X}\in\mathbb{R}^{n\times c\times h\times w}$ is the input images in the form of a tensor, $\hat{\Phi}\in\mathbb{R}^{n\times d}$ is the flattened output feature map from the feature extractor, and $\hat{y}\in\mathbb{R}^{n\times m}$ is the output probability matrix with $m$ being the number of classes.

Let $(X,y)$ be the source dataset, ($\tilde{X},\tilde{y}$) be the target dataset with $\tilde{y}$ unknown, $\Phi=f(X)$ be the feature map of $X$, $\tilde{\Phi}=f(\tilde{X})$ be that of $\tilde{X}$. Let $\Phi^{\prime}(\Phi,k)=U\Sigma^{+}V^{\top}$ be the reduced $\Phi$ and $\tilde{\Phi}^{\prime}(\tilde{\Phi},k)=\tilde{U}\tilde{\Sigma}^{-}\tilde{V}^{\top}$ be the reduced $\tilde{\Phi}$, using the same dimensionality reduction defined previously. To perform unsupervised domain adaptation using label alignment w.r.t. $\hat{\Phi}$, we transform and combine objective (2.2) and (16) to form the below objective function:

where $\lambda$ is a hyperparameter controlling the strength of label alignment to the target domain, and we include $f$ in objective (19) since we want to train the entire network $f(g(X))$ end-to-end.

Note that the first term in (19) is simply the classification loss on the reduced $\Phi$ and is not limited to mean-squared-error loss. It can be replaced by other loss functions such as the cross-entropy loss if desired.

Also note that the dimentionality reduction on $\Phi$ and $\tilde{\Phi}$ depends on a suitable choice of $k$ for the construction of $\Sigma^{+}$ and $\tilde{\Sigma}^{-}$. In the original work of Imani et al. (2022), $\Phi$ is a constant representation matrix irrespective of the optimization, and thus $k$ can be extracted by manually analyzing the principal components of $\Phi$. However in this case this is not feasible as $\Phi$ now varies according to $f$.

To address the above problem, we borrow some intuitions from Imani et al. (2021). We make $k$ a variable and observe that the loss term in (2.2) will be large if we choose $k>>k^{*}$ keeping all other terms constant, with $k^{*}$ being the theoretical optimal label alignment rank. Thus, following our previous derivations, minimizing the first term in (19) w.r.t. $k$ only will have the effect of approximating $k\approx k^{*}$.

Expanding upon this idea, we make $k$ a learnable parameter for our optimization objective in (19). Furthermore, in practice in our experiment, we found insignificant performance difference when alternating the minimization of (19) to be w.r.t. $f\&g$ and $k$ versus joint minimization of (19) w.r.t. $f$, $g$, and $k$ all at once. Thus, we transform (19) into the following final objective:

where the last term regularizes the learned $k$ to be small which is desired, and $\gamma$ is a hyperparameter controlling the weight of this regularization.

Additionally, to make (20) differentiable w.r.t. $k$, we perform soft-gating on $\{\sigma_{i}|i\in\{1,...,d\}\}$ and $\{\tilde{\sigma_{i}}|i\in\{1,...,d\}\}$ using the sigmoid function to approximate selective indexing for the construction of $\Sigma^{+}$ and $\tilde{\Sigma}^{-}$:

where $i\in[1,d]$ is the index of both $\sigma_{i}$ and $\tilde{\sigma_{i}}$, $k\in[0,1]$ is our aforementioned $k$ but normalized, and $\beta>0$ is a hyperparameter controlling the smoothness of the gating.

In practice, performing optimization of (20) on large dataset is infeasible for DNNs, and batch optimization with a batch of data sampled from the dataset is used instead. To make our algorithm applicable to DNNs in general, we facilitate this pattern with objective (20) with the assumption that the batch is large enough to be representative of our dataset.

Combining the aforementioned thoughts, Algorithm 1 is the final resulted pseudocode encompassing our general strategy.