Tensor Dropout for Robust Learning

TL;DR

This paper introduces tensor dropout, a method that enhances robustness and efficiency of CNNs by imposing low-rank tensor structures, demonstrating superior performance on image classification and brain MRI datasets.

Contribution

The paper proposes tensor dropout with low-rank tensor structures, providing a novel approach for robust, compact, and efficient neural networks with theoretical validation.

Findings

Outperforms existing methods on ImageNet and CIFAR-100

Achieves state-of-the-art accuracy on UK Biobank brain MRI dataset

Significantly improves robustness to noise and adversarial attacks

Abstract

CNNs achieve remarkable performance by leveraging deep, over-parametrized architectures, trained on large datasets. However, they have limited generalization ability to data outside the training domain, and a lack of robustness to noise and adversarial attacks. By building better inductive biases, we can improve robustness and also obtain smaller networks that are more memory and computationally efficient. While standard CNNs use matrix computations, we study tensor layers that involve higher-order computations and provide better inductive bias. Specifically, we impose low-rank tensor structures on the weights of tensor regression layers to obtain compact networks, and propose tensor dropout, a randomization in the tensor rank for robustness. We show that our approach outperforms other methods for large-scale image classification on ImageNet and CIFAR-100. We establish a new state-of-the-art accuracy for phenotypic trait prediction on the largest dataset of brain MRI, the UK Biobank brain MRI dataset, where multi-linear structure is paramount. In all cases, we demonstrate superior performance and significantly improved robustness, both to noisy inputs and to adversarial attacks. We rigorously validate the theoretical validity of our approach by establishing the link between our randomized decomposition and non-linear dropout.

Tables

Table 1. TABLE I : Classification accuracy on ImageNet with Bernoulli Tucker R-TRL. TRL ( θ 𝜃 \theta = 1) is the corresponding deterministic architecture.

		Accuracy
Model	$\theta$	Top-1 (%)	Top-5 (%)
ResNet	✗	77.1	93.4
TRL [16]	1.0	77.1	93.5
Ours	0.9	77.7	93.7
Ours	0.8	77.7	93.7
Ours	0.7	77.4	93.6
Ours	0.6	78.0	93.8

Table 2. TABLE II: Real-valued network performance on ImageNet for FGSM, BIM and PGD attacks with  λ ∈ { 2 , 8 , 16 } 𝜆 2 8 16 \lambda\in\{2,8,16\} . We report classification accuracy in all cases.

Attack		Method
Type	$\lambda$	Baseline	Ours
			$\theta =0.8$	$\theta =0.7$	$\theta =0.6$
Clean		$77.1$	$77.7$	$77.4$	$78.0$
(no attack)
FGSM	2	$26.1$	$26.0$	$35.3$	$47.4$
	8	$11.1$	$15.4$	$26.0$	$41.6$
	16	$8.5$	$14.1$	$24.0$	$39.5$
BIM	2	$26.1$	$26.0$	$35.3$	$47.3$
	8	$1.0$	$4.0$	$9.9$	$26.1$
	16	$0.1$	$1.0$	$2.8$	$13.2$
PGD	2	$26.0$	$26.0$	$35.6$	$47.1$
	8	$0.8$	$4.5$	$11.3$	$27.9$
	16	$0$	$1.4$	$3.8$	$13.5$

Table 3. TABLE III : Classification accuracy for UK Biobank MRI. The ResNet models with R-TRL significantly outperforms the version with a fully-connected (FC) layer.

Architecture	Regression	MAE
3D-ResNet	FC	2.96 years
3D-ResNet	Tucker	2.70 years
Ours	Randomized Tucker	2.65 years
Ours	Randomized CP	2.58 years

Table 4. TABLE IV : Classification accuracy for CIFAR-100 with a ResNet and various regression layers for classification.

ResNet classification	Accuracy
FC	75.88 %
FC + dropout	75.84 %
Tucker	76.02 %
CP	75.77 %
Randomized Tucker	76.05 %
Randomized CP	76.19 %

Equations

$$j=\sum _{\begin{array}{c}k=0,\\ k\ne n\end{array}}^N{i}_k\times \prod _{\begin{array}{c}m=k+1,\\ m\ne n\end{array}}^N{I}_m.$$

$$j=\sum _{k=0}^N{i}_k\times \prod _{m=k+1}^N{I}_m.$$

$${\left(\mathcal{X}{\times }_n\mathbf{M}\right)}_{[n]}=\mathbf{M}{\mathcal{X}}_{[n]}\in {\mathbb{R}}^{I_0\times \cdots \times I_{n-1}\times R\times I_{n+1}\times \cdot \times I_N}$$

$$⟨\mathcal{X},\mathcal{Y}{⟩}_N=\sum _{i_0=0}^{I_0}\sum _{i_1=0}^{I_1}\cdots \sum _{i_n=0}^{I_N}{\mathcal{X}}_{i_0,i_1,\cdots ,i_n}{\mathcal{Y}}_{i_0,i_1,\cdots ,i_n}$$

$$\mathcal{X}=\sum _{k=0}^{R-1}\underset{\text{rank-1 components}}{\underbrace{{\lambda }_k{\mathbf{u}}_k^{(0)}\circ {\mathbf{u}}_k^{(1)}\circ \cdots \circ {\mathbf{u}}_k^{(N)}}},$$

$$\mathcal{X}$$

$$=[[\mathcal{G};{\mathbf{U}}^{(0)},\cdots ,{\mathbf{U}}^{(N)}]]$$

$$y^{(k)}$$

$$\mathcal{W}=\mathcal{G}{\times }_0{\mathbf{U}}^{(0)}{\times }_1{\mathbf{U}}^{(1)}\cdots {\times }_N{\mathbf{U}}^{(N)}$$

$$\frac{1}{S-1}\sum _{k=0}^{S-1}{\left({\mathbf{y}}^{(k)}-⟨\tilde{\mathcal{W}},{\mathcal{X}}^{(k)}⟩\right)}^2,$$

$$\tilde{\mathcal{W}}=\tilde{\mathcal{G}}{\times }_0{\tilde{\mathbf{U}}}^{(0)}\times \cdots {\times }_N{\tilde{\mathbf{U}}}^{(N)}$$

$$\begin{array}{rl}\tilde{\mathcal{W}}=& \mathcal{G}{\times }_0{\mathbf{M}}^{(0)}\times \cdots {\times }_N{\mathbf{M}}^{(N)}\\ & {\times }_0\left({\mathbf{U}}^{(0)}({\mathbf{M}}^{(0)}{)}^{\top }\right)\times \cdots {\times }_N\left({\mathbf{U}}^{(N)}({\mathbf{M}}^{(N)}{)}^{\top }\right)\\ & =[[\tilde{\mathcal{G}};{\tilde{\mathbf{U}}}^{(0)},\cdots ,{\tilde{\mathbf{U}}}^{(N)}]]\end{array}$$

$$\begin{split}&\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}\\ &=\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta^{N}\langle{\mathcal{W}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\\ &+\tiny{\frac{\theta^{N}(1-\theta^{N})}{S-1}\sum^{S-1}_{k=0}}\langle{\mathcal{G}^{\star 2}\times_{0}(\mathbf{U}^{(0)})^{\star 2}\dots\times_{N}(\mathbf{U}^{(N)})^{\star 2}},{(\mathcal{X}^{(k)})^{\star 2}}\rangle.\end{split}$$

$$\begin{array}{rl}\tilde{\mathcal{W}}& =\sum _{k=0}^{R-1}{\tilde{\mathbf{U}}}_k^{(0)}\circ \cdots \circ {\tilde{\mathbf{U}}}_k^{(N)}\end{array}$$

$$\begin{array}{rl}\tilde{\mathcal{W}}=& [[\mathbf{\lambda };{\mathbf{U}}^{(0)},\cdots ,{\mathbf{U}}^{(N)}]]\end{array}$$

$$\underset{{\mathbf{U}}^{(0)},\cdots ,{\mathbf{U}}^{(N)}}{\min }{\left({\mathbf{y}}^{(k)}-⟨[[\mathbf{\lambda };{\mathbf{U}}^{(0)},\cdots ,{\mathbf{U}}^{(N)}]],{\mathcal{X}}^{(k)}⟩\right)}^2$$

$$\begin{split}&\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\llbracket{}{\mathbf{\bm{\lambda}};\,\mathbf{U}^{(0)},\cdots,\mathbf{U}^{(N)}}\rrbracket{}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}\\ &=\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta\langle{\llbracket{}{\mathbf{U}^{(0)},\cdots,\mathbf{U}^{(N)}}\rrbracket{}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\\ &+\tiny{\frac{\theta(1-\theta)}{S-1}\sum^{S-1}_{k=0}}\langle{\llbracket{}{(\mathbf{U}^{(0)})^{\star 2},\cdots,(\mathbf{U}^{(N)})^{\star 2}}\rrbracket{}},{(\mathcal{X}^{(k)})^{\star 2}}\rangle.\end{split}$$

$${\mathbf{M}}^{(i)}(j,:)=\{\begin{array}{ll}0& \text{mbox}ifj>K\\ {\mathbf{Id}}_m(r,:),m\in [0..R_i]& \text{mbox}otherwise\end{array}$$

$$\begin{array}{r}{\tilde{\mathcal{W}}}_{i_0,\dots ,i_N}=\sum _{r_0=1}^{R_0}\cdots \sum _{r_N=1}^{R_N}{\mathbf{g}}_{r_0,\dots ,r_N}{\lambda }_{r_0}^{(0)}\dots {\lambda }_{r_N}^{(N)}{\mathbf{U}}_{i_0{r}_0}^{(0)}\dots {\mathbf{U}}_{i_N{r}_N}^{(N)}\end{array}$$

$$\begin{split}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}&=\frac{1}{S-1}\sum^{S-1}_{k=0}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}\\ &=\underbrace{\frac{1}{S-1}\sum^{S-1}_{k=0}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)}\,\bm{\big{]}}^{2}}_{\textrm{E}}+\underbrace{\frac{1}{S-1}\sum^{S-1}_{k=0}\operatorname{Var}\left[{\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle}\right]}_{\textrm{V}}\end{split}$$

$$\begin{split}\textrm{E}&=\frac{1}{S-1}\sum^{S-1}_{k=0}{\left(\mathbf{\bm{y}}^{(k)}-\sum_{i_{0}}\cdots\sum_{i_{N}}\mathcal{X}_{i_{0},\dots,i_{N}}^{(k)}{\sum_{r_{0}}\cdots\sum_{r_{N}}\mathbf{\bm{g}}_{r_{0},\dots,r_{N}}\mathbb{E}_{\,}\bm{\big{[}}\,{\lambda_{r_{0}}^{(0)}\dots\lambda_{r_{N}}^{(N)}}\,\bm{\big{]}}}\mathbf{U}_{i_{0}r_{0}}^{(0)}\dots\mathbf{U}_{i_{N}r_{N}}^{(N)}\right)}^{2}\\ &=\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta^{N}\langle{\mathcal{W}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\text{ as }\lambda^{(0)},\dots,\lambda^{(N)}\text{ are i.i.d. }\end{split}$$

$$\begin{array}{rl}\text{V}& =\frac{1}{S-1}\sum _{k=0}^{S-1}\left({(-1)}^2\sum _{i_0}\cdots \sum _{i_N}{({\mathcal{X}}_{i_0,\dots ,i_N}^{(k)})}^2\sum _{r_0}\cdots \sum _{r_N}{\mathbf{g}}_{r_0,\dots ,r_N}^2\mathrm{Var}\left[{\lambda }_{r_0}^{(0)}\dots {\lambda }_{r_N}^{(N)}\right])({\mathbf{U}}_{i_0{r}_0}^{(0)}{)}^2\dots ({\mathbf{U}}_{i_N{r}_N}^{(N)}{)}^2\right)\\ & =\frac{1}{S-1}{\theta }^N(1-{\theta }^N)\sum _{k=0}^{S-1}\left({\sum _{i_0}\cdots \sum _{i_N}({\mathcal{X}}_{i_0,\dots ,i_N}^{(k)})}^2{\mathcal{G}}^{\star 2}{\times }_0({\mathbf{U}}^{(0)}{)}^{\star 2}\cdots {\times }_N({\mathbf{U}}^{(N)}{)}^{\star 2}\right)\\ & =\frac{{\theta }^N(1-{\theta }^N)}{S-1}\sum _{k=0}^{S-1}⟨{\mathcal{G}}^{\star 2}{\times }_0({\mathbf{U}}^{(0)}{)}^{\star 2}\cdots {\times }_N({\mathbf{U}}^{(N)}{)}^{\star 2},({\mathcal{X}}^{(k)}{)}^{\star 2}⟩)\end{array}$$

$$\begin{split}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\,\bm{\big{]}}=\,\,&\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta^{N}\langle{\mathcal{W}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\\ &+\frac{\theta^{N}(1-\theta^{N})}{S-1}\sum^{S-1}_{k=0}\langle{\mathcal{G}^{\star 2}\times_{0}(\mathbf{U}^{(0)})^{\star 2}\dots\times_{N}(\mathbf{U}^{(N)})^{\star 2}},{(\mathcal{X}^{(k)})^{\star 2}}\rangle.\end{split}$$

$$\begin{array}{rl}\tilde{\mathcal{W}}& =\sum _{k=0}^{R-1}{\tilde{\mathbf{U}}}_k^{(0)}\circ \cdots \circ {\tilde{\mathbf{U}}}_k^{(N)}\\ & =\sum _{k=0}^{R-1}{\lambda }_k{\mathbf{U}}_k^{(0)}\circ \cdots \circ {\mathbf{U}}_N^{(0)}\\ & =[[\mathbf{\lambda };{\mathbf{U}}^{(0)},\cdots ,{\mathbf{U}}^{(N)}]]\end{array}$$

$${\tilde{\mathbf{U}}}_{i_k,r}^{(k)}=\sum _{j=0}^{R-1}{\mathbf{U}}_{i_k,j}^{(k)}{\mathbf{M}}_{j,r}.$$

$${\tilde{\mathbf{U}}}_{i_k,r}^{(k)}={\lambda }_r{\mathbf{U}}_{i_k,r}^{(k)}.$$

$${\tilde{\mathcal{W}}}_{i_0,i_1,\cdots ,i_N}=\sum _{r=0}^{R-1}{\lambda }_k\left({\mathbf{U}}_k^{(0)}\circ \cdots \circ {\mathbf{U}}_k^{(N)}\right)$$

$$\begin{split}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}&=\frac{1}{S-1}\sum^{S-1}_{k=0}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)^{2}}\,\bm{\big{]}}\\ &=\underbrace{\frac{1}{S-1}\sum^{S-1}_{k=0}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,{\left(\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle\right)}\,\bm{\big{]}}^{2}}_{\textrm{E}}+\underbrace{\frac{1}{S-1}\sum^{S-1}_{k=0}\operatorname{Var}\left[{\mathbf{\bm{y}}^{(k)}-\langle{\mathcal{\tilde{W}}},{\mathcal{X}^{(k)}}\rangle}\right]}_{\textrm{V}}\end{split}$$

$$\begin{split}\textrm{E}&=\frac{1}{S-1}\sum^{S-1}_{k=0}{\left(\mathbf{\bm{y}}^{(k)}-\sum_{i_{0}=1}^{I_{0}}\cdots\sum_{i_{N}=1}^{I_{N}}{\sum_{r=0}^{R-1}\mathbb{E}_{\,}\bm{\big{[}}\,{\lambda_{n}}\,\bm{\big{]}}\mathbf{U}_{i_{0}r_{0}}^{(0)}\dots\mathbf{U}_{i_{N}r_{N}}^{(N)}}(\mathcal{X}^{(k)})^{\star 2}\right)}^{2}\\ &=\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta^{N}\langle{\mathcal{W}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\end{split}$$

$$\begin{array}{rl}\text{V}& =\frac{1}{S-1}\sum _{k=0}^{S-1}{(-1)}^2\sum _{i_0}\cdots \sum _{i_N}{({\mathcal{X}}_{i_0,\dots ,i_N}^{(k)})}^2\sum _{r=0}^{R-1}\mathrm{Var}\left[\theta \right]({\mathbf{U}}_{i_0,r}^{(0)}{)}^2\dots ({\mathbf{U}}_{i_{N}r}^{(N)}{)}^2\\ & =\frac{\theta (1-\theta )}{S-1}\sum _{k=0}^{S-1}⟨({\mathbf{U}}^{(0)}{)}^{\star 2},\dots ,{\mathbf{U}}^{(N)}{)}^{\star 2},({\mathcal{X}}^{(k)}{)}^{\star 2}⟩\end{array}$$

$$\begin{split}\mathbb{E}_{\,\mathbf{\bm{\lambda}}}\bm{\big{[}}\,\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\langle{\llbracket{}{\mathbf{\bm{\lambda}};\,\mathbf{U}^{(0)},\cdots,\mathbf{U}^{(N)}}\rrbracket{}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\,\bm{\big{]}}=&\frac{1}{S-1}\sum^{S-1}_{k=0}\left(\mathbf{\bm{y}}^{(k)}-\theta\langle{\llbracket{}{\mathbf{U}^{(0)},\cdots,\mathbf{U}^{(N)}}\rrbracket{}},{\mathcal{X}^{(k)}}\rangle\right)^{2}\\ &+\frac{\theta(1-\theta)}{S-1}\sum^{S-1}_{k=0}\langle{\llbracket{}{(\mathbf{U}^{(0)})^{\star 2},\cdots,(\mathbf{U}^{(N)})^{\star 2}}\rrbracket{}},{(\mathcal{X}^{(k)})^{\star 2}}\rangle\end{split}$$

Full text

I Introduction

Deep neural networks have evolved into powerful predictive models with remarkable performance on computer vision tasks [1, 2, 3]. Such models are usually over-parameterized, involving an enormous number (typically tens of millions) of parameters. This is much larger than the typical number of available training samples, making deep networks prone to overfitting [4]. Coupled with overfitting, deep networks lack robustness to small adversarial or noise perturbations [5]. In computer vision tasks, such as image classification or regression, perturbed examples are often perceived identically to the original ones by humans while lead arbitrarily different predictions by networks. In addition, not only are these neural networks lacking in robustness, they are also over-confident when (wrongly) predicting on noisy inputs [6]. These shortcomings pose a possible obstacle for mass deployment of systems relying on deep learning in sensitive fields such as medical image analysis for disease prediction and expose an inherent weakness in their reliability.

To improve the robustness of deep neural networks to (adversarial or random) perturbations and prevent them from overfitting, several methods that essentially affect the parameters of deep networks via regularization have been investigated. In particular, the link between regularization and robustness of deep neural networks to adversarial perturbations have been recently established in [7, 8]. These regularization mostly fall into two categories: structural changes to the architecture to make them inherently more robust and regularization methods that constrain directly the parameters. The latter category is well studied in the context of deep learning. In this context, regularization can be applied in the form of added randomness during training and/or testing. This can be applied to the activations through dropout [9] or to the weights (DropConnect [10]). Regularization functions (e.g., $\ell_{2}$- or $\ell_{1}$- norm) can also be directly applied to network’s parameters [11, 12, 13, 14].

However, a more powerful class of regularization focuses on structural changes to the architecture of the network itself. The most notable such approach is based on tensor methods and accounts for the structure in the data. Machine learning is in essence the learning from data, and tensor methods allow to fully leverage the structure in that data. This is because the unknown transformation from input image to the output label is in general a tensor map. The data typically manipulated by deep neural networks is typically 3 dimensional (e.g. RGB images, brain MRI) or 4 dimensional (e.g. videos, functional MRI images, etc). Preserving multi-dimensional structure is crucial for performance. By limiting the network to matrix operations (e.g. with flattening layers followed by one or more fully-connected), we are ignoring this structure, resulting in deteriorated performance [15, 16, 17]. In contrast, tensor methods allow to leverage that structure to improve the model, reduce the number of parameters and improve computational efficiency. One way this is done is by leveraging the multi-linear correlations in the network [18, 19, 20, 16, 21].

In this paper, we introduce a novel higher-order randomized factorization method to improve the performance and robustness of deep networks to adversarial and random perturbations. We apply this stochastic decomposition to tensor regression layers which we use to entirely replace flattening and fully-connected layers. This preserves the structure of multi-dimensional data (e.g., images and MRI data) by expressing an output tensor as the result of a tensor contraction between the input tensor and some low-rank regression weight tensors. Instead of using the deterministic reconstruction of the regression weight tensor, our randomized tensor decomposition leads to the stochastic reduction of the rank during both training and inference. This is related to non-linear dropout on tensor factorization, which, by randomly dropping units during training, prevents over-fitting. However, rather than dropping random elements from the activation tensor, this is done on the rank of the regression weight tensor.

We conduct thorough experiments in image classification and phenotypic trait prediction from MRI analysis. Specifically, we apply our method to the UK Biobank brain MRI dataset for estimating brain-age, a metric which has been associated with a range of diseases and mortality [22], and could be an early predictor for Alzheimer’s disease [23]. A more accurate and robust metric of brain condition can consequently lead to more accurate disease diagnoses.

Summary of contributions:

• •

We propose tensor dropout, a novel method, a novel stochastic tensor decomposition where non-linear dropout is applied in the latent subspace spanned by a low-rank factorization.

• •

We apply tensor dropout to tensor regression layers and show that it improves the inductive bias of convolutional neural networks (CNNs) by fully leveraging the structure in the data via a stochastic tensor decomposition.

• •

We demonstrate superior performance against both regular deep network architectures with fully-connected layers and networks with tensor regression layers that do not incorporate our proposed randomized decomposition.

• •

We establish a new state-of-the-art for large scale regression from MRI data on the UK Biobank dataset, the largest dataset of brain MRI. We demonstrate that our model is significantly more robust to noise in the input, as occurs naturally during capture.

• •

Our method makes neural networks significantly more robust to adversarial noise, without adversarial training.

• •

We show that our proposed method implicitly regularizes the tensor decomposition. We establish theoretically and empirically the link between tensor dropout and the deterministic low-rank tensor regression.

• •

We validate all aspects of our method on the CIFAR-100 image classification dataset in thorough ablation studies.

II Related work

In this section, we review the most closely related work to ours.

Network regularization and dropout. Several methods that improve generalization by mitigating overfitting have been developed in the context of deep learning. The interested reader is referred to the work of [24] and the references therein for a comprehensive survey of regularization techniques for deep networks. The most closely related regularization method to our approach is Dropout [9], which is probably the most widely adopted technique for training neural networks while preventing overfitting. Concretely, during training, each unit (i.e., neuron) is equipped with a binary Bernoulli random variable and only the network’s weights whose corresponding Bernoulli variables are sampled with value $1$ are updated at each back-propagation step. At each iteration, those Bernoulli variables are re-sampled again and the weights are updated accordingly. Our proposed regularization method can be interpreted as dropout on low-rank tensor regression, a fact which we proved in Section (IV).

Randomized tensor decompositions. Tensor decompositions exhibit high computational cost and low convergence rate when applied to massive multi-dimensional data. To accelerate computation, and enable them to scale, randomized approaches have been employed. A randomized least squares algorithm for CP decomposition is proposed by [25], which is significantly faster than traditional one. In [26], CP is applied on a small tensor generated by multi-linear random projection of the high-dimensional tensor. The CP decomposition of the large-scale tensor is obtained by back projection of the CP decomposition of the small tensor. [27] introduce a fast yet provable randomized CP decomposition that performs randomized tensor contraction using FFT. Methods in [28, 29] are highly computationally efficient algorithms for computing large-scale CP decompositions by applying random projections into a set of small tensors, derived by subdividing a tensor into a set of blocks. Fast randomized algorithms that approximate Tucker decomposition using Sketching have also been investigated [30, 31]. More recently, a randomized tensor ring decomposition that employs tensor random projections has been developed in [32]. The most similar method to ours is that of [25], where elements of the tensor are sampled randomly, and each factor of the decomposition updated in an iterative manner. By contrast, our method allows for end-to-end training, and applies randomization on the fibers of the tensor, effectively randomizing the rank of the weight tensor.

III Preliminaries and tensor regression layers

In this section, we introduce the notations and preliminary notions necessary to introduce our stochastic rank regularization, including tensor regression layers.

Notation: We denote $\mathbf{\bm{v}}$ to be vectors (1${}^{\text{st}}$order tensors) and $\mathbf{M}$ to be matrices (2${}^{\text{nd}}$order tensors). $\mathbf{Id}$ is the identity matrix. We denote as $\mathcal{X}$ tensors of order $N\geq 3$, and denote its element $(i,j,k)$ as $\mathcal{X}_{i_{0},i_{1},\cdots,i_{N-}}\,\text{ or }\,\mathcal{X}(i_{0},i_{1},\cdots,i_{N-})$. A colon is used to denote all elements of a mode e.g. the mode-0 fibers of $\mathcal{X}$ are denoted as $\mathcal{X}_{\,\bm{:}\,,i_{1},\cdots,i_{N}}$. The transpose of $\mathbf{M}$ is denoted $\mathbf{M}^{\top}$. For any tensor $\mathcal{X}$, we denote $\mathcal{X}^{*2}$ the Hadamard power of 2, which squares each element of the tensor. Finally, for any $i,j\in\mathbb{N},i<j,\,[{i}\mathrel{{.}\,{.}}\nobreak{j}]$ denotes the set of integers $\{i,i+1,\cdots,j-1,j\}$, and $i\nonscript\mskip-4.0mu plus -2.0mu minus -4.0mu\mkern 5.0mu\mathbin{\operator@font div}\penalty 900\mkern 5.0mu\nonscript\mskip-4.0mu plus -2.0mu minus -4.0muj$ the integer division of $i$ by $j$.

Tensor unfolding: Given a tensor, $\mathcal{X}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}}$, its mode-$n$ unfolding is a matrix $\mathbf{X}_{[n]}\in\mathbb{R}^{I_{n},I_{M}}$, with $M=\prod_{\begin{subarray}{c}k=0,\\ k\neq n\end{subarray}}^{N}I_{k}$, defined by the mapping from $(i_{0},i_{1},\cdots,i_{N})$ to $(i_{n},j)$, with

[TABLE]

Tensor vectorization: Given a tensor, $\mathcal{X}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}}$, we can flatten it into a vector $\text{vec}(\mathcal{X})$ of size $\left(I_{0}\times\cdots\times I_{N}\right)$ defined by the mapping from element $(i_{0},i_{1},\cdots,i_{N})$ of $\mathcal{X}$ to element $j$ of $\text{vec}(\mathcal{X})$, with

[TABLE]

Mode-n product: For a tensor $\mathcal{X}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}}$ and a matrix $\mathbf{M}\in\mathbb{R}^{J\times I_{n}}$, the n-mode product of a tensor is a tensor of size $\left(I_{0}\times\cdots\times I_{n-1}\times J\times I_{n+1}\times\cdots\times I_{N}\right)$ and can be expressed using unfolding of $\mathcal{X}$ and the classical dot product as:

[TABLE]

Inner product: For two tensors of matching sizes, $\mathcal{X},\mathcal{Y}\in\mathbb{R}^{I_{0}\times\cdots\times I_{N}}$, we denote by $\langle{\mathcal{X}},{\mathcal{Y}}\rangle\in\mathbb{R}$ the contraction of $\mathcal{X}$ by $\mathcal{Y}$ along each mode.

[TABLE]

Kruskal tensor: Given a tensor $\mathcal{X}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N-1}}$, the Canonical-Polyadic decomposition (CP), also called PARAFAC, decomposes it into a sum of $R$ rank-1 tensors. The number of terms in the sum, $R$, is known as the rank of the decomposition. Formally, we find the vectors $\mathbf{u}^{(0)}_{k},\mathbf{u}^{(1)}_{k},\cdots,\mathbf{u}^{(N)}_{k}$, for $k=[{0}\mathrel{{.}\,{.}}\nobreak{R-1}]$, as well as an optional vector of weights $\mathbf{\bm{\lambda}}\in\mathbb{R}^{R}$ to absorbs the magnitude of the factors such that:

[TABLE]

For any $k\in[{0}\mathrel{{.}\,{.}}\nobreak{{N}}],$ these vectors $\mathbf{u}^{(k)}_{r},r\in[{0}\mathrel{{.}\,{.}}\nobreak{R-1}]$ can be collected in matrices, called factors or the decomposition, and defined as $\mathbf{U}^{(k)}=\left[\begin{matrix}\mathbf{u}^{(k)}_{0},\mathbf{u}^{(k)}_{1},\cdots,\mathbf{u}^{(k)}_{R-1}\end{matrix}\right].$ The decomposition can be denoted more compactly as $\mathcal{X}=\llbracket{}{\mathbf{\bm{\lambda}};\,\mathbf{U}^{(0)},\cdots,\mathbf{U}^{({N})}}\rrbracket{}$.

Tucker tensor: Given a tensor $\mathcal{X}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}}$, we can decompose it into a low rank core $\mathcal{G}\in\mathbb{R}^{R_{0}\times R_{1}\times\cdots\times R_{N-}}$ by projecting along each of its modes with projection factors $\left(\mathbf{U}^{(0)},\cdots,\mathbf{U}^{({N})}\right)$, with $\mathbf{U}^{(k)}\in\mathbb{R}^{R_{k},I_{k}},k\in(0,\cdots,{N})$. The tensor in its decomposed form can then be written:

[TABLE]

Note that the Kruskal form of a tensor can be seen as a Tucker tensor with a super-diagonal core.

Tensor diagrams: In order to represent easily tensor operations, we adopt the tensor diagrams, where tensors are represented by vertices (circles) and edges represent their modes. The degree of a vertex then represents its order. Connecting two edges symbolizes a tensor contraction between the two represented modes. Figure 1 presents a tensor diagram of the tensor regression layer and its stochastic rank-regularized counter-part.

Tensor regression layers: Reducing the storage and computational costs of deep networks has become critical for meeting the requirements of environments with limited memory or computational resources. To this end, a surge of network compression and approximation algorithms have recently been proposed in the context of deep learning. By leveraging the redundancy in network parameters, methods such as [18, 19, 20, 16] employ low-rank approximations of deep networks’ weight matrices (or tensors) for parameter reduction. Network compression methods in the frequency domain [33] have also been investigated. An alternative approach for reducing the number of effective parameters in deep nets relies on sketching, whereby, given a matrix or tensor of input data or parameters, one first compresses it to a much smaller matrix (or tensor) by multiplying it by a (usually) random matrix with certain properties [34, 35].

A particularly appealing approach to network compression, especially for visual data (and other types of multidimensional and multi-aspect data), is tensor regression networks [16]. Deep neural networks typically leverage the spatial structure of input data via series of convolutions, point-wise non-linearities, pooling, etc. However, this structure is usually wasted by the addition, at the end of the networks’ architectures, of a flattening layer followed by one or several fully-connected layers. A recent line of study focuses on alleviating this using tensor methods. [15] proposed tensor contraction as a layer, to reduce the size of activation tensors, and demonstrated large space savings with this layer. Recently, tensor regression networks [16] propose to replace fully-connected layers entirely with a tensor regression layer (TRL). This tensor regression layers preserves the structure by expressing an output tensor as the result of a tensor contraction between the input tensor and some low-rank regression weight tensors.

More specifically, we assume we have a set of $S$ input activation tensors $\mathcal{X}^{(k)}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}}$ and corresponding target labels $y^{(k)}\in\mathbb{R}$, with $k\in[{0}\mathrel{{.}\,{.}}\nobreak{N}]$. Not that here, for clarity, we introduce the tensor regression layer for scalar targets. In practice the target is typically a vector or a tensor. A tensor regression layer estimates the regression weight tensor $\mathcal{W}\in\mathbb{R}^{I_{0}\times I_{1}\times\cdots\times I_{N}},$ expressed under some low-rank decomposition. For instance, using a rank–$\left(R_{0},\cdots,R_{N}\right)$ Tucker decomposition, we have:

[TABLE]

with $\mathcal{G}\in\mathbb{R}^{R_{0}\times\cdots\times R_{N}}$, $\mathbf{U}^{(k)}\in\mathbb{R}^{I_{k}\times R_{k}}$ for each $k$ in $[{0}\mathrel{{.}\,{.}}\nobreak{N}]$ and $\mathbf{U}^{(N)}\in\mathbb{R}^{O\times R_{N}}$.

In addition to preserving and leveraging the structure in the input, TRL allow for large space savings without sacrificing accuracy. [36] explore the same model with various low-rank structures on the regression weights.

IV Tensor dropout

In this section, we introduce our tensor dropout and apply it to tensor regression layers. Specifically, we propose a new stochastic rank-regularization, applied to low-rank tensors in decomposed forms. This formulation is general and can be applied to any type of decomposition. We introduce it here, without loss of generality, to the case of Tucker and CP decompositions.

Randomized tensor regression layer We propose a novel randomized decomposition on $\mathcal{W}$, which applies dropout in the latent subspace spanned by a tensor decomposition. For instance, for an $N^{\text{th}}$–order regression weight with a Tucker structure, we can define for each $k\in[{0}\mathrel{{.}\,{.}}\nobreak{N}]$, a sketch matrix $\mathbf{M}^{(k)}\in\mathbb{R}^{R_{k}\times R_{k}}$ (e.g. a random projection or column selection matrix). This can then be used to sketch the factors $\mathbf{U}^{(k)}$ of the decomposition as $\mathbf{\tilde{U}}^{(k)}=\mathbf{U}^{(k)}(\mathbf{M}^{(k)})^{\top}$ be a sketch of factor matrix , and the core tensor $\mathcal{G}$ as $\mathcal{\tilde{G}}=\mathcal{G}\times_{0}\mathbf{M}^{(0)}\times\cdots\times_{N}\mathbf{M}^{(N)}$.

Applying this to tensor regression, we can apply this tensor dropout technique to the regression weights. Given an activation tensor $\mathcal{X}\in\mathbb{R}^{I_{0}\times\cdots\times I_{N}}$ and a set of $S$ target labels $\mathbf{\bm{y}}^{(k)}$, a Randomized-Tensor Regression Layer (R-TRL) is obtained from equation III and aims at minimizing the empirical risk:

[TABLE]

where $\mathcal{\tilde{W}}$ is a stochastic low-rank approximation of Tucker decomposition. In other words, in addition to the low-rank structure of the weights, we apply our tensor dropout in the latent subspace spanned by the decomposition.

For instance, in the case of a Tucker R-TRL, we have:

[TABLE]

This is the core of our proposed R-TRL, which incorporates tensor dropout within a TRL. Even though several sketching methods have been proposed, we focus here on R-TRL with two different types of binary sketching matrices, namely binary matrix sketching with replacement and binary diagonal matrix sketching with Bernoulli entries, which we detail below.

IV-A Bernoulli Tucker randomized tensor regression

For any $n\in[{0}\mathrel{{.}\,{.}}\nobreak{N}]$, let $\mathbf{\bm{\lambda}}^{(n)}\in\mathbb{R}^{R_{N}}$ be a random vector, the entries of which are i.i.d. Bernoulli($\theta$), then a diagonal Bernoulli sketching matrix is defined as $\mathbf{M}^{(n)}=\text{diag}(\mathbf{\bm{\lambda}}^{(n)})$.

When the low-rank structure on the weight tensor $\mathcal{\tilde{W}}$ of the TRL is imposed using a Tucker decomposition, the randomized Tucker approximation is expressed as:

[TABLE]

The main advantage of considering the above-mentioned sampling matrices is that the products $\mathbf{\tilde{U}}^{(k)}=\mathbf{U}^{(k)}(\mathbf{M}^{(k)})^{\top}$ or $\mathcal{\tilde{G}}=\mathcal{G}\times_{0}\mathbf{M}^{(0)}\times\cdots\times_{N}\mathbf{M}^{(N)}$ are never explicitly computed, we simply select the elements from $\mathcal{G}$ and the corresponding factors.

Interestingly, in analogy to dropout, where each hidden unit is dropped independently with probability $1-\theta$, in the proposed randomized tensor decomposition, the columns of the factor matrices and the corresponding fibers of the core tensor are dropped independently and consequently the rank of the tensor decomposition is stochastically dropped.

The tensor dropout acts as an implicit regularizer on the regression, by limiting the rank, at each iteration. This can be shown by examining the expectation of the stochastic loss, which can be expressed deterministically as the unrandomized empirical loss, plus an additional regularization term.

Theorem 1**.**

Tensor Dropout with Tucker decomposition is a deterministic regularized loss. (Proof in Appendix A)

The minimization objective (equation 4) can be reformulated by expanding the tensor contractions, and the expectation of the minimization objective becomes:

[TABLE]

IV-B Bernoulli CP randomized tensor regression

An interesting special case of equation (5) is when the weight tensor $\mathcal{\tilde{W}}$ of the TRL is expressed using a CP decomposition. In that case, we set $\mathbf{M}=\mathbf{M}^{(0)}=\cdots=\mathbf{M}^{(N)}=\mathop{\operator@font diag}\nolimits(\mathbf{\bm{\lambda}})$, with, for any $k\in[{0}\mathrel{{.}\,{.}}\nobreak{R}]$, $\lambda_{k}\sim\text{Bernoulli}(\theta).$

Then a randomized CP approximation is expressed as:

[TABLE]

The above randomized CP decomposition on the weights is equivalent to the following formulation (see proof in Appendix B):

[TABLE]

Based on the previous stochastic regularization, for an activation tensor $\mathcal{X}$ and a corresponding label vector $\mathbf{\bm{y}}$, the optimization problem for our tensor regression layer with stochastic regularization is given by:

[TABLE]

In addition, the above stochastic optimization loss can be rewritten as a deterministic regularized one:

Theorem 2**.**

Tensor Dropout with CP decomposition is equivalent to a deterministic regularized loss. (Proof in Appendix B)

[TABLE]

IV-C R-TRL with replacement

Previously, we focus on the R-TRL with Bernoulli sampling only. Our model is more general and can be applied to many different sampling settings. Here, we introduce one such case: the R-TRL with binary sketching matrix (i.e. with replacement). Specifically, we first choose $\theta\in[0,1]$.

Mathematically, we then introduce the uniform sampling matrices $\mathbf{M}^{(0)}\in\mathbb{R}^{R_{0}\times R_{0}},\cdots,\mathbf{M}^{(N)}\in\mathbb{R}^{R_{N}\times R_{N}}$. $\mathbf{M}_{j}$ is a uniform sampling matrix, selecting $K_{j}$ elements, where $K_{j}=R_{j}\nonscript\mskip-4.0mu plus -2.0mu minus -4.0mu\mkern 5.0mu\mathbin{\operator@font div}\penalty 900\mkern 5.0mu\nonscript\mskip-4.0mu plus -2.0mu minus -4.0mu\theta$. In other words, for any $i\in[{0}\mathrel{{.}\,{.}}\nobreak{N}]$, $\mathbf{M}^{(i)}$ verifies:

[TABLE]

In practice this is done efficiently by selecting directly the correct elements from $\mathcal{G}$ and its corresponding factors.

V Experimental evaluation

In this section, we present the numerical simulations, experimental setting, databases used and implementation details. We experimented on several datasets across various tasks, namely image classification and MRI-based regression. All methods were implemented 111code will be released upon acceptance to reproduce all results using PyTorch [37] and TensorLy [38]. For all adversarial attacks, we used Foolbox [39].

V-A Numerical experiments on synthetic data

Here, we empirically demonstrate the equivalence between our stochastic rank regularization and the deterministic regularization based formulation of the dropout.

To do so, we first created a random regression weight tensor $\mathcal{W}$ to be a third order tensor of size $(25\times 25\times 25)$, formed as a rank–$10$ Kruskal tensor, the factors of which were sampled from an i.i.d. Gaussian distribution. We then generated a tensor of $5000$ random samples, $\mathcal{X}$ of size $(5000\times 25\times 25\times 25)$, the elements of which were sampled from a Normal distribution. Finally, we constructed the corresponding response array $\mathbf{\bm{y}}$ of size $5000$ as: $\forall i\in[{1}\mathrel{{.}\,{.}}\nobreak{5000}],\,\mathbf{\bm{y_{i}}}=\langle{\mathcal{X}_{i}},{\mathcal{W}}\rangle$. Using the same regression weight tensor and same procedure, we also generated $1000$ testing samples and labels.

We use this data to train a CP R-TRL with our Tensor Dropout with rank $15$, using both our Bernoulli stochastic formulation (equation 9 in main paper) and its deterministic counter-part (equation 10 in main paper). We train for $500$ epochs, with a mini-batch size of $100$, and an initial learning rate of $10e-4$, which we decrease by a factor of $10$ every $200$ epochs. Figure 2 shows the loss function as a function of the epoch number. As expected, both formulations are identical.

V-B Large Scale Image Classification

In this section, we report results for large-scale image classification on the ImageNet dataset.

ImageNet [40]

is a large dataset for image classification composed of $1.2$ million training images and $50,000$ images for validation. We evaluate the classification error in terms of top-1 and top-5 classification accuracy on a $224\times 224$ single center crop from the raw input images. For training, we use a ResNet-101 and follow the same procedure and setting as [3, 16]. Note that the original ResNet paper [3] reports 10-crop validation error on ImageNet. Since then, single-crop has become the more common method of processing. The ResNet authors reported their single-crop results in their up-to-date code repository [41].

The results, Table I, show that our method outperforms the baselines, with higher classification accuracy even for large values of $\theta$. This suggests the stochastic rank regularization helps constrain the learned representations to be more general. That leads to reduced over-fitting and better out-of-sample generalizability. The improvement of tensor dropout is greater than that of regular dropout, which has been reported to be around 0.3 percentage points in previous studies on ResNets trained on ImageNet [42, 43].

Robustness study

In addition to improving performance by reducing over-fitting, we demonstrate that our proposed stochastic regularization makes the model more robust to perturbations in the input, for both random noise and adversarial attacks.

In particular, we tested robustness to adversarial examples produced using three different types of attack: Fast Gradient Sign Method (FGSM) [44], Basic Iterative Method (BIM) [44] and Projected Gradient Descend (PGD) [45] in Foolbox [39]. All attacks were untargeted. We note, that in [45], the authors showed that a method robust to PGD attack is reasonably robust to any form of first order gradient based attacks. In this method, the sign of the optimization gradient multiplied by the perturbation magnitude is added to the image in a single iteration. The perturbations we used are of magnitudes $\lambda\times 10^{-3},\lambda\in\{2,8,16\}$ . For an input image normalized in the range [0, 1], $\lambda$ will be divided by 255. For the iterative methods BIM and PGD, we follow [44] setting the step size to $1$ and the number of iterations to ${\text{min}(\lambda+4,1.25\lambda)}$. This fixed computational budget provides a compromise between the resources available to the attacker and the computational cost of the experiments.

As the results from Table II show, the proposed approach is significantly more robust to adversarial attacks across the entire range of adversarial noise tested. This further highlights that R-TRL is able to generalize better.

Moreover, in figure 4(a), we report the classification accuracy as a function of the added adversarial noise. Specifically, following [46], we sampled $1,000$ unseen test images. The models were trained without any adversarial training and adversarial noise was added using the Fast Gradient Sign method. Our Tucker R-TRL model is much more robust to adversarial attacks.

Adversarial defenses through randomness have been explored before, in different forms. For example, perturbing the input with isotropic Gaussian noise gives substantial and certifiable robustness to attacks on ImageNet [47], and adding random noise to each layer of the network provides defense against strong attacks [48]. However, these prior methods differ fundamentally to the one we present here. The randomness in our method is concerned with the rank of the factor tensors during training only and not perturbing the input or weights. Theoretically, our method is more related to regular dropout as we show in section V-D.

V-C Phenotypic trait prediction from MRI data

Contemporary neurological research often uses scalar measures of brain or grey matter volumes over time as a proxy for brain neuronal volume. For example, to provide estimates of neuronal loss over time [49]. However, the global and local structure of brain tissue also change [50], causing structural and textural changes that are not captured by scalar volume measures. CNNs are more suited for learning these intricate and complex relationships. However, the highly multi-dimensional mappings come with added risk of overfitting. Model generalizability is therefore paramount in this domain, providing a concrete demonstration of the strengths of the tensor dropout.

In a regression setting, we investigate the performance of our method on a very large MRI dataset on the task of age prediction.

The difference between an individual’s chronologial age and the age as predicted by a trained model is often referred to as brain-age. This metric which has been associated with diseases [23] and increased risk of mortality [22]. A more accurate and robust metric of brain condition can consequently lead to more accurate disease diagnoses.

In addition to the general neuroradiological motivations outlined above, this task is particularly interesting since MRI volumes are large 3D tensors, all modes of which carry important information. The spatial information is traditionally discarded during the flattening process, which we avoid by using a tensor regression layer. In these experiments we train the entire model, but any pre-trained network can be easily modified post-hoc to make use of the TRL.

UK Biobank brain MRI dataset [51]

is the world’s largest MRI imaging database of its kind. The aim of the UK Biobank Imaging Study is to capture MRI scans of vital organs for $100,000$ primarily healthy individuals by 2022. Associations between these images and lifestyle factors and health outcomes, both of which are already available in the UK Biobank, will enable researchers to improve diagnoses and treatments for numerous diseases.

The data we use here consists of T1-weighted $182\times 218\times 182$ MR images of the brain captured on a 3 T Siemens Skyra system. $11,500$ are used for training, $3,800$ are used for validation and $3,800$ samples are used to test. The target label is the age for each individual at the time of MRI capture. We use skull-stripped images that have been aligned to the MNI152 template [52] for head-size normalization. We then center and scale each image to zero mean and unit variance for intensity normalization.

Results:

For MRI-based experiments we implement an 18-layer ResNet with three-dimensional convolutions. We minimize the mean squared error using Adam [53], starting with an initial learning rate of $10^{-4}$, reduced by a factor of 10 at epochs 25, 50, and 75. We train for 100 epochs with a mini-batch size of 8 and a weight decay ($\mathrm{L}_{2}$ penalty) of $5\times 10^{-4}$. For Tucker-based R-TRL we used a tensor with rank $128\times 6\times 7\times 6$. For CP-based R-TRL we used a Kruskal tensor with $82$ components. As previously observed, our randomized tensor regression network outperforms the 3D-ResNet baseline by a large margin, Table III. To put this into context, the current state-of-art for convolutional neural networks on age prediction from brain MRI on most datasets is an MAE of around 3.6 years [54].

Robustness study:

We tested the robustness of our model to white Gaussian noise added to the MRI data. Noise in MRI data typically follows a Rician distribution but can be approximated by a Gaussian for signal-to-noise ratios (SNR) greater than $2$ [55]. As both the signal (MRI voxel intensities) and noise are zero-mean, we define $\mathrm{SNR}=\frac{\sigma_{\mathrm{signal}}^{2}}{\sigma_{\mathrm{noise}}^{2}}$, where $\sigma$ is the variance. We incrementally increase the added noise in the test set and compare the error rate of the models.

The ResNet with R-TRL is significantly more robust to added white Gaussian noise compared to the same architectures without it (figure 5). At signal-to-noise ratios below 10, the accuracy of a standard fully-connected ResNet is worse than a naive model that predicts the mean of training set (MAE = $7.9$ years). Brain morphology is an important attribute that has been associated with various biological traits including cognitive function and overall health [56, 57]. By keeping the structure of the brain represented in MRI in every layer of the architecture, the model has more information to learn a more accurate representation of the entire input. Randomly dropping the rank forces the representation to be robust to confounds. This a particularly important property for MRI analysis since intensities and noise artifacts can vary significantly between MRI scanners [58]. Randomized tensor regression layers enable both more accurate and more robust trait predictions from MRI that can consequently lead to more accurate disease diagnoses.

V-D Ablation Studies on CIFAR-100

In the image classification setting, we perform a thorough study of our method on the CIFAR-100 dataset. We empirically compare our approach to both standard baseline, traditional tensor regression, and regular dropout, and assess the robustness of each method in the face of adversarial noise.

CIFAR-100 [59]

consists of 60,000 $32\times 32$ RGB images in 100 classes, divided into $50,000$ images for training and $10,000$ for testing. We pre-processed the data by centering and scaling each image and then augmented the training images with random cropping and random horizontal flipping.

We compare the randomized tensor regression layer to full-rank tensor regression, average pooling and a fully-connected layer in an 18-layer residual network (ResNet) [3]. For all networks, we used a batch size of $128$ and trained for $400$ epochs, and minimized the cross-entropy loss using stochastic gradient descent (SGD). The initial learning rate was set to $0.01$ and lowered by a factor of $10$ at epochs $150$, $250$ and $350$. We used a weight decay ($\mathrm{L}_{2}$ penalty) of $10^{-4}$ and a momentum of $0.9$.

Classification results:

Table IV presents results obtained on the CIFAR-100 dataset, on which our method matches or outperforms other methods, including the same architectures without R-TRL. Our regularization method makes the network more robust by reducing over-fitting, thus allowing for superior performance on the testing set.

A natural question is whether the model is sensitive to the choice of rank and $\theta$ (or drop rate when sampling with repetition). To assess this, we show the performance as a function of both rank and $\theta$ in figure 7(a). The reduction in rank is presented as the $\text{compression ratio}=\frac{\text{size of full tensor}}{\text{size of factorized cores and factors}}$. As can be observed, there is a large surface for which performance remains the same while decreasing both parameters (note the logarithmic scale for the rank). This means that, as demonstrated, choosing good values for CIFAR-100 was not a problem. Rank insensitivity cannot be guaranteed across all tasks and datasets, but rank-selection was not found to be an issue for ImageNet or MRI tasks.

Performance as a function of rank and $\theta$ in replacement R-TRL

To illustrate the generality of our approach, which does not depend on the Bernoulli sampling, we perform a similar experiment with a different randomization: instead of using a Bernoulli random variable, we sample components with replacement according to a uniform sampling matrix (figure  7(b)). As for the Bernoulli case, there is a large surface for which performance remains the same while decreasing both parameters.

Comparison with regular dropout

One question is whether our proposed tensor dropout is more robust than traditional dropout applied directly to the weights. To test this, we apply FGSM adversarial perturbations to each method, with varying magnitudes $\lambda\times 10^{-3},\lambda\in\{1,2,4,8,16,32,64,128\}$. We sample $1,000$ images from the test set [46]. The models were trained without any adversarial training, on the training set, and adversarial noise was added to the test samples using the Fast Gradient Sign method. The results of which can be seen in figure 6(a). Our model is much more robust to adversarial attacks. Intuitively our method is able to leverage redundancies in the latent subspace, without creating holes in the weights, unlike dropout. In addition, since the randomization is used only during training (and not at test time), this forces the latent decomposition to be over-complete and account for noise, thus rendering the model more robust to perturbation.

VI Conclusion

We introduced tensor dropout, a novel randomized tensor decomposition, suitable for end-to-end training of tensor regression layers. Adding stochasticity on the rank during training renders the network significantly more robust and leads to better performance. This results in networks that are more resilient to noise, both adversarial and random, without any addition such as adversarial training. Our results demonstrate superior performance on a variety of real-life, large-scale challenging tasks, including MRI data and images, as well as increased robustness. Finally, we establish the link between this randomized TRL and dropout on the deterministic TRL.

Acknowledgment

This research has been conducted using the UK Biobank Resource under Application Number 18545.

Appendix A Connection between stochastic and deterministic regression

In this section, we study the empirical risk minimization of Tensor Dropout. We establish the link between Tensor Dropout with stochastic Tucker decomposition and the deterministic regularized loss. We do this by proving the equality between equations 4 and 7 (Theorem 1).

Proof of Theorem 1.

The minimization objective in equation 4 can be reformulated by expanding the tensor contractions in Tucker form as:

[TABLE]

By considering the expectation over the empirical risk, we get:

[TABLE]

The expectation E and variance V can then be written separately, using the Tucker expansion from Equation 13:

[TABLE]

[TABLE]

Therefore,

[TABLE]

∎

Appendix B

In this section, we study Tensor Dropout’s empirical risk minimization and establish the link between Tensor Dropout with stochastic CP decomposition and the corresponding deterministic regularized loss (Theorem 2). To do so, we first establish the equality between 8 and 9. We then prove Theorem 2, which is the equivalence of equations 4 and 11.

CP decomposition with tensor dropout.

Here we prove the following equivalence (equations 8 and 9):

[TABLE]

This can be seen by looking at the individual elements of the sketched factors. Let $k\in[{0}\mathrel{{.}\,{.}}\nobreak{N}]$ and $i_{k}\in[{0}\mathrel{{.}\,{.}}\nobreak{I_{k}}],r\in[{0}\mathrel{{.}\,{.}}\nobreak{R-1}]$. Then

[TABLE]

Since $\mathbf{M}=\text{diag}(\mathbf{\bm{\lambda}})$, i.e. $\forall i,j\in[{0}\mathrel{{.}\,{.}}\nobreak{R-1}],\mathbf{M}_{ij}=0$ if $i\neq j$, and $\lambda_{i}$ otherwise, we get

[TABLE]

It follows that $\mathcal{\tilde{W}}_{i_{0},i_{1},\cdots,i_{N}}=\sum_{r=0}^{R-1}\lambda_{k}\mathbf{U}^{(0)}_{k}\circ\cdots\circ\lambda_{k}\mathbf{U}^{(N)}_{k}$ Since $\lambda_{r}\in\{0,1\},$ we have

[TABLE]

∎

Proof of Theorem 2.

Here we prove that Tensor Dropout with stochastic CP decomposition is equivalent to deterministic regularized loss. By considering the expectation over the empirical risk, we get:

[TABLE]

For Tensor dropout where weights are constructed with CP factors: $\mathcal{\tilde{W}}=\sum_{r=0}^{R-1}\lambda_{n}\mathbf{U}_{i_{0}r_{0}}^{(0)}\dots\mathbf{U}_{i_{N}r_{N}}^{(N)}$ the minimization objective can be written as:

[TABLE]

[TABLE]

Therefore,

[TABLE]

∎