On the Versatile Uses of Partial Distance Correlation in Deep Learning

More general use cases. While the above discussion was cast as comparing two networks, it is representative of a broad basket of tasks in deep learning. (a) Consider the problem of learning fair representations [73, 17, 72, 45] where the model must be invariant to one (or more) sensitive attributes. We seek latent representations, say $\Psi_{\rm pred}(X)$ for the prediction task, which minimizes mutual information w.r.t. the latent representation relevant for predicting the sensitive attribute $\Psi_{\rm sens}(X)$. Indeed, if information regarding the sensitive attribute is partially preserved or leaks into $\Psi_{\rm pred}(X)$, the relative entropy will be low [50]. Observe that this calculation is possible partly because the latent space specifies the same probability space for the two distributions. (b) The setting is identical in common approaches for learning disentangled representations, where disentanglement is measured via various information theoretic measures [8, 1, 21, 62]. If we now segue back to comparing two different networks, but without the convenience of a common coordinate system to measure divergence, the options turn out to be limited. (c) Recently, in trying to understand whether vision Transformers “see” similar to convolutional neural networks [57], one option utilized recently was a kernel-based representation similarity, in a layer-by-layer manner. What we may actually want is a mechanism for conditioning – for example, if one of the models is thought of a “nuisance variable”, we wish to check the residual in the other after the first has been controlled for (or marginalized out). Importantly, this should be possible without assuming that the probability distributions live in the same space (or networks $\Theta_{X}$ and $\Theta_{Y}$ are the same).

A direct application of CCA? Consider two different feature spaces $(\mathcal{X}$ and $\mathcal{Y})$, say in dimensions $\mathbb{R}^{p}$ and $\mathbb{R}^{q}$, pertaining to feature activations from two different models. Comparison of these two feature spaces is possible. One natural choice is canonical correlation analysis (CCA) [5], a generalization of correlation, specifically suited when $p\neq{q}$. The idea has been utilized for studying representation similarity in deep neural network models [49], albeit in a post-training setting for reasons that will be clear shortly, as well as for identifying more efficient training regimes (i.e., can lower layers be sequentially frozen after a certain number of timesteps). CCA has also been shown to be implementable within DNN pipelines for multi-view training, called DeepCCA [4], although efficiency can be a bottleneck limiting its broader deployment. A stochastic version of CCA suitable for DNN training with mini-batches has been proposed very recently, and strong experimental evidence was presented [48], also see [25]. Given that a stochastic CCA is now available, its extensions to the partial CCA setting are not yet available. If successful, this may eventually provide a scheme, suitable for deep learning, for controlling the influence of one model (or a set of variables) on another model.

This work. The starting point of this work is a less widely used statistical concept to measure the correlation between two different feature spaces $(\mathcal{X},\mathcal{Y})$ of different dimensions, called distance correlation (and the method of dissimilarities). In shallow settings, CCA and distance correlation offers very similar functionality – for the most part, they can be used interchangeably although distance correlation would also need specification of distances (or dissimilarities). In other words, CCA may be easier to deploy. On the other hand, deep variants of CCA involve specialized algorithms [4, 48]. Further, deep versions of partial CCA have not been reported. In contrast, as long as feature distances can be calculated, the differences between the shallow and the deep versions of distance correlation are minimal at best, and adjustments needed are quite minor. These advantages carry over to partial distance correlation, directly enabling conditioning one model w.r.t. another (or using such a term as a regularizer). The main contribution of this paper is to study distance correlation (and partial distance correlation) as a powerful measure in a broad suite of tasks in vision. We review the relevant technical steps which enable its instantiation in deep learning settings and show its broad applications ranging from learning disentangled representations to understanding the differences between what two (or more) networks are learning to training “mutually distinct” deep models (akin to earlier works on $M$ best solutions to MAP estimation in graphical models [19, 71]) or training $M$ diverse models for foreground-background segmentation as well as other tasks [27].

Benefits of Distance Correlation. In many applications, the notion of distances or dissimilarities appears quite naturally. Motivated by the need for a scheme that can capture both linear and non-linear correlations when provided with such dissimilarity information, in [66], the authors proposed a new measure of dependence between random vectors, called distance correlation. The key benefits of distance correlation are:

1. 1.

   The distance correlation $\mathcal{R}$ satisfies $0\leq\mathcal{R}\leq 1$, and $\mathcal{R}=0$ if and only if $X,Y$ are independent.

2. 2.

   $\mathcal{R}(X,Y)$ is defined for $X$ and $Y$ in arbitrary dimensions, e.g., $\mathcal{R}(X,Y)$ is well-defined when $X$ is of dimension $p$ while $Y$ is of dimension $q$ for $p\neq q$.

where $k,l\in\{1,\cdots,n\}$. Similarly, we can define $b_{k,l}=\|Y_{k}-Y_{l}\|$, and $B_{k,l}=b_{k,l}-\bar{b}_{k,\cdot}-\bar{b}_{\cdot,l}+\bar{b}_{\cdot,\cdot}$, and based on these quantities we have.

Examples. We show a few simple 2D examples to contrast Pearson Correlation and Distance Correlation in Fig. 1. Notice that if the relationship between the two random variables is not linear, Pearson Correlation might be small while Distance Correlation remains meaningful.

Extensions to conditioning. Given three random variables $X$, $Y$, and $Z$, we want to measure the correlation between $X$ and $Y$ but “controlling for” $Z$ (thinking of it as a nuisance variable), i.e., we want to estimate $\mathcal{R}(X|Z,Y|Z)=\mathcal{R}^{*}(X,Y;Z)$. Such a quantity is key in existing approaches in disentangled learning, deriving invariant representations and understanding what one or more networks are learning after concepts learned by another network have been accounted for. Consider how this task would be accomplished in linear regression. We would project $X$ and $Y$ into the space of $Z$, and only use the residuals to measure the correlation. Nonetheless, defining partial distance correlation is more involved – in [65], the authors introduced a new Hilbert space where we can define the projection of distance matrix. To do so, the authors calculate a $\mathcal{U}$-centered matrix $\tilde{A}$ from the distance matrix $(a_{k,l})$ so that the inner product of the $\mathcal{U}$-centered matrices will be the distance covariance.

Notation. We use $\Theta_{X},\Theta_{Y}$ to denote the parameters of the neural networks, and $X,Y$ as features extracted by the respective neural networks. Let the minibatch size be $m$, and the dataset $\mathcal{D}=(\mathcal{D_{X}},\mathcal{D_{Y}})$ be of size $n$. We use $(x_{t},y_{t})_{t=1}^{T},x_{t}\subset\mathcal{D_{X}},y_{t}\subset\mathcal{D_{Y}}$ to represent the data samples at step $t$, $T$ is the total number of training steps. The distance matrices $A_{t},B_{t}$ are computed when given $X_{t},Y_{t}$ using (A(1)), which is of dimension $m\times m$ for each minibatch. Further, we use $(X_{t})_{k}$ to represent the $k^{th}$ element in $X_{t}$. And $(A_{t})_{k,l}$ is the $k^{th}$ row and $l^{th}$ column element in the matrix $A_{t}$. The inner-product between two matrices $A,B$ is defined as $\langle A,B\rangle=\sum_{i,j}^{m}(A)_{i,j}(B)_{i,j}$.

Objective function. Consider the case where we minimize DC between two networks $\Theta_{X},\Theta_{Y}$. Since the parameters between $\Theta_{X},\Theta_{Y}$ are separable, we can use the block stochastic gradient iteration in [70] with some simple modifications.

We slightly abuse the notation of $\Theta_{X}(x)$ as applying the network $\Theta_{X}$ onto data $x$, and reuse $A$ to simplify the notation $A(\Theta_{X};x)$ and the distance matrix. We can rewrite the expression (with $A$, $B$ defined above) using:

where $(x,y)$ are the minibatch of samples from the data space $(\mathcal{D_{X}},\mathcal{D_{Y}})$.

where $f(\Theta_{X},\Theta_{Y};x,y)$ is $\langle A,B\rangle$ and $\gamma(\Theta_{X})$ encodes the convex constraint of network $\Theta_{X}$: $\max_{x\subset\mathcal{D_{X}}}\langle A,A\rangle\leq m$. Similarly, $\gamma(\Theta_{Y})$ encodes $\max_{y\subset\mathcal{D_{Y}}}\langle B,B\rangle\leq m$. $\Phi(\Theta_{X},\Theta_{Y})$ is the constrained objective function to be optimized.

Block stochastic gradient iteration. We adjust Alg. 1 from [70] to our case in Alg. 2. Since we will need the entire minibatch $(x_{t},y_{t})$ to compute the objective function, there will be no mean term when computing the sample gradient $\tilde{\mathbf{g}}_{X}^{t}$. Further, since both blocks $(\Theta_{X},\Theta_{Y})$ are constrained, line $3,5$ will use (5) from [70]. The detailed algorithm is presented in Alg. 2.

Rationale. Recently, several efforts have explored generating of adversarial examples that can transfer to different networks and how to defend against such attacks [15, 63, 6]. It is often observed that an adversarial sample for one trained network is relatively easy to transfer to another network with the same architecture [15]. Here, we show that even for as few as two networks (same architecture; trained on the same data), we can, to some extent, prevent adversarial examples from transferring between them by seeking independent features.

Setup. We formulate the problem considering a classification task as an example. Given two deep neural networks with the same architecture denoted as $f_{1}(\cdot),f_{2}(\cdot)$, we train them using image-label pairs $(x,y)$ using the cross-entropy loss $\text{Loss}_{\text{CE}}$. If we train $f_{1}$ and $f_{2}$ using only the cross-entropy loss, the adversarial examples generated on $f_{1}$ can relatively easily transfer to $f_{2}$ (see the performance of “Baseline” in Table 1). To enforce $f_{1}$ and $f_{2}$ to learn independent features, let the extracted feature of $x$ in some intermediate layer of $f$ be given as $g(x)$ (in this section we use the feature before the last fully connected layer as an example). We can still train $f_{1}$ using $\text{Loss}_{\text{CE}}$, and then, we train $f_{2}$ using,

where $\alpha$ is a constant scalar and Loss${}_{\text{DC}}$ is the distance correlation from Def. 1. Note that we do not require $g_{1}(x)$ and $g_{2}(x)$ to be in the same dimension, so in principle we could easily use features from different layers for these two networks.

Experimental settings. We first conduct experiments on CIFAR10 [39] using Resnet 18 [28]. We then use four different architectures (mobilenet-v3-small [32], efficientnet-B0 [67], Resnet 34, and Resnet152) and train them on ImageNet [40]. For each network architecture, we first train two networks using only $\text{Loss}_{\text{CE}}$. Next, we train a network using only $\text{Loss}_{\text{CE}}$ before training a second network using the loss in (13). On CIFAR10, we utilize the SGD optimizer with momentum $0.9$ and train for $200$ epochs using an initial learning rate $0.1$ with a cosine learning rate scheduler [53]. The mini-batch size is set to $128$. On ImageNet [40], we train for $40$ epochs using an initial learning rate $0.1$, which decays by $0.1$ every 10 epochs. The mini-batch size is $512$. Our $\alpha$ in (13) is set to $0.05$ for all cases. For each combination of the dataset and the network architecture, we train two networks $f_{1}$ and $f_{2}$, after which we generate adversarial examples on $f_{1}$ and use them to attack $f_{2}$ and measure its classification accuracy. We construct a baseline by training $f_{1}$ and ${f_{2}}_{\rm Baseline}$ without constraints. And train ${f_{2}}_{\rm Our}$ using (13) to learn independent features w.r.t. $f_{1}$. We report performance under two widely used attack methods: fast gradient sign method (FGM) [26] and projected gradient descent method (PGD) [47], where the latter is considered among the strongest attacks. The scale $\epsilon$ of the adversarial perturbation is chosen from $\{0.03,0.05,0.1\}$ and the maximum number of iterations of PGD is set to $40$.

Results. The results are shown in Table 1. We see that we get significant improvement in accuracy over the baseline under adversarial attacks, with comparable performance on clean inputs. Notably, our method achieves more than $10\%$ absolute improvement in accuracy under PGD attack on Resnet-18 and Mobilenet-v3-small. This provides evidence supporting the benefits of enforcing the networks to learn independent features using our distance correlation loss.

Overview. This experiment studies disentanglement [30, 36, 8, 44, 21]. It is believed that the image data are generated from low dimensional latent variables – but isolating and disentangling the latent variables is challenging. A key in disentangled latent variable learning is to make the factors in the latent variables independent [2]. Distance correlation fits perfectly and can handle a variety of dimensions for the latent variables. When the distance correlation is $0$, we know that the two variables are independent.

Experimental settings. We follow [21] which focuses on semi-supervised disentanglement to generate high-resolution images. In [21], one divides the latent variables into two categories: (a) attributes of interest – a set of semantic and interpretable attributes, e.g. hair color and age; (b) residual attributes – the remaining information. Formally, $x_{i}=G(f_{i}^{1},...,f_{i}^{k},r_{i})$, where $G$ is the generator that uses the factors of interest $f_{i}^{l}$ and the residual to generate image $x_{i}$.

Results. (Shown in Fig. 5) Our model shows the ability to change specific attributes without affecting residual features, such as posture (also see supplement). If you use distance correlation for disentanglement, please give credit to the following paper [43] which discusses a nice demonstration of distance correlation helps content style disentanglement. We were not aware of this paper when we wrote the paper last year and thank Sotirios Tsaftaris for communicating his findings with us.

Acknowledgements. Research supported in part by NIH grants RF1 AG059312, RF1 AG062336, and RF1AG059869, and NSF grant CCF #1918211.
The authors thank Grace Wahba for introducing us to the Gabor Szekely’s work on distance correlation in 2014.