Learning Low-Rank Tensor Cores with Probabilistic ℓ0-Regularized Rank Selection for Model Compression

1 Introduction

In the last decade, deep neural networks have achieved excellent results on a variety of ill-posed machine learning tasks in various fields, including image classification and object detection [He et al., 2016] in computer vision (CV), neural machine translation (NMT) and language modeling [Vaswani et al., 2017] in the field of natural language processing (NLP), and music source separation [Rouard et al., 2023]. Many pioneering models have a large number of parameters, which are considered to enhance the feasibility of learning optimization and the expressive power of the model in later studies [Arora et al., 2018]. For example, a ResNet-50 has 25.6M parameters and needs 97.5MB memory storage. The large storage and computation overhead associated with a huge number of model parameters makes it difficult to apply deep neural networks on resource-constrained devices, such as smartphones, wearables, and IoT devices. Several works [Ba and Caruana, 2014] have shown that modern neural networks are highly redundant because most parameters are attributed to a small subset of the learned parameters. This has given rise to the study of various types of model compression techniques, including low-rank approximation [Yaguchi et al., 2019], knowledge distillation [Hinton et al., 2015], pruning, and quantization [Han et al., 2015].

As a promising low-rank technique, tensor decomposition efficiently represents high-order arrays using decomposed tensor cores. It has been an attractive tool in quantum system modeling [Or´us, 2019], data mining [Kwon et al., 2023], and applied mathematics applications [Dolgov et al., 2021], and has been applied to compress neural networks [Novikov et al., 2015]. A compact model in the decomposed form can be used to replace the original model with little performance loss. Thanks to the good expressive power of tensor decomposition, tensor decomposition-based compression methods usually result in large compression rates. However, they usually suffer from the rank selection problem, which is particularly serious in the successful chain-structured TensorTrain [Novikov et al., 2015] and Tensor-Ring [Wang et al., 2018] methods due to their large rank search spaces.

In the data completion problem, it has been extensively studied that ranks are essential influence factors to the expressive power of decomposition [Hashemizadeh et al., 2020]. A similar phenomenon [Li et al., 2021] holds for deep compression since tensor ranks determine the shapes of decomposed cores and thus control the trade-off between model accuracy and compression rate. Early works [Yin et al., 2022] typically select ranks manually for each decomposed layer, where all ranks are set to equal and treated as one hyperparameter. However, it is intuitively over-simplified and usually leads to sub-optimal performance. Several automated nonuniform rank selection methods are proposed, turning out to be better than the simple uniform scheme. However, existing methods rely on iterative large-scale low-rank tensor approximations [Gusak et al., 2019], retraining [Li et al., 2021], Bayesian network design [Hawkins and Zhang, 2021], or integer approximate searching algorithms such as genetic algorithms [Li et al., 2021] and reinforcement learning [Cheng et al., 2020], which complicate the optimization and require significantly additional training overhead and memory consumption.

To overcome these limitations, in this paper, we propose to apply an approximate Bernoulli gate to each of the slices of tensor cores and apply probabilistic ℓ0 regularization on the ranks of decomposed tensor cores during training. The proposed ℓ0-regularized scheme requires only a little extra overhead than training a vanilla tensorized neural network and manages to find a superior set of ranks than baselines. We also introduce an initialization scheme. The technical contributions are summarized as follows.

• We propose a novel gated form of tensor contractions (illustrated in Fig. 2) and a probabilistic ℓ0-regularized rank selection method that allows simultaneous parameter learning and rank selection, as well as flexible combination with arbitrary tensor decompositions.
• We introduce the scheme of optimization using an efficient and scalable gradient-based optimizer with the help of the Gaussian-based Bernoulli continuous relaxation.
• Our method has been successfully applied for tensorized fully-connected layers, convolutional layers, and embedding layers, and comprehensive experiments on various CV and NLP tasks show its efficiency.

2 RelatedWork

Tensor Methods for Deep Compression. Tensor decomposition has been successfully applied to deep model compression. Various types of decompositions, e.g. Tucker [Gusak et al., 2019], CP [Lebedev et al., 2014], Tensor-Train (TT) [Novikov et al., 2015; Yin et al., 2022], Tensor-Ring (TR) [Wang et al., 2018], and Tensor Network [Hayashi et al., 2019], are used in designing significantly compact neural network layers. Early works focus on fully-connected layers [Novikov et al., 2015; Yang et al., 2017] and convolutional layers [Garipov et al., 2016; Wang et al., 2018], while more recent works generalize to tensorized word embedding [Hrinchuk et al., 2020] and self-attention [Ma et al., 2019]. Our proposed method can be well combined with existing techniques and is not limited to specific types of decomposition, layers, and tasks, as will be shown in the methodology and experiments.

Tensor Rank Selection. A major problem of tensor-based deep compression is how to determine the ranks of the compact tensorized representation of layers. Determining the tensor rank is NP-hard on CP decomposition [Hillar and Lim, 2013], and it is likewise non-trivial on other decompositions, which is commonly circumvented by setting all ranks equal in model compression. This uniform rank selection scheme sacrifices more performance than non-uniform ones [Kodryan et al., 2023]. To achieve non-uniform rank selections for a better compression-performance trade-off, MUSCO [Gusak et al., 2019] iteratively performs low-rank approximation with ranks determined using variational Bayesian matrix factorization [Nakajima et al., 2010] and fine-tuning. [Dai et al., 2023] jointly learns full and low-rank Tucker models and utilizes Tucker decomposition at each iteration to reduce ranks. [Cheng et al., 2020] treats ranks of decomposition as actions and performance plus compression ratio as rewards, and then searches TR ranks using a reinforcement learning agent. [Li et al., 2021] proposes to progressively search TR ranks using neural architecture search techniques with genetic algorithms. Bayesian approaches [Hawkins and Zhang, 2021; Kodryan et al., 2023] model the distributions of TT ranks and learn the proper ranks via smart prior choices. Current approaches typically treat the rank selection problem as combinatorial optimization and apply approximate search algorithms to solve it, resulting in significantly additional computational costs for searching and re-training. In contrast, our proposed method requires only a little additional overhead than training a vanilla tensorized neural network, without sacrificing prediction performance.

Regularization for Sparsity. Prior works in the context of feature selection and neural network pruning have extensively investigated embedded regularization approaches. To mention a few relevant, LASSO [Tibshirani, 1996] selects feature subsets with convex ℓ1 constraints. Grouped ℓ1 regularization allows structured neural network pruning [Wen et al., 2016]. Non-differentiable ℓ0 is introduced by its relaxed alternatives, such as Concrete [Maddison et al., 2016], Hard Concrete [Louizos et al., 2018; Zhen et al., 2022] and stochastic gate approximation [Yamada et al., 2020]. Despite the difficulty of non-convex optimization, ℓ0-based approaches outperform the convex ℓ1 in both feature selection [Yamada et al., 2020] and model compression [Louizos et al., 2018]. Therefore, it is promising to impose ℓ0 regularization to learn the appropriate tensor rank of the weight parameters and prune tensorized neural networks.

3 Preliminary

Notations. A tensor, in machine learning literature, refers to a multi-dimensional array. Throughout the paper, 'x', 'X', and 'X' denote scalar, matrix, and tensor respectively. The (i1, . . . , iN)-th entry of N-order tensor A is denoted as A(i1, . . . , iN). We use the notation [n] = {1, . . . , n}, where n is a positive integer. The ℓ0 (quasi-)norm ∥ · ∥0 refers to the number of non-negative entries of a matrix.

Contraction. Tensor contraction is a mathematical operation that extends the concept of matrix multiplication to higher-dimensional tensors. It involves summing the products of corresponding elements along matched indices, resulting in a new tensor with reduced dimensions. A formal definition can be found in [Kodryan et al., 2023]. For instance, given two tensor G1 ∈ Rr1×d1×r2 and G2 ∈ Rr2×d2×r1 , the output tensor A ∈ Rd1×d2 is obtained through tensor contractions along two matched indices as follow:

Tensor Decomposition. In this paper, we adopt TR [Wang et al., 2018] as the main decomposition method for compression, and the same idea applies to other methods as well. A TR decomposition [Zhao et al., 2016] represents X ∈ RI1×···×Id with cores {Gi}di=1 as follows:

where Gi ∈ RRi-1×Ii×Ri for i ∈ [d] and R0 = Rd. The intermediate dimensions {Ri}di=1 are the TR ranks. Constraining R0 = Rd = 1 yields a TT decomposition model [Oseledets, 2011]. With the rest intermediate ranks equal, a TR decomposition can be seen as a linear combination of TT decompositions and thus has stronger expressive power.

4 Methodology

Suppose that X is the random variable in the input domain RI with corresponding variable Y in the output domain RO , let fθ(·) be the L-layer neural network predictor parameterized by weights θ = {Wl}Ll=1 . Given a loss function L, the expected risk minimization can be formulated as:

where EX [·] is the expectation over the distribution P(X).

4.1 Tensorized Neural Network Layers

We start by parameterizing the predictor using decomposed cores. To compress the most common layers involving linear transformations (including fully-connected layers 1 and embedding layers), following [Novikov et al., 2015; Wang et al., 2018], the weight matrix W ∈ RO×I is reshaped to a high-order tensor W ∈ RO1×···×On×I1×···×Im , where O = ∏ni=1 Oi and I = ∏mi=1 Ii. Then W can be represented by TR cores {Gi}m+n i=1 of ranks {Ri}m+n i=0 . An arbitrary input sample x ∈ RI is correspondingly reshaped to X ∈ RI1×···×Im . Thus, the reshaped output Y ∈ RO1×···×On is calculated through contractions between the cores and X:

where

Once Y is obtained, it can be flattened to y ∈ RO .

The original convolutional layer has the weight of a kernel tensor K ∈ RO×I×K×K and deals with the input x of shape H × W × I, where O, I, and K are the number of output channels, the number of input channels, and the kernel size, respectively. Similar to (4), K and X are reshaped to RO1×···×Om×I1×···×Im×K×K and RH×W×I1×···×Im , respectively, where O = ∏ni=1 Oi and I = ∏mi=1 Ii. Then, K is represented with TR cores {Ui}ni=1 ∪ {G} ∪ {Vj}mj=1 , where Ui ∈ RRi-1×Ii×Ri , G ∈ RRn×K×K×Rn+1 , and Vj ∈ RRj+n×Oj×Rj+n+1 . Then, the original convolutional operation, formulated by

1 Biases are omitted in this part to simplify notation.

where h = (h′ -1)S +d1 -P and w = (w′ -1)S +d2 -P , with S ∈ N being the stride, and P ∈ N being the padding size, can be reformulated by a three-step alternative:

(8)

Note that in (5), (7), and (9), TR ranks determine the number of parameters and the expressive power of the model, whose search space could be potentially large with a large number of cores. Setting all ranks equal could be over-simplified while manually tuning each of the ranks is impractical. Therefore, we consider an automatic rank selection method for tensorized neural networks.

4.2 Automatic Rank Selection via Probabilistic ℓ0 Regularization

In this section, we propose an ℓ0-based rank selection formulation through approximate Bernoulli gates, to remove the rank redundancy without performance loss in (5), (7), and (9). Following the decomposition scheme in Sec. 4.1, θ in (3) is parameterized with cores {{G(l) i }nl i=1 }Ll=1 , where nl is the number of cores of the l-th layer. Based on (3), we define the following low-rank constrained optimization problem:

where Q is some target rank.

Gated Tensor Contraction with ℓ0 Regularization. Since (10) is unsolvable with gradient methods due to the nondifferentiable nature of the rank function, one may replace the rank minimization with trace norm penalty [Tomioka et al., 2010]. However, penalizing on the trace norm is SVD-based and thus computationally intractable in large-scale neural networks. To overcome the limitation, we introduce Bernoulli gates applied to slices of tensor cores to shrink their ranks and propose the gated tensor contraction. Specifically, we propose to insert a diagonal matrix Zi ∈ RRi×Ri between each pair of adjacent factors (Gi,Gi+1) to be contracted, as shown in Fig. 2. Diagonal entries of Zi are independent random variables and satisfy Zi(j, j) ∼ Bern(πij). Then the contraction between Gi and Gi+1 is replaced by:

where Zi acts as gates on the ranks of cores Gi and Gi+1. The learning problem of (10) is then transformed to an ℓ0regularized optimization problem:

where f is parameterized by θ′ = {{G(l) i , {π(l) ij }R(l) i j=1 }nl i=1 }Ll=1 and λ is the non-negative hyperparameter. Note that if Zi(j, j) = 0, the slices Gi(:, :, j) and Gi+1(j, :, :) can be zeroed out and pruned without changing the result of tensor contractions. Therefore, (12) can be seen as a structured pruning on tensorized neural networks that brings an actual reduction in both parameters and computations. A simple and conventional way is to relax ℓ0 to ℓ1. However, practically, introducing an ℓ1 penalty into gradient methods does not result in sparse solutions [Ziyin and Wang, 2023] and requires post-training thresholding.

Approximate Bernoulli Gates for ℓ0 Regularization. The optimization of (12) involves sampling from discrete distribution and suffers from high variance 2 . Inspired by stochastic gates [Yamada et al., 2020], we propose to use the Gaussianbased Bernoulli continuous relaxation, which makes the rank selection and the end-to-end learning of a tensorized neural network stable and scalable. For simplicity of notations, we consider one layer with n tensor cores of the neural network and omit the index l. An approximate Bernoulli gate is defined by Zi(j, j) := zi,j = max(0,min(1, µi,j + σϵi,j)), where ϵi,j ∼ N(0, 1) and σ is fixed during training. Thus, its gradient with respect to µi,j can be computed using standard backpropagation, and the expected ℓ0 regularizer of Zi becomes

2 An ablation study is done in the supplementary material. Optimizing discrete distribution is feasible but unstable and inefficient. Code and supplementary materials are available at https: //github.com/ctxGou/Tensor-L0-Compression.

where Φ(·) is the standard Gaussian CDF. Given the training data {xk, yk}Kk=1 , the Monte Carlo (MC) approximation of the regularized empirical risk minimization of (12) is

where the set of variables is θ′ = {{G(l) i , {µ(l) ij }R(l) i j=1 }nl i=1 }Ll=1 , z[m] i,j = max(0,min(1, µi,j + σϵ[m] i,j )) is the m-th MC sample, ϵ[m] i,j ∼ N(0, 1), and M is the number of samples. The gradient of the objective of (14) w.r.t. µi,j is estimated as

where ˆGi = Gi ×3 Zi. The optimization incorporated with the ℓ0 norm can be solved using ubiquitous stochastic gradient methods through the approximations. Compared to a vanilla tensorized neural network parameterized with tensor cores, our formulation introduces minor extra parameters with their number linear to the sum of tensor ranks.

Learned Ranks and Pruning. During the evaluation phase, the injected noises ϵi,j are removed for deterministic prediction, leading to zi,j = max(0,min(1, µi,j)). There is no need to round the learned gates after training since (1) most of the gates converge to 0 and 1, and (2) rounding of the gates scales some slices of the cores, resulting in output distortion and possible performance degradation. The trained gates {Zi}ni=1 in (14) indicate the learned ranks, because the slices Gi(:, :, j) and Gi+1(j, :, :) can be zeroed out and pruned if Zi(j, j) = 0. The learned rank ˆRi is defined as follows:

where I(·) is the indicator function. It is worth noting that sparsity does not imply instant compression. To this end, we create a compact representation (denoted by G∗i ) of the core tensor ˆGi ∈ RRi-1×d×Ri by excluding the pruned slices, leading to size reduced to ˆRi-1 × d × ˆRi.

4.3 Initialization

Training a tensorized neural network from scratch requires proper weight initialization. We propose to initialize all µij's by some real numbers larger than 1, so all gates are one at the beginning, making the initialization methods on original tensorized neural networks applicable. For the sake of convergence and stable training, model weights are designed to maintain variance during forward propagation and during backward propagation. Specifically, consider a TR fullyconnected layer or embedding layer with weights {Gi}n+m i=1 , and ranks {Ri}n+m i=1 , the input and output are denoted by X ∈ RI1×...×Im and Y ∈ RO1×...×On , respectively. We make an assumption on X and {Gi}n+m i=1 that their entries are i.i.d., and thus the following equation holds:

following Proposition 3.3 of [Pan et al., 2022], where I = ∏mi=1 Ii. Similar to Kaiming Initialization [2016], we consider the impact of ReLU activation and seek σ2 (Y) = 2σ2 (X) to avoid exponentially reducing or magnifying the magnitude of inputs. Additionally, we simply assume σ(G1) = · · · = σ(Gn+m). Therefore, (17) becomes

In a similar way, the variance of tensorized convolutional layer TR cores with {Gi}n+m+1 i=1 and ranks {Ri}n+m+1 i=1 is derived as follows:

where K is the kernel size.

Tensor cores are initialized to uniform distribution following the variance suggested as (18) and (19). In some rare cases, we observe that the learned rank ˆRi drops to 0 with all gates turning to 0's, for which the whole layer is zeroed out and the data flows are destroyed. Such glitches can be circumvented by stopping updating the diagonal entries of Zi once their values are lower than a pre-defined minimal rank.

5 Experiments

The proposed approach applies to various layers that can be parameterized with factorized tensor formats. Therefore, in this section, we evaluate the proposed method on various learning models for a wide range of applications. To evaluate the fundamental availability of our method on fully-connected layers, we initiate the experimental study with a basic 2-layer MLP on the MNIST dataset [LeCun et al., 1998], following the structure of [Hawkins and Zhang, 2021] (Sec. 5.1). Then we extend to convolutional layers based on LeNet-5, which is used in [Wang et al., 2018] (Sec. 5.2). To show the robustness and scalability of the proposed approach on deeper networks, we further investigate the performance of ResNet-32/110 [He et al., 2016;

Wang et al., 2018] on the CIFAR-10 dataset [Krizhevsky et al., 2009] (Sec. 5.3), which is a particularly common setting in deep compression literature. Moreover, we evaluate our approach to embedding layers of LSTMs [Hochreiter and Schmidhuber, 1997] and Transformers [Vaswani et al., 2017] on NLP tasks, including text sentiment classification and neural machine translation (Sec. 5.4).

Methods for Comparison. We compare our proposed method with tensor decomposition baselines with uniform ranks [Wang et al., 2018], and cutting-edge rank selection approaches [Cheng et al., 2020; Gao et al., 2020; Li et al., 2021; Kodryan et al., 2023]. We implemented Tensor-Ring Net [Wang et al., 2018] and conducted the experiments of MARS [Kodryan et al., 2023] in Sec. 5.1 by their released code.

5.1 Tiny Example: 2-Layer MLP

In the first experiment, we evaluate our method on a toy 2-layer MLP on the MNIST image classification task using 60k/10k train/test samples. We follow the same settings of structures in [Hawkins and Zhang, 2021; Kodryan et al., 2023], which train the network and select ranks in an endto-end fashion. A tensorized neural network with two fullyconnected layers [Novikov et al., 2015] of 784 × 625 and 625 × 10 is parameterized with two TT-matrix representations of sizes (7, 4, 7, 4) × (5, 5, 5, 5) and (25, 25) × (5, 2), respectively. Initial ranks are set to 20, as recommended in [Kodryan et al., 2023]. We also extend the 2-layer network to TR decomposition of shapes (7, 4, 7, 4, 5, 5, 5, 5) and (5, 5, 5, 2) for two layers, respectively. We set σ = 1 and the MC sampling number M = 1.

Results are shown in Table 1. We test our method and reproduce the results in [Kodryan et al., 2023] (denoted as MARS) by running the provided codes, where soft and hard represent different hyperparameter choices. Both models are trained for 300 epochs using Adam with decaying learning rates from 1e - 2. The results indicate that with a comparable performance loss, ours learns a tensorized neural network of much lower ranks, suggesting the effectiveness of the ℓ0based regularization.

The two results of Ours (TR) are obtained by using different values of λ's (λ = 0.003 and λ = 0.001). The hyperparameter λ controls the strength of regularization on the sparsity of the diagonal entries of Zi. It helps to dramatically decrease the size of the search space, compared to manual tuning. As we will further see in the following experiments in Sec. 5.3, the search of λ can be further averted using a budget-aware scheme.

5.2 LeNet-5 onMNIST

We follow the network architecture described in [Wang et al., 2018]. The model is a modified LeNet-5-like CNN, which incorporates two convolutional layers with 20 and 50 output channels, respectively, and two linear layers with 320 and 10 neurons, respectively, resulting in a cumulative parameter count of 429K in an uncompressed state. To embed TR decomposition to the network and present a fair comparison, we follow the decomposition settings of [Wang et al., 2018]. We train the nets using Adam with decaying learning rates from 5e - 3. In this experiment, we set 3σ = 0.5, M = 1, and initial ranks as 25.

Table 2 reports the prediction performance and compression rates of ours and other rank selection approaches. Our method achieves a competitive accuracy-compression tradeoff. It is worth noting that PSTRN [Li et al., 2021] utilizes a genetic algorithm search scheme, and therefore requires an evolutionary phase in which thousands of models are trained to search the appropriate sets of ranks. [Li et al., 2021] reports that the search processes cost about 5 GPU days on NVIDIA V100s, which is much more time-consuming than our end-toend rank selection approach.

Figure 3 further shows the learned ranks of the proposed method and PSTRN. PSTRN learns very small ranks in some cores but large ranks in other cores, while the learned ranks within a layer of ours are closer, resulting in a larger compression rate. Whether the optimal rank should be more even or vice versa, which could be relevant to the function of each tensor core in tensorized neural networks, is still an open problem, and we leave it for future research.

3 A hyperparameter sensitivity analysis and experiments on KMNIST and FMNIST are presented in the supplementary material.

5.3 ResNets on CIFAR-10

The combinations of ResNets [He et al., 2016] and CIFAR10 [Krizhevsky et al., 2009] are popular settings in the literature on deep compression. CIFAR-10 [Krizhevsky et al., 2009] consists of 50k images for training and 10k images for testing. We utilize SGD optimizer with a decaying learning rate initiating from 0.1. The models are trained for 200 epochs with batch size 128. The stochasticity introduced by the gates Zi's makes the ℓ0 approximation theoretically and empirically practical, but it also brings difficulties to the training of tensor core parameters due to the perturbation during the training phase. Thus, on the relative deep network ResNet, we set the variance of sampled noises σ = 10-3 , M = 1, and initial ranks as 15. The choice of λ directly determines the final compression rate, which brings inconvenience in practice since one may need to do cross-validation to meet a specific trade-off between compression and accuracy. We circumvent the problem by stopping updating the parameters of Zi's and removing their stochasticity after the target compression ratio is reached. For a fair comparison, we follow [Wang et al., 2018; Cheng et al., 2020; Li et al., 2021] by using the same TR decomposition (the same shapes and orders of core tensors) of ResNet-32 proposed in [Wang et al., 2018]. We also conduct an experiment with ResNet-110 on CIFAR-10.

Results are shown in Tables 3 and 4. Our method outperforms the uniform rank baselines and rank selection approaches at a compression ratio of around 5×. Specifically, ours has a 0.5% accuracy advantage over Tensor-Ring Net baselines with a uniform rank of 10 at a similar compression rate, which again validates the superiority of the non-uniform rank scheme. While it is reported in [Li et al., 2021] that the search processes for ResNet-32 of PSTRN cost about 3.2 GPU days on NVIDIA V100, ours requires about ∼ 0.11 GPU days on NVIDIA 2080Ti, which is ∼ 1.18× time of training a TR Net with a uniform rank of 15 and achieves a 0.33% accuracy gain under 5× compression rate. The search process of PSTRN is done with genetic algorithms and therefore becomes extremely time-consuming when the search space is large. In addition, PSTRN adopts a coarsegrained layer-wise rank selection scheme. In contrast, our method can search the rank of every tensor core with little extra overheads.

5.4 Embedding Layers

Word embedding is one of the core components of neuralbased NLP models, which encapsulates linguistic information into vectors. However, as the vocabulary and embedding size increase, the number of its parameters can become large. For instance, neural machine translation models typically employ embeddings with a vocabulary size of ∼30k and an embedding size of ∼1024, totaling around 3 × 107 parameters.

A compact TT representation of embedded layers is proposed in [Hrinchuk et al., 2020], replacing the original huge layers and drastically reducing the number of parameters. The implementation of this method is the same as the TT-format linear layer without bias and can be easily extended to TR. We first compare our probabilistic ℓ0based rank selection scheme with [Hrinchuk et al., 2020; Kodryan et al., 2023] on TT-embeddings in a text sentiment classification task, and then validate its effectiveness on TRembeddings [Hrinchuk et al., 2020] in a NMT task.

Text Sentiment Classification. We follow [Hrinchuk et al., 2020; Kodryan et al., 2023] by using the fine-grained five-class Stanford Sentiment Treebank (SST) dataset 4 . The dataset consists of 8.54k texts for training and 2.21k texts for testing. We take the most frequently appearing 17200 words and embed them into vectors of 256 dimensions using embedding layers of the original shape 17200 × 256. Embedding layers are followed by a bidirectional 2-layer LSTM with a hidden size of 128 and a dropout rate of 0.5.

4 Details at https://nlp.stanford.edu/sentiment/.

In our experiment, initial ranks are set as 16. To compare with [Kodryan et al., 2023; Hrinchuk et al., 2020], we use TTembeddings of the same shape (10, 10, 12, 15) × (4, 4, 4, 4). As shown in Table 5, our method achieves similar accuracy but a higher compression ratio than MARS and outperforms baselines in both metrics.

Neural Machine Translation. In this experiment, we use a standard Transformer [Vaswani et al., 2017] and different embedding layers to conduct NMT on the IWSLT'14 DeEn [Cettolo et al., 2014] dataset, which consists of 170k German-English sentence pairs for training and 6.75k pairs for testing. We use a vocabulary of 32768 tokens via a joint source and target Byte-Pair Encoding (BPE). A full hyperparameter list can be found in the supplementary material.

Results are shown in Table 6. We use a TR-embedding of shape (32, 32, 32) × (8, 8, 8) for TR and Ours and search the ranks initiated from 32. On this challenging task, our method learns a more compact layer with a preferable BLEU score compared to the TR baselines with uniform ranks.

6 Conclusion

In this paper, we present an ℓ0-regularized formulation for jointly learning model weights and decomposition ranks of tensorized neural networks in an end-to-end manner and propose to efficiently solve this by probabilistic approximations of ℓ0 (quasi-)norm. Comprehensive experiments are done to validate that the proposed scheme learns more compact models. Our future work includes applying the method to other layers, e.g. tensorized self-attention, and combining existing orthogonal techniques, such as knowledge distillation [Wang et al., 2021; Li et al., 2022; Gu et al., 2022].

Acknowledgments

Canh Hao Nguyen was supported by MEXT KAKENHI Grant Number 22K12150. Hiroshi Mamitsuka was supported by MEXT KAKENHI Grant Numbers 21H05027 and 22H03645.