CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor Optimization
Authors: Zi Yang
Ziyue Liu
Samridhi Choudhary
Xinfeng Xie
Cao Gao
Siegfried Kunzmann
Zheng Zhang
Venue: NeurIPS 2024 poster (2024)
Keywords: Compressed Training
Memory-efficient
Tensor
Primary area: optimization_for_deep_networks

Abstract: Training large AI models such as LLMs and DLRMs costs massive GPUs and computing time. The high training cost has become only affordable to big tech companies, meanwhile also causing increasing concerns about the environmental impact. This paper presents CoMERA, a **Co**mputing- and **M**emory-**E**fficient training method via **R**ank-**A**daptive tensor optimization. CoMERA achieves end-to-end rank-adaptive tensor-compressed training via a multi-objective optimization formulation, and improves the training to provide both a high compression ratio and excellent accuracy in the training process. Our optimized numerical computation (e.g., optimized tensorized embedding and tensor-vector contractions) and GPU implementation eliminate part of the run-time overhead in the tensorized training on GPU. This leads to, for the first time, $2-3\times$ speedup per training epoch compared with standard training. CoMERA also outperforms the recent GaLore in terms of both memory and computing efficiency. Specifically, CoMERA is $2\times$ faster per training epoch and $9\times$ more memory-efficient than GaLore on a tested six-encoder transformer with single-batch training. Our method also shows $\sim 2\times$ speedup than standard pre-training on a BERT-like code-generation LLM while achieving $4.23\times$  compression ratio in pre-training.
With further HPC optimization, CoMERA may reduce the pre-training cost of many other LLMs.  An implementation of CoMERA is available at <https://github.com/ziyangjoy/CoMERA>.

CoMERA: Computing- and Memory-Effiiiet Training via Rank-Adaptive Tensor Optimization
Zi Yang University at Albany, SUNY zyang80albany. edu
Ziyue Liu University of California at Santa Barbara ziyueliubucsb.ed
Samridhi Choudhary
Xinfeng Xie Meta xinfeng@meta.ccm com
Cao Gao
Amazon Alexa AI samridhc manrdhheeamzzon.o com
Meta caogao@meta. com
Siegfried Kunzmann Amazon Alexa AI kunzman@amazon. com
Zheng Zhang
University of California at Santa Barbara zhengzhanggece. ucsb.ed.
Abstract
Training large Al models such as LLMs and DLRMs costs massive GPUs and computing time. The high training cost has become only affordable to big tech companies, meanwhile also causing increasing concerns about the environmen- tal impact. This paper presents COMERA, Computing- and Gemory-Efficiee training method via Bank-AAaptive tensor optimization. CoMERA achieves rank- adaptive compressed (pre)-training a multi-objective optimization for- mulation and improves the training to provide both a high compression ratio and excellent accuracy in the training process. Our optimized numerical computation (e.g., optimized tensorized embedding and tensor-network contractions) and GPU implementation eliminate part of the run-time overhead in the tensorized training on GPU. This leads to, for the first time, 2 3x speedup per training epoch compared with standard training. CoMERA also outperforms the recent GaLore in terms of both memory and computing efficiency. Specifically, COMERA is 2 faster per training epoch and more memory-efficient than Galore on tested six-encoder transformer with single- batch training. Our method also shows 2x speedup than standard pre-training on a BERT-like code-generation LLM while achieving 4.23 compression ratio in pre-training. With further HPC optimization. CoMERA may reduce the pre- pre-traing trainini cost of many other LLMs. An implementation of CoMERA is available httpps: teattps://githen giggth..b....... www.wooaawoooo/////oooooooooooao CCMMER
Introduction
Deep neural networks have gained success in solving numerous engineering problems. These approaches usually use huge number of variables to parametrize network, and require massive hardware resources to train the model. For instance, the Deep Learning Recommendation Model (DLRM) released by Meta (which is smaller than the practical product) 32 has 4.2 billion parameters; GPT-3 has 175 billion parameters. OpenAl shows that the computing power required for key AI tasks has doubled every 3.4 months ① since 2012. Training large language model like ChatGPT and L. from scratch often takes several weeks or months on thousands GPUs (3.22 2
The majority of this work done when the first author postdoc UC Santa Barbara.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

100
ITT.
CIMERA
Caloree
1IT
inemoo Semmory (MB)
10022
time (min)
(4400
5961
771
40-
5000
4130
2000
4532
20
142
19.3
15.8
165
112
155
A3
NIA
NA
12
64
128
32
64
128
batch size
batch size
Training time per epoch.
(b) Peak memory consumption.
Figure 1: Training time and total memory cost of CoMERA, GaLore 41 and LTE (20 on six- encoder transformer with varying batch sizes. The experiment is done on Nvidia RTX 3090 GPU.
Large Al models often have much redundancy. Therefore, numerous methods have been developed to reduce the cost of AI inference [5] [13 16, 31 30 [7]. However, training large Al models
(especially from scratch) remains an extremely challenging task. Low-precision training 19 27 12
[36] has been popular in on-device setting, but its memory reduction is quite limited. Furthermore, it is hard to utilize ultra low-precision training on GPU since current GPUs only support limited precision formats for truly quantized training. Chen [6] employed the idea of robust matrix factorization to reduce the training cost. Similar low-rank matrix approximation techniques have been applied to train large AI models including large language models [4] 20). Among them, GaLore [41] can train the 7B LLaMA model on an RTX 4090 GPU using a single-batch and layer- wise setting. However, this setting lead to extremely long training time, which is infeasible in practical settings.
Compared with matrix compression, low-rank tensor compression has achieved much higher com- pression ratios on various neural networks [[] 22
25
29
40 43]. This idea has been studied post-training compression [1, fixed-rank
33
37
in structure search for compact representations
training 55 33 24 zeroth-order training 42 and parameter-efficient fine tuning 33 The recent
provides a rank-adaptive lensor-compressed training from Bayesian perspective. However, to achieve reduction in both memory and training time (especially transformers), two open problems need to be addressed. Firstly, more robust rank-adaptive tensor-compressed training model is desired, since the method in [4 [55 relies on heuristic fixed-rank warm-up training. Secondly, while modern GPUs are well-optimized for large-size matrix computations, they are un- friendly for low-rank tensor-compe compressed training. Specifically, most operations in tensor-compressed training are small-size tensor contractions, which can cause significant runtime overhead on GPUs though the theoretical computing FLOPS is very low. As a result, far have reported real training speedup on GPU. This issue was also observed in [1]] SVDinsTNN 43 controls tensor ranks search for compact tensor structure of given tensor. HEAT [I] tensor decompositions for post-training model compression of trained models. Detailed comparisons with these works are shown in Appendix A.l
work 14 15
know papers
Paper Contributions. In this work, we propose CoMERA, a tensor- compressed training method that can achieve, for the first time, simultaneous reduction of both memory and runtime on GPU. Our specific contributions are summarized as follows.
Multi-Objective Optimization for Rank-Adaptive Consor-Compressed Training. We propose a multi multi-objective optimization formulation to balance the compression ratio and model accuracy and customize the model for specific resource requirement. One by-product of this method the partial capability of automatic architecture search: some layers are identified as unnecessary and can be completely removed by rank-adaptive training.
Performance Optimization of Pensor-Compressed Training. While comsor-compressed training greatly reduces the memory cost and computing FLOPPS, it often slows down the practical training GPU. We propose three approaches to achieve real training speedup: ① optimizing the lookup process of tensorized embedding tables, 2 optimizing the sensor-network contractions in both forward and backward propagation, 3 eliminating the GPU backend overhead via CUDA Graph. Experimental Results. We evaluate our method on the end-to-end training of transformer with six encoders and the deep learning recommendation system model (DLRM). On these two benchmarks, our method achieves 80x and 99x compression ratios respectively, while maintaining the testing accuracy of standard uncompressed training. CoMERA also achieves 2 3x speedup per training epoch compared with standard training methods on the transformer model. In preliminary study of L.LM pre-training, COMERA shows 1.9x to 2.3x speedup in different pre-training stages on CodeBERT [21. while achieving 4.23x overall model reduction in the pre-training process.

---

melri--vector
matrix--matrix multiplication
multiplication
Order-1
Order-2
Order-3
Order-4
multi-sensor contraction
contraction
Figure 2: (a) Tensors. (b) Tensor contractions.
Figure 1 compares COMERA with GaLore [1 and the recent LoRA-based training method LTE (20 on the six-encoder transformer. When data and back-end memory cost are considered, COMERA's memory consumption is 9x less than GaLore in the single-batch training adopted in [1]] and it uses the least memory under all batch sizes. Our method is 2 3x faster than GaLore and LTE in each training epoch, although COMERA has yet been fully optimized on GPU.
While this work focuses on reducing the memory and computing cost of training, it can also reduce the communication cost by orders of magnitude: only low-rank tensorized model parameters and gradients need to be communicated in a distributed setting. The CoMERA framework can also be implemented resource-constraint edge devices to achieve energy efffcient on-device learning.
Background
A E R X7IXxXX-.-XHA is indexed as A = (dip-in)) and said of tensor A is defined as -n Pi In tensor networks, the order-d tensor A is represented node with edges. Some tensor network representations are illustrated in Fig. (a). Tensor Contraction. Let A E T. xan.xmaxx--ae Xm2x--X4 and B I8 xl2x......... xfm be two tensors with n,, = 4- The tensor contraction product C A ee, B has dimension Hig Hgf and the entries are
The tensor 26
28
to have order and dimension R1,--.; Rdd- The Frobenius
||||
Tha
f=
This definition can be naturally generalized to multiple pairs. Figure 2(b) illustrates some tensor contractions. For general operations among multiple tensors, we use Py'oorc einsum the following
einsum Sm = T, ,A, Am.),
(2)
where each S; is a string of characters that specifies the dimension of A,, The output tensor B is obtained by summing over all other dimensions that are not in T. In the following, we show a few commonly used einsum operations. The Tensor- Train decomposition as in (3) is
A = cinsum(nuri.
Gal). 1d-1nd R1712 ná, [9, ,
For the batched matrices A E pbxmxk B E bx the batched matrix multiplication is AB einsum (bmnk, bkn = bmn, [A, B)), where Cli, :,:] A(i :..
Tensor Decomposition. In this paper, will mainly tensor-train (TT) 35 and tensor-train matrix (TTM) (34) decomposition for compressed neural network training. TT (35 represents the
tensor A E I XmyXX-
set of small-size cores 91,..., Ga such that g: R Xn, xr, and
A 91 X3,1 G2 X3,1 X3,1 Gs.
(3)
The tuple (ro. ........ra is the TT rank the TT decomposition (3) and must satisfy To Ta TTM considers an order-2d tensor B of dimension
n2 X
x na, and represents B as
B = F, X F2 X11 X4,1 Fa.
(4)
where F, R5-1 X00-11XmXxxxx for 1...., d and rn Fd 1. Figure shows the tensor-network representations of TT and TTM decomposition.
In tensor-compressed neural networks, large weight matrices are reshaped to high-order tensors and compressed into small tensor cores in TT or TTM format. The weights of linear layers are often compressed into the TT format due to its efficiency in tensor-vector multiplications. The TTM format more suitable for embedding tables whose dimension is highly unbalanced.
3

---

Rd-1
Ri
Y22
Fld-1
map-1
(a)
(b)
Figure 3: Tensor networks for (a) tensor-train and (b) tensor-train-matiix decompositions.
3
The CoMERA Training Framework
The size of the lensor-compressed neural networks can be adjusted by modifying the tensor ranks. However, it also brings in an important problem: how can we determine the tensor ranks automatically for given resource limit? We propose multi-objective optimization address this issue.
3.1 Multi-Objective Training Model
Modified TT Representation. We consider the tensor-comppessee training for generic neural
network. Suppose that the neural network parameterized f(x] g4 -2 1=))) where {9}
RD xxn xr3 }dz compress the original weight W, Let {xx, } be training data and L be the loss function. The training is to minimize the following objective function
N
min
L
Σ
S(xa
1 (ar))
(5)
(9)
We modify the TT compression and control the ranks of gi..... Gú. by set of diagonal matrices [Dj e RRix j=|- Specifically, let W; be the reshape of W, and the compression of W, is
W: gi X3,1 D X2,1 9 X3,1 X3,1 1,1 D d X2,1 Gú.
(6)
Now the tensor cores for W; have S, n| D llo + n²y D d, llo
j==
d4
n2 D
1o
D
llo
variables. For simplicity, we denote = (9)... .... 1i and = (Di D 12
Multi-Objective Optimization. We intend to minimize both the loss and compressed network size, which can be formulated as multi-objective optimization ming.D (L(g. (I, ((D)) where
S(D) Σm= S, S,D). In most simultancously. Therefore, we look for a Pareto point (g', DD), meaning that there exist no g and
cannot find point that minimizes the loss and model size
D such that Lig, D) s L(g", D*), S(D) S(D)), and at least one of inequalities is strict.
3.2 Training Methods
We convert a multi-objective optimization to single-objective one via scalarization. We use different scalarization methods at the early and late stage of training. The late stage is optional, and it can further compress the model to enable efficient deployment Early Stage. At the early stage of CoMERA, aggressively pruning ranks dramatically hurts the convergence. Hence, we start the training with the following linear scalarization formulation g
cesource-constraint platforms.
min g.D
L(g, D) (D))
It is still hard to solve 7 since S(D) uses || which is nonsmooth. Therefore, we replace || llo by the norm and get the convex relaxation
Š(D) = Σ
D la n/a D
lla + m |
We note that S(D) can be arbitrarily close to 0 while keeping L D) unchangeable, since the corresponding slices of TT factors be scaled accordingly. Therefore, a direct relaxation of the scalarization (IT) does not have minimizer. To address this issue, we add an lz regularization 2 EP | to the relaxation and get the formulation
min 9.D
L D) + S + ||
The optimizer of Problem 9 is Pareto point for constrained problem, shown in the following.

---

Proposition 3.1. For all > 0, > O, there exists some constant C. > O such that the solution to the problem 9 is a Pareto point of the following multi-objective optimization problem
min 9.D
(L,
subject to || C.
(10)
Proof. See Appendix A.2 for the complete proof. Late Stage (Optional). The carly-stage training can provide us with a Pareto point, but we cannot may continue training the model towards a preferred loss Lo and a preferred model size So for deployment requirements. This can be achieved by the achievement scalarization
control where the Pareto point is. In the late stage of CoMERA,
8 that leads to Pareto point close So):
min max {wi (L( - Lo), w2(S(D) So)} } + p(L(G, + S(D)). G.D
(11)
Here W1, > O scale the objectives into proper ranges, and > 0 is small constant. After relaxing S(D) to S(D) and adding the regularization term, we get the following problem
min max (wr (T(G, D) To), (S(D) -So) + p(L D) + S(D)) 3| S.D
(12)
where 3 > is positive constant. Note that the S(D) inside max is not relaxed now for accurate comparisons. When wi (L(G, Lo) 2 w2(((D) So). we consider the following problem
min wi (LG,D)) Lo) + p( (GL D) S(D)) ?
(13)
g.D
We rum a step of a gradient- muddentt-based algorithm on this problem. When (L(G Lo) < w2(((D)
So), we relax the SDD again and get the following problem
min w2(S(D) So) + p(L(G D) + S(D)) 9 S.D
(14)
and run step of gradient-based algorithm on this problem. The Algorithm T is summarized in Appendix A.3 The late stage optimization can be independently applied to a trained tensor- compressed model for further model size reductions.
4 Performance Optimization of COMERA
While COMERA can greatly reduce training variables and memory cost, the low-rank and small-size tensor operations in COMÉRA are not efficiently supported by GPU. This often slows the training process. This section presents three methods to achieve real training speedup on GPU.
4.1 Performance Optimization of TTM Embedding Tables.
Embedding tables are widely used to transfer dis- crete features into continuous hidden space. The row size embedding tables is usually much larger than the column size, making TTM compression more suitable than the TT format. In the following, we use an order-4 TTM embedding table illustrate how to accelerate the lookup process.
selest infiess
We consider embedding table T E R A look- up operation selects the submatrix T[Z::] R for the index set {ix Jam This operation fast and inexpensive. However, the full embedding table itself extremely memory- neemoryconsuming. Suppose that m m m2mg?44 = Mingngg44 then we reshape represent it in TTM format
lock
butched GEMM
Figure 4: Optimized TTM embedding lookup.
T into tensor T E m x Xmxx4xx, and
9 X4,1 G2 X4,1 G3 X4,1 G1-
(15)
The compressed embedding table does not have the matrix T explicitly. We convert each row index 4 I to tensor index vector (2f, 2 21 24) and denote Z, {2/11 1 then T can be computed by contracting the tensors (G Z, :. 1 where each has size rz-1 x b x n, x ny. The
[.[:, Z,,;; :] stores many duplicated values especially when the set I is large. Therefore, directly computing the tensor contractions can cause much computing and memory overhead.
5

---

covirgroood optinited orde
527M
5294
5294
memmory (MBS)
villee opt cadalaa
rime
4.01
3100
2066 2500
201 2.12
2000
1763
0.56
644
D-47
1.:1
0,82 642
161
TW SU1
10000
30000
40000
10000
20000 butch sizze
40000
batch sizee
(a) Speed-up
(b) Memory reduction
Figure 5: Performance of optimized TTM embedding table lookup. The labels uncompressed, pro- posed approach, optimized order, unique indices, without optimization represent standard embedding with sparse gradients, the new method in 4.1 the method that only the unique order, the method that only uses the unique indices, and the method without optimization, respectively.
We optimize the tensor contraction by eliminating the redundant computation at two levels. Row- index level. We construct the index set I, = {is) containing all unique indices in I. We can easily obtain :] from T :I. Tensor-index level. The reduced index set Zu leads to 2k associated tensor index vectors (=f 22 23 )4 but at most m m2 pairs of (z 22) and m3M4 pairs
of (2f, x are unique. For instance, (2,3,1.3) and (2,3,2,4) are common in (2.3), ), so we only compute (2,3) entry once. Therefore, we can consider all unique pairs (2 25) and c xf) and compute
A einsumm (rom11111, Mim2222 (mimg)(M172)rs) r2, [91, 92));
(16)
A -insum (r2mgnars, r3m4n444 = (2(mm34)((nmm. [9,, G4l).
(17)
For each in E Iu, let ( be the coordinate of is for size (m1 mmm.), We denote Ji =
Uilk=1 and J2 1 16 === then compute the unique rows of T as
T | = einsum (b(m)12)r2) (acrrr,rrbbbnu r2b(g)4)) = b(1in2nan4) [A, A J2.:))).
(18)
Figure summarizes the whole process of TTM embedding table look-up. This approach can be easily applied to higher-order embedding tables by first grouping some small tensor cores to obtain intermediate tensors and then utilizing them to compute unique row vectors.
Performance. We demonstrate the optimized TTM embedding tables on single RTX 3090 GPU, We consider an embedding table of TTM shape [[80, 50, 54, 50], [4, 4, 4, 2]] and rank 32, extracted from a practical DLRM model. shown Figure our proposed method achieves about 5x speed-up up and 2-3 x memory saving than the standard TTM embedding without any optimization. The uncompressed embedding with sparse gradients is faster than our approach since our TTM embedding table requires extra computation, but it uses much more memory than the TTM embedding table.
4.2 Contraction Path Optimization for TT-Vector Multiplications
Next, we optimize the forward- and back- propagation of linear layers in the TT format. We consider the linear layer Y XW, where Y E 1oxXN W E N X N1 The W is compressed into the Tensor-Tren Train format: W ......., G2a] E I X--X2 where G; E IR Xo,-ixx;xrr and Ni nlg N₂ 1d4 722.. The forward- propaagation in the einsum form is
Y= XW einsum(bn; (onsumm(nn....ae ... n.. .,,..., S24 + bm441..... R22d 1 S....... G2a))
(19)
where E bxxn xna xna the reshaping of and S, denotes r;- n,r, Suppose that the gradient to Y is gy. then the gradients to G: and X can be computed as follows:
gg. cinsum (bm) ... ...nd, na, br n22, Si..... S1
S25 + S,,
X. I1..... 9. 9.4............... (2u)). gx cinsum (bma+).... ... n22, Si,..., S2a - bn m na, [gy. 91,..., 92a)).
(20)
(21)
In total, 2d + 2 contraction sequences are needed for the TT-format forward- and back- propagation. To reduce the computational costs, critical to find an optimal or near-optimal contraction path.

---

Large batch case. We denote A, := 91 x ... x
X x... x Gu. B: Ga+1 X Gati: = 924-it+ x x G2a. which are all computed sequentially. In practice, we only need to compute A.. A d. Ba, B and store the intermediate results. The 19h then computed the following way
Ti = einsum (Im .., nara = bra, {X, Adl)
(22)
y cinsum((ry Tafla+ ... ..2d + bn.... ... nid [Ti, Bal).
(23)
In backward propagation, the gradients computed in the following way:
The gradient gx is computed as
U1 einsum (onuuumbnn++1.... m2dd rand+1 ... R2dd bra: |gy, Bal)
(24)
gx
Cinsum(bras Ti... Raldd d + bn .nd,
nd, [u1, Anl).
(25)
'The gradients gg. for i 2 d+ can be computed as
T2 cinsum(br: bn ... .2d = rand+1 ... M2d [T,, gyy)
(26)
89. einsum(ramd
122d, 1 rddd+1 ... 10:--11-111 7:74+11 Rodd
(27)
Ti-1 [T2, 1- B (24-1)))]
Similarly, the gradients 8g, for I d can be computed as
U2 einsum (bra, bn ........ Ad + ramh ... na, (() gg. = cinsum (ram na, n1 -4-1 7:0++ Hardd
(28)
(29)
11-11 [22 A:, A (d-:)))-
The contraction paths forward- and back- propagation are summarized in Appendix A.5 Analysis. The proposed empirical path is near- optr-oltimal for large batch sizes. The following result analyzes the contraction path for forward-propagatton propagation.
Proposition 4.1. Suppose that the TT ranks satisfy ro <1 S ra rd-1 ra-1>....
and the batch size b is large enough. There exist groups {5: where S. = (9 9 containing consecutive lensor cores for 0 ÿ1 <...<<. jj < Jb+1 2d. Then, the contraction path with the least number of flops for the forward-propggation (19 first contracts the tensor cores in each St obtain V; with dimension r;; X Rij+11 tensor X with tensors {V:)/=. in the sequential order.
1 X fa1 and then contract the impul
Proof. See Appendix A.4 the complete proof.
Proposition 4.1 limplies that the optimal path first contracts some consecutive tensor cores and then contracts obtained tensors with the input tensor sequentially. The groups {S. depend on the dimensions, ranks, and batch size. The proposed contraction path satisfies the property shown in Proposition 4.1 and has flops roughly b(n1 ... ny fld+1 R2a)ra. The optimal contraction path has flops about bn: ... nach + bma+1 ... 02dC2, where C1, C2 are some constants. Hence, the proposed near-optimal and has comparable complexity to the optimal path. Suppose the optimal path is different from the proposed empirical path. Then the optimal path will likely involve a few more large intermediate tensors, which pose more memory costs during training and inference especially for static computational graphs. The empirical path is good choice to balance time and memory consumption. Similar arguments can be applied to the contractions for back pack-propagation.
When the batch size is small, the optimal path may have much fewer flops. However, the exccution time is almost the same the proposed path since all the operations too small. Hence, we can use the proposed path for most batch sizes. See Appendix A.6 for more analysis.
4.3 GPU Performance Optimization via CUDA Graph
While COMERA consumes much less computing FLOPS than standard uncompressed training, it can be slower on GPU if not implemented carefully. Therefore, is crucial to optimize the GPU performance to achieve real speedup. Modern GPUs are highly optimized for large-size matrix multiplications. However, the small-size tensor contractions COMERA are not yet optimized on GPU and require many small-size GPU kernels, causing significant runtime overhead. During the training. Cuda Graph launches and executes the whole computing graph rather than launching a large number of kernels sequentially. This can eliminate lots of back-end overhead and lead to significant training speedup. is more suitable for CoMERA since tensor-compressed training has much more small kernels than uncompressed training. This is just an initial step of GPU optimization. We expect that a more dedicated GPU optimization can achieve a more significant training speedup.
7

---

Table Result of Transformer on MNLI of batch size 128.
validation
total size (MB)
compressed size (MB)
uncompressed training
62.2%
256 (1x)
253 (1x)
CoMERA (early stage)
63.3%
5.9 (43x)
3.4 (74x)
CoMERA (late stage), target ratio: 0.8
62.2%
4.9 (52x)
2.4 (105x)
CoMERA (late stage), target ratio: 0.5
62.1% %
3.9 (65x)
1,4 (181x)
CoMERA (late stage), target ratio: 0.2
61.5%
3.2 (80x)
0.7 (361x)
Table 2:
The change of ranks of layers in the fifth encoder block.
before training
early-stage
stage rank
Q-layer in attention
(12, 30, 30, 30, 12)
(12, 30, 30, 30, 12)
00,0, 0,0,
K-layer in attention
31,,3030, 30, 12)
(12, 30, 30, 12)
00.0.0.0.0)
V-layer in attention
(12,30, 30, 12)
(12, 30, 30, 12)
(9, 11,
FC-layer in attention
(12, 30, 30, 30, 12)
(12, 30, 29, 30,
(9,8, 10,8,8,
#1 linear-layer in Forw-rorward
(12, 30,30,00000 30, 16)
(0.0.0.0.0)
(0.0.0.0.0)
linear-layer in Feed Forward
(16,30, 30, 30, 12)
(0,0,0,0,0)
0,,0,0,0,0)
Training Results
In this section, we test the performance of CoMERA on a few benchmarks, including domain- specific LLM. Our experiments are Nvidia RTX 3090 GPU with 24GB RAM.
5.1 A Medium-Size Transformer with Six Encoders
We first consider six-encoder transformer. The embedding tables and all linear layers are repre- sented as tensor
0.64
validation accurccy
in the training process as detailed in Appendix AA.7 We train this model
0.60
the MNLI dataset (38 with the maximum sequence length 128 and compare the accuracy, resulting model size, and training time of CoM- ERA with the standard uncompressed training. CoMERA Accuracy and Compression Perfor- mance. Table I summarizes the training results. The carly-stage training of CoMERA achieves 74x compression ratio on all tensorized layers, and the validation accuracy is even higher than the uncompressed training. Figure 6 shows the validation accuracy of CoMERA. In the late stage of COMERA, we set different target compression ratios for more aggressive rank pruning. The target compression ratios for the tensorized layers rather than for the whole model. The late-stage training can reach the desired compression ratio with very little accuracy drop. The smallest model has compression ratio of x0 for the whole model due to a 361 x compression on the tensorized layers with slightly worse accuracy.
0.58
0.56
incomppoeed tniring
LOMERA
2
6
10
14
16
18
spoch
Figure 6: Behavior of early-stage CoMERA train-
ing on the MNLI dataset.
Architecture Search Capability of CoMERA. major challenge in training is architecture search: shall we keep certain layers of model? Interestingly, COMERA has some capability of automatic architecture search. Specifically, the ranks some layers become zero the training, and thus the whole layer can be removed. For the target compression ratio 0.2, the whole second last encoder and some linear layers in other coders are completely removed after late-stage rank-adaptive training. The change of ranks of
CUMBRA w/ CLDMGngh
sacceperseel CUBAMisph wit C.lDMiigg
25
CAMBBA ELDAingh
Eine (mis)
14.3
15.6
16.1
15.3
6.4
7.4
32
batch size
Figure 7: Training time per epoch for the six-
layers in the 5th encoder is shown in Table [
encoder transformer model on the MNLI dataset.
Training Time. As shown in Figure 7 COMERA with CUDAGraph achieves around 3x speed-up than uncompressed training. COMERA without CUDAGraph can take much longer time in small batch-size setting due to the launching overhead of too many small kernels. The uncompressed training with CUDAGraph takes longer time than the one without CUDAGraph. This is because CUDAGraph requires all batches to have the same sequence length, and the consequent computing overhead more than the time reduction of CUDAGraph. In contrary, CoMERA has much fewer computing FLOPS and the computing part accounts for much smaller portion of the overall

---

250
COORIAA
CLABREA
1500
smemory (MD)
30000
time (G)
0900
15000
3UT
1gc
10000
$229
300-
3A17
1000))
2000
48000
SC000)
2000 butch size
40000
batch size
(a) Training time per epoch.
(b) Memory cost.
Figure 9: Performance of optimized COMERA on training DLRM.
runtime. Empirically CoMERA is 2 - 3x faster in the whole training than uncompressed training for transformers on a single GPU, but we do not have theoretical guarantees about the number epochs although they similar in our experiments. Appendix [A.8 provides provvide more details about the run-time comparison on this benchmark, showing that COMERA is still faster than standard training even the compression ratio is close
1.
5.2 A DLRM Model with 4-GB Model Size
We further test COMERA DLRM 32 released by Meta Criteo Ad Kaggle dataset [23]. We
compress the ten largest embedding tables into the TTM format as in Section 4.1 All fully connected 128 are compressed into TT format. The model is trained for epochs.
layers with sizes
Effect of Optimized TTM Embedding. The training time per epoch and peak memory cost are shown in Figure 9 Our optimized TTM lookup speeds up the training process by around 2x and remarkably reduces the memory cost by 5X.
normal CE
0.900
Overall Performance of CoMERA. Table 3 shows the testing accuracy, testing loss (measured normal- ized CE). memory costs, and model sizes of CoMERA and uncompressed training. CoMERA achieves similar accuracy as the uncompressed training, while COMERA compresses the whole model by 99x and saves 7x peak memory cost (with consideration of the data and backend overhead) in the training process. The reduction model size and memory cost mainly comes from the compact TTM tensor representation of large embedding tables. Standard uncompressed training is faster than CoMERA since DLRM is an embedding intensive model, and the computation in the embedding table look-up rather than matrix multiplications. However, CoMERA uses much less memory, saving 6.9X, 4.8X, and 3.1X memory for batch sizes 10000, 2000, and 4000, respectively. Furthermore, COMERA has a similar convergence and needs fewer iterations than standard training for DLRM, as shown in Figure
2000
4000
6000
5000
iteration
Figure 8: NCE loss curve of DLRM on the validation dataset.
5.3 Comparison with GaLore and LTE
We compare method with two recent low-rank compressed training frameworks: GaLore [41 and LTE (2)- GaLore [41] re- duces the memory cost by performing SVD compression the gradient, and LTE rep- resents the weights as the sum of parallel low-rank matrix factorizations. We evaluate their memory costs and training times per epoch on the six-encoder transformer model under different batch sizes. We do not compare the total training time because the training epochs of various methods are highly case-dependent. The COMERA and GaLore achieve almost the same validation accuracy. 64%, on the MNLI dataset. However, the LTE approach does not converge on the task using its default setting.
Table 3: Training results the DLRM model with batch size 10, 000.
CoMERA uncompressed
accuracy normalized CE
78.68%
78.76%
0.793
0.792
model size (GB)
4.081
0.041 (99x)
peak memory (GB)
18.275
2.612 (7x)
Training Time Per Epoch. We use rank 128 for the low-rank gradients in GoLore, and rank 32 and head number 16 for the low-rank adapters in L.T. For fair comparison, all methods are executed with CUDA graph to reduce the overhead of launching CUDA kernels. The runtimes per training epochs are reported in Figure a)) For the LTE, we only report the results for batch sizes 32, 64, 128 since it requires the batch size to be a multiple the head number. Overall, our COMERA is about 2x faster than GaLore and 3x faster than LTE for all batch sizes, because the forward and backward propagation using low-rank censor-netword contractions dramatically reduce the computing FLOPS.

---

COMERA
CodeBERT
Trairing Lozs
Phase
Phase
Steps
Figure 10: Pre-training loss curves of CodeBERT and COMERA.
Memory Cost. Figure (b) shows the memory cost of all three training methods. In the single- batch setting as used in [41] our CoMERA method is 9x more memory-efficient than Galore on the tested case (with consideration of data and back-end cost). As the batch size increases, the memory overhead caused by data and activation functions becomes more significant, leading to less memory reduction ratios. However, our proposed CoMERA still uses the least memory.
We run the experiments on the RTX 3090 GPU. The work GaLore [1 uses the RTX 4090 GPU for experiments, also compare them the RTX 3090 GPU. The results are in Appendix A.9
5.4 Preliminary LLM Pre- Training Results: Case Study on CodeBERT
To show the benefit of CoMERA in pre-training (domain-specifie LLMs, we follow the setup from CodeBERT (TO) to pre-train BERT-like model for code generation. The pre-traiming dataset is the CodeSearchNet 221, a collection of 2M (comment, code) pairs and 6M pure code sequences from open- libraries with 6 types of programming languages. We pre-train Code large (357M) and its CoMERA (84M) variant using the masked language modeling (MLM) objective and compare their training loss Figure TO We achieve up to 12.72x compression on tensorized layers and 4.23x overall compression with final loss of 0.40 vs 0.28. There is a small gap between the final losses. However, this does not necessarily imply performance degradation on downstream tasks, based our observation on BERTLARG LARGE ( shown in Appendix A. Furthermore, CoMERA is 2.3x and 1.9x faster than standard pre-training Phase and Phase respectively, when evaluated on the Nvidia RTX 3090 GPU. Our current CoMERA implementation is still slower than standard pre-tra pre-training on massive GPUs, since no performance optimization has been done on HPC.
Conclusions and Future work
This work has presented CoMERA framework to reduce the memory and computing time of training AI models. We have investigated rank-adaptive training via multi-objective optimization to meet specific model sizes while maintaining model performance. We have achieved real training speedup on GPU via three optimizations: optimizing the tensorized embedding tables, optimizing the con- traction path in tensorized forward and backward propagation, and optimizing the GPU latency via CUDAGraph. The experiments on transformer model demonstrated that COMERA can achieve 2 - 3x speedup per training epoch. The model sizes of the transformer and DLRM model have been reduced by 43x to 99x in the training process, leading to significant peak memory reduction (e.g., 7x total reduction in large- large-batch training of DLRM on single GPU). Our method has also outperformed the latest Galore and LTE frameworks in both memory and runtime efficiency. More importantly, our method has demonstrated significant speedup and model compression in pre-training CodeBERT, domain-specific L.LM for automatic code generation. We have also observed further speedup by combining CoMERA with mixed- mixed-precisson computation. The discussions and some preliminary results are Appendix(A.I1. AAII]
Unlike large-size matrix operations, the small low-rank tensor operations used in CoMERA are not yet supposuppoteed by existing GPU kernels. The performance CoMERA can be further boosted significantly after comprehensive GPU and HPC optimization. The existing optimizers, e.g. Adam, are well-studied for uncompressed training. However, COMERA has very different optimization landscape due to the tensorized structure. Therefore, it is also worth studying the optimization algorithms specifically for CoMERA in the future.
7
Acknowledgement
The pre-training task used resources of the National Energy Research Scientific Computing Center, DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CHI123 using NERSC award ASCR-ERE ERCAP0030039.
10

---

References
[1] Al and compute.
OpenAl I in 2019.
[2] ChatGPT and generative AI
are
booming, but the costs
can be extraordinary.
https://www.cn cnbc. cot.://3///m//13//////m3
chatgpt- and- generative ai ai- are-
- how.onggg-...aaa.eeeeeeaaaaaaeaa
html Accessed: 2023-03-15.
[3] Update: ChatGPT runs
10K
Nvidia
training GPUs with potential for
thousands more.
http:://ww..ierreeeeeees www.reeeeeeeeeeeeiccc....mm com/ /senooss https://wwx
chatgpt-runs 10k tww.magggess--eeecaaaaaaoooooaao
Accessed:
2023-03-15.
[4] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. arXiv preprint Siti:::000.141666 2020.
[5] Giuseppe G Calvi, Ahmad Moniri, Mahmoud Mahfouz, Qibin Zhao, and Danilo P Mandic. Com- pression and interpretability of deep neural networks via tucker tensor layer. arXiv:1903.06133, 2019.
[6] Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Re. Pixelated Butterfly: Simple and efficient sparse training for neural network models. arXiv preprint arXiv:21 2021.
[7] Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. Compressing neural networks with the hashing trick. International Conference on Machine Learning. pages 2285-2294, 2015.
[8] Kalyanmoy Deb, Karthik Sindhya, and Jussi Hakanen. Multi-objective optimization. In Decision sciences, pages 161-200. CRC Press, 2016.
[9] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proc. Association for Computa- tional Linguistics: Human Language Technologies, pages 4171 2019.
[10] Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng. Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. Codebert: A pre-trained model for programming and natural languages. arXiv preprint arXiv:2002.08155. 2020.
[11] Jiaqi Gu, Ben Keller, Jean Kossaifi, Anima Anandkumar, Brucek Khailany, and David Pan. Heat: Hardware-efficient automatic tensor decomposition for transformer compression, arXiv preprint arXiv:2211.166: 167.7 2022.
[12] Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In International Conference on Machine Learning, pages 1737-1746, 2015.
[13] Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural net- works with pruning, trained quantization and huffman coding. arXiv preprint 150.00051 2015.
[14] Cole Hawkins, Xing Liu, and Zheng Zhang. Towards compact neural networks via end-to-end training: A bayesian tensor approach with automatic rank determination. SIAM Journal on Mathematics of Data Science, 4(1):46-71, 2022.
[15] Cole Hawkins and Zheng Zhang. Bayesian tensorized neural networks with automatic rank selection. Neurocomputing, 453: 2021.
[16] Qinyao He, He Wen, Shuchang Zhou, Yuxin Wu, Cong Yao, Xinyu Zhou, and Yuheng Zou. Rffective quantization methods for recurrent neural networks. arXiv preprint arXiv: 16] (61.. 10176. 2016.

---

[17] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint 150..0233,.
[18] Oleksii Hrinchuk, Valentin Khrulkov, Leyla Mirvakhabova, Elena Orlova, and Ivan Oseledets. Tensorized embedding layers for efficient model compression. arXiv preprint (tth::991..1117. 2019.
[19] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran Ei- Yaniv, and Yoshua Bengio. Quan- tized neural networks: Training neural networks with low precision weights and activations. The Journal of Machine Learning Research, 18(1):68996888, 2017.
[20] Minyoung Huh, Brian Cheung, Jeremy Bernstein, Phillip Isola, and Pulkit Agrawal. Training neural networks from scratch with parallel low-rank adapters. arXiv preprint arXiv2240..16828, 2024.
[21] Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Marc Brockschmidt. CodeSearchNet challenge: Evaluating the state of semantic code search. arXiv preprint VrXiv:1909009430 2019.
[22] Max Jaderberg, Andrea Vedaldi, and Andrew Zisserman. Speeding up convolutional neural networks with low rank expansions. arXiv preprint arXiv::405..85 2014.
[23] Olivier Chapelle Jean-Baptiste Tien, joycenv. Display advertising challenge, 2014.
[24] Valentin Khrulkov, Oleksii Hrinchuk, Leyla Mirvakhabova, and Ivan Oseledets. Tensorized embedding layers for efficient model compression. arXiv preprint arXiv:1901.1078/ 2019.
[25] Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Compression of deep convolutional neural networks for fast and low power mobile applications. arXiv preprint arXiv: /5/ 1.055/5111.000.. 2015.
[26] Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review, 51(3):455-500 Aug. 2009.
[27] Urs Köster, Tristan Webb, Xin Wang, Marcel Nassar, Arjun Bansal, William Constable, Oguz Elibol, Scott Gray, Stewart Hall, Luke Horoff, et al. Flexpoint: An adaptive numerical format for efficient training of deep neural networks. NIPS, pages 1742-1752, 2017.
[28] Joseph M Landsberg. Tensors: geometry and applications. Representation theory, 380243): 2012.
[29] Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition. arXiv preprint arXiv:1412.6553, 2014.
[30] Jian-Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. In Proceedings of the IEEE international conference on computer vision, pages 5058-5066, 2017.
[31] Pavlo Molchanov, Arun Mallya, Stephen Tyree, luri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 1 11264-11272 2019.
[32] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sun- daraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. 2019.
arXiv preprint
[33] Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neural networks. In Advances in neural information processing systems, pages 442-450, 2015,
[34] Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. In Advances Neural Information Processing Systems 28, pages 442-450, 2015.
12

---

[35] Ivan V Oscledets. Tensor-train decomposition. SIAM Journal on Scientific Computing. 33(5):2295-2317 2317, 2011.
[36] Xiao Sun, Naigang Wang. Chia- Yu Chen, Jiamin Ni, Ankur Agrawal, Xiaodong Cui, Swagath Venkataramani, Kaoutar El Maghraoui, Vijayalakshmi Viji Srinivasan, and Kailash Gopalakr- ishman. Ulra-low precision 4-bit training of deep neural networks. NIPS, 33, 2020.
[37] Andros Tjandra, Sakriani Sakti, and Satoshi Nakamura. Compressing recurrent neural network with tensor train. CoRR, 0os//705.00803 2017.
[38]
Adina Williams, Nikita Nangia, and Samuel Bowman. broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pages 1112-1122. Association for Computational Linguistics, 2018.
[39] Yifan Yang, Jiajun Zhou, Ngai Wong, and Zheng Zhang. LoRETTA: Low-rank economic tensor-train adaptation for satra-lo--prrmeett fine-tuning of large language models. arXiv preprint arXiv:2402. 11417, 2024.
[40] Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-tral train recurrent neural networks for video classification. In Proc. the 34th International Conference Machine Learning. volume 70, pages 3891-3900, Sydney, Australia, 06-11 Aug 2017.
[41] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient LLM training by gradient low-rank projection, 2024.
[42] Yequan Zhao, Xinling Yu, Zhixiong Chen, Ziyue Liu, Sijia Liu, and Zheng Zhang. Tensor- compressed pack-propagationn-fre training for preprint arXiv:2308. 09858, 2023.
neural networks. arXiv
[43] Yu-Bang Zheng. Xi-Le Zhao, Junhua Zeng, Chao Li, Qibin Zhao, Heng-Chao Li, and Ting Huang. Svdinstn: A tensor network paradigm for efficient strueture search from regularized modeling perspective. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22254-26263,, 2024.
13

---

A Supplementary Material
A.1 Comparison with existing works
uses sparse diagonal matrices and the & regularization to control tensor ranks for a compact tensor structure of a given tensor. In contrast, our work compresses weights during end-to-end training without any prior information on the tensor (i.e., model parameters). Both works use sparse diagonal matrices and & terms to control tensor ranks. Using diagonal matrices to control the ranks of matrices is very common, like SVD. The & norm is also widely used to induce sparsity in various models, like compressed sensing and Lasso regression. is natural to combine these two techniques control tensor ranks, regardless of tensor formats. Moreover, problem as a more generic multi-objective problem and uses a two-stage algorithm to solve it. The formulation in SVDinsTNN is similar to linear scalarization approach in early stage. Our work further uses the achievement scalarization in the late stage to find a model close to our preferred model performance and size.
SVDinsTN 43
work formulates the
HEAT T also considers contraction optimization for post-training model compression of trained models. In contrast, CoMERA considers end-to-end ensor-compressee training, where model parameters are known prior to training. In addition, HEAT only discusses the single path optimization for forward propagation in CP format. We have optimized d + 2 contraction paths jointly in both forward- and back- propagation in TT format. Since these contractions can be coupled, we have also minimized the overall computation costs by reusing intermediate results.
A.2 Proof of Proposition 3.1
Proof. The objective function in 9 is bounded below by 0. Hence, the problem infimum value f". Let (g D*1 1=1 be a sequence such that lim L(g D* S(
(9) has a finite
3|| f". The sequence must be bounded because of the I regularization of D and the & regularization of G. As a result, the sequence has a cluster point (9", D*) which a minimizer of the 9 Let C = || The relaxation (9) is equivalent to the constrained optimization problem.
min G.D
L(G, D) + +D)
(30)
s.t.
| C.
implies that the solution to the training problem (9 Pareto point of the mbltivobjecttvv optimization problem (M
A.3 Algorithm for Late Stage Optimization in Section 3.2
The algorithm for the late stage optimization in Section 3.2 is summarized in Algorithm |
Algorithm 1 Solve relaxed scalarization problem
12
Input: Initializations Go, Do. constants Lo, So, W1, W2, p. 3, and an optimization algorithm O. Output: Tensor cores Gr and rank-control parameters Dy.
for t O.....T T do w1 (L D.) Lo) 2 w2(S(D)) - Su) then
The optimization algorithm O runs step the problem
(3)
else
The optimization algorithm O runs step the problem 14
end
end for
A.4 Proof of Proposition 4.1
Proof. For convenience, let V; be string of characters to specify the dimension of V., C. be the set of tensor cores used obtain V,, and X; the tensor by contracting X with V Vi..... V.. denoted by the string X;.
14

---

We first show that V1 must be in the proposed format. Suppose otherwise for contradiction. Let g
be the first tensor core used to obtain Vn. If 1, then we write Vi V V2² where the tensor V
corresponding to VT is obtained by contractions of longest consecutive tensor cores containing 9: in
the set C₁. Let Z V. X, Z2 V22 X. The number of flops for the contraction between X'
If we first contract X with V1 and then contract the obtained tensor with and Vi is x(21)=(2a)
(
VP, the number of flops is =(X)x(12) + ((2))=x(Z2)
(x(V) which is less than (V)))((+)))(() since ((2)))(22)
a(Z1)
Z2 + Vi2. It contradicts our assumption that this is the optimal path. If d 2 >> let S be the lensor generated by the longest consecutive tensor cores containing - and used in the optimal path. A better path is to first contract V1 with S obtain W. then contract W with X' and all other unused parts in the optimal path. It is better because the number of flops for contracting W and N no greater than that for contracting V1 and X' and the new path reduces the number of flops in the remaining contractions. The reduction in the remaining contractions is more than the potential flop increase in obtaining W when the batch size is big enough. Finally, we consider the case that i >> d. Let 9j be the first tensor core used to obtain 22. If j > d, then first contracting V1 and V2 and then all other parts is better choice. Otherwise, if j I d, let V be the tensor contracted by the consecutive tensors containing 9j in C2. The tensor V can be represented by the contraction of V and another tensor V generated by the remaining tensor cores in C2. In this scenario, we can contract V with V1 to get S, then S with X' to get W, then 13 with W, and finally the obtained tensor with V 3,- .3..... ... Vy- It is not hard verify contracting V and V1, S and S and X, and W and V uses less flops than directly contracting X' with V and 1/2 directly when the batch size b is large. Summarizing everything above, we can conclude that Vi must be in the proposed format.
The contraction of X; and Vi+1....: VR has the similar structure to the contraction of X and Vi...., Vg- By applying the same proof, conclude that the tensors V, must be in the format stated in the proposition and we will contract the input tensor X' with the tensors ten te........ V3 in the sequential order.
A.5 Algorithm for Contraction Path in Section 4.2
The empirical near-optimal contraction path for tensor-compressed training is shown in Algorithm
Figure II presents the tensor diagrams for contraction paths of TT forward- and back- propagation as discussed in Section 4.2
Algorithm 2 Empirical path for tensor-compressed forward- and back- propagation Forward Input: Tensor cores 9I...., Gzd and input matrix X. Forward Output: Output matrix Y, and intermediate results.
1: Reshape the matrix X to the tensor X.
2: Compute As, A_ di Ba, Bdd in the sequential order and store intermediate results of
[A., B. B 1d 1-= for back-propagattio.
3: Compute Tj as (22) to store it for propagatto.......
4: Compute y as in (23) and reshape it to the appropriate matrix Y.
Backward Input: Inputs of Forward, stored results from Forward, and output gradient gy.
Backward Output: Gradients gx :g.....: .g...
1: Reshape the gradient gy to the tensor gy.
2: Compute 44 and gx as in (24). (25) and store U4, for future use.
3: Compute Eg, for i 2 d + 1 as in (26). (27) using stored tensors.
4: Compute gg, for S d as in
using stored tensors.
A.6 Small Batch Case for Contraction Path in Section 4.2
Small batch case. The empirical contraction path in Algorithm 2 eeliminates the batch size dimen- sion b early, so it is nearly optimal when the batch size is large. We may search for better path using a greedy search algorithm to minimize the total operations. the pairs that output the smallest tensors. Such a choice can quickly eliminate large intermediate dimensions to reduce the total number of operations. When the batch size large, the searched path is almost identical to the empirical path in Algorithm 2 which eliminates the batch size dimension b
each iteration, we prioritize
15

---

Mell
Bad
91
Ad
Ba
Pael Rar
7i
8a
(a) TT- vector forward proppaatiootnionn
71g
Td+11
72d-1
51
Skd
Sid+1
Ad
B
gx
44
(b) TT-vector mack-propagation: compute gx
1h
R2
84
7
Gdi
5
Ti-1
P11
75g
rd
7
Sp.
A
A_
74
TT-vector back back-propagation: compute gG, Siss
91
G2
Gd+1
G
G
71-1
80.
8
B
(36-1)
B
(d) TT-vector back-propagation: compute gG,, 1sis 2d
Figure 11: Tensor diagrams for contraction paths of TT forward- and back- propagation.
early. The searched path may differ from Algorithm for small batch sizes, but their execution times on GPU are almost the same. This is because the tensor contractions for smaller batch sizes have minor impact on the GPU running times. Consequently, despite certain tensor contractions in the empirical path being larger than those in the optimal path, the actual GPU execution times between them exhibit only negligible differences. Therefore, the empirical contraction path Algorithm 2 is adopted for all batch sizes in CoMERA.
A.7 Compression Settings for the Experiment in Section 5.1
The compression settings for the experiment in Section 5.1 are shown in Table 4
16

---

Table Tensorized setting for the Transformer model in CoMERA.
Format
linear shape
lensor shape
rank
embedding attention
TTM
(30527,768)
(64,80,80660)
30
(768,768)
(12.8,888.8812))
feed-forward
TT
(768,3072)
(12.8.8.12.16,16)
30
compesssion
compesssion 1.5% compeession
compresion unoomproooo
20
18.5
time (min)
15.53
16.6
16.4
13.94
14.35
11.16
12.44
11.56
10-
9.7
7.2
7.79
6.4
6.21
32
128
batch size
Figure 12: Per epoch training time of COMERA on MNLI for various compression ratios.
A.8
Per Epoch Training Time of CoMERA on MNLI for Various Compression Ratios
Table 12 shows the per-epoch training time of CoMERA on MNLI dataset for different compression ratios. The acceleration is more obvious for larger compression ratios. When the compression ratio is greater than 1, CoMERA always has speedup. When the compression ratio approaches 1, the time of CoMERA approaches that of uncompressed training.
A.9 Comparison with GaLore and LTTE on single RTX 4090 GPU
single RTX 4090 GPU for experiments in the original paper, also the experiments on the RTX 3090 GPU and compare the results. Figure 13 presents the training time and peak memory consumption. Compared to RTX 3090, the training memory and takes less training time, and COMERA still the fastest method and consumes the least memory among all three techniques. The memory savings reported in Figure in our paper. The speed-up factors are almost identical for batch sizes 32, 64, and 128. For batch size 1, our method is 1.2x faster and 1.7x faster than GaLore RTX 3090, respectively. The difference is that RTX 4090 GPU significantly accelerates matrix multiplications of batch size 1, while it does not accelerate that much for smaller tensor contractions. We find that r = 30 matrix multiplication on RTX 3090 has a similar speedup for both batch sizes, whereas the same multiplication on RTX 4090 only has speedup for batch 32 and does have any speedup for batch 1. We would like to note that it might be caused by that different GPU platforms have different backend overhead, which can become more dominant as computation decreases to batch=1. We will continue optimizing GPU-level kernels to accelerate small tensor contractions and expect to see similar speedup.
Since GaLore 41
RTX 4090 similar
almost the
the results
A.10 COMERA Pretraining Result on BERTLARGE LARGE
CoMERA on Original BERTLARGE- Our results on CodeBERT rather preliminary as only pre- pre-training loss is available. For the original BERT [ (336M) and its CoMERÁ (125M) variant that we trained by using Wikipedia (2500M words), we achieve up to 6.36x compression on tensorized layers and 2.69x overall compression. with final loss of 1.45 1.26. On downstream tasks, COMERA ourperforms BERTLARGE on SST-2 (accuracy: 92.10% vs 91.74%) and MRPC (accuracy: 86.82% vs 86.00%) underperforms BERTLARGE SQuAD (fl: 88.7 76% vs 90.68%).
17

---

60
CoMERA
Caloree
1TT
more significant acceleration after memoyy (MB)
1006
time (min)
40
733
4636
20
5734
4A
111
34
6.
34
6M
182
NA
1791
32
64
128
128
batch size
batch size
Training time per epoch.
(b) Peak memory consumption.
Figure 13: Training time and total memory cost of CoMERA, GaLore 41 and LTE (20 on six-encoder transformer with varying batch sizes. The experiment done Nvidia RTX 4090 GPU.
Table Speed-up of seecedppeeision computation on tensor-compressed linear layers.
tensor-vector
matrix-vector
shape (b.m,n)
FP8-mix
FP32
HP8
FP32
(10000.1024.1024
1.95
1.63
1.02
2.82
1,0000,1024,1000
2.02
3.37
1.93
5.41
1.0000.1024,1022
2.55
6.82
4.27
10.93
1.97
3.96
3.69
10.28
(20000, 4024,406))
2.96
8.32
7.26
20.93
(40000,0024,409)
5.47
17.13
15.27
45.60
A.11 Discussion: Mixei-Precisoon CoMERA
Modern GPUs offer low- oow-precision computation to speed up the training and inference. natural to combine low-rank tensor compression and quantization to achieve the best training efficiency. However, CoMERA involves many small-size low-rank tensor contractions, and a naive low-precision implementation may even slow down the training due the overhead caused by precision conversions.
To resolve the above issuc, we implement mixed preccisio computation in CoMERA based on one simple observation: large-size contractions enjoy much more benefits of low preciisiooon computation than small-size ones. This is because the overhead caused by precision conversions can dominate the runtime in small-size contractions. In large- batch censor-compressed training, small- and large-size tensor contractions can be distinguished by whether the batch size dimension is involved. In general, contraction with the batch b is regarded large and is computed in low precision. Otherwise, it is regarded small and is computed in full- pell-sreciioon. The actual mixed- seeeddprecision algorithm depends on the contraction path used in the forward- and back- propagations of COMERA.
Runtime. We evaluate the mixed- precision forward and backward propaga- tions of CoMERA the NVIDIA L4 GPU. We consider single
validation accuracy
0.625
FP8 precision on
0.600
linear layer. The shapes (1024, 1024) and
0.575
(1024, 4096) are converted to the TT shapes (16, 8, 8, 8, 8, 16) and (16, 8, 8, 16, 16, 16) respectively, and the ranks both 32. The total execution time for 1000 forward and backward propagations
Mixed-prexision
0.550
FP32 COMERA
10
14
16
18
20
epoch
shown in Table 5) The FP8 tensor tensor-comppnt compressed linear layer has about 3x speed-up up compared to the FP8 vanilla linear layer when the batch size and layer size large. When the batch size is small, the FP8 vanilla linear layer is faster. This because the tensor-compressed linear layer consists of a few sequential computations that are not well supported by current GPU kemels. We expect optimizing the GPU kernels.
Figure 14: Convergence of mixed- preecsson CoMERA on the six-encoder transformer.
Convergence. We use the mixed-precision CoMERA to train the DLRM model and the six-encoder transformer. The result on DLRM is shown Table T The convergence curve of the six-encoder transformer is shown in Figure 14 The experiments demonstrate that the accuracy of FP8 training is
18

---

Table Training results of mixed- -reeiippreeiiio CoMERA on DLRM (batch size=10,,000
accuracy
normalized CE
HP32 COMERA FP8/FP32 mixed-precision COMERA
78.76%
0.792
78.88%
0.793
similar to FP32 training. However, we did not much acceleration of using FP8 in the experiments. This is mainly because of ( the computation overhead of slow data type casting between FP32 and FP8; 2 the sequential execution of small tensor contractions that not well supported by current GPUs; 3 the relatively small sizes linear layers in the tested models. We will investigate these problems in the future, and are optimistic to see significant acceleration larger models.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction accurately reflect the methods in Section 3 and Section and and the numerical results in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We mentioned that current GPU-level optimization is not optimal yet and there still exist lots of opportunities for improvement.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
20

---

Answer: [Yes]
Justification: The complete assumptions and proofs are provided in the Appendix, e.g.
Appendix A.. and Appendix A.6
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proof' in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper clearly describes the proposed methods in Section 3 and Section 4 The experimental settings are provided in Section 5
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
21

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We used the open-sourced data for experiments. An implementation of CoMERA is available at
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis.co cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Coaee//dddeSebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental settings are provided in Section 5
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The paper focuses on the computation-effiiient training. We mainly report the resource usage and training results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computer resources are provided in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform. in every respect, with the NeurlPS Code of Ethics
Question:
Answer: [Yes]
Justification: Our paper about new computational algorithms for efficient training of large models. It conforms with the NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our paper proposes new methods for efficient training of large models. The research has the possibility significantly reduce the training costs of large models and also accelerate the inference. It is discussed in Section [ad Section [
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper is about efficient training methods and poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, they are properly credited, mentioned, and respected.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
24

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper is about methods and algorithms for efficient training. No new assets are released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper is about methods and algorithms for efficient training. It does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: paper is about methods and algorithms for efficient training. does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
25

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "hjIKLHtxGz", "reviewer": "Reviewer_1f6L", "summary": "This paper leverages the tensor decomposition concept in model training and tries to address two questions: the first is how to enable rank-adaptive, and the second one is how to obtain real efficiency. This paper achieves a 2 − 3× speedup per training epoch compared with standard training with on-par accuracy.", "questions": "See weakness before.", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* An easy paper to follow\n* Although tensor decomposition is not new, they successfully enable real speed-up with several performance optimization techniques.\n* Real speed up and on-par accuracy on small models.", "weaknesses": "* The problem formulation in 3.1 and how to solve the l0 norm is not quite new. \n* Contraction Path Optimization is also not a new topic. The authors should cite previous work, e.g. [1] and discuss differences, and highlight the uniqueness.\n* Does this training take more time in terms of convergence? And could the author discuss more about the training difficulty of tensorized layers and normal layers.\n[1] Gu, Jiaqi, et al. \"Heat: Hardware-efficient automatic tensor decomposition for transformer compression.\" arXiv preprint arXiv:2211.16749 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "***Weakness 1**: “The problem formulation in 3.1 and how to solve the l0 norm is not quite new.”*\n\nResponse: Thanks a lot for the comment! \n* Our novelty is to formulate the problem of balancing performance and size as a **multi-objective** problem. The **linear scalarization** for gently pruning ranks and achievement scalarization for deployment requirements are applied in early-stage and late-stage to solve the problem. **It is an early try that this multi-objective optimization approach** is applied to satisfy deployment requirements. While the formulation invovles l0 and l1 regularization as in single-objective optimization, **the objectives are different from those in existing literature since the mathematical analysis is entirely different.** \n* Using the l1 relaxation alone can cause both theoretical and practical issues:\n    * First, the l1 relaxation cannot effectively control the tensor ranks because magnitudes of TT factors can keep growing. Hence, an extra l2 norm regularization term for the tensor cores is added. Mathematically, the new problem with the l2 norm regularization is equivalent to a new constrained multi-objective optimization problem (10) in the paper whose tensor cores are bounded. The equivalence and related theoretical analysis are shown in Proposition 3.1 and its proof. \n    * Second, the l1 relaxation does not correctly reflect the model size in the achievement scalarization problem in late-stage. Thus, we propose to use the l0 norm for the comparison between performance gap and size gap. This **l0-based measure is used as a switching condition** *between* Eq. (13) and Eq. (14), then the **l1 norm is used in numerical implementation** *inside* the solver of Eq. (13) and (14). This completely differs from an l1-relaxed single-objective optimization which does not require l0-norm any more and does not require switching.\n    \n***Weakness 2**: Difference from Contraction Path Optimization in [1] HEAT paper.*\n\nResponse: Thanks! The contraction optimization in [1] is very different from ours. We will cite it and discuss differences. The key differences are summarized in the following table and described below:\n* **Post-training compression VS compression in training.** The paper [1] considered **compression after training**, where trained model parameters were already given. Ours considers **end-to-end tensor-compressed training**, where no model parameters are known prior to training. It can reach computation savings and training acceleration, whereas post-training compression cannot. \n* **Different tensor formats.** The paper [1] uses **CP, Tucker, and Tensor-Train-Matrix** formats to compress linear layers. In contrast, we use the **Tensor-Train** format for linear layers, and TT matrix for embedding tables. \n* **One contraction path in forward VS d+2 coupled contraction paths in forward and backward.** The paper [1] only discussed the **single path optimization for forward propagation in CP format**. We have optimized **d+2 contraction paths jointly in both forward- and back- propagation in TT format**. Since these contractions can be coupled, we have also minimized the overall computation costs by reusing intermediate results.  Finally, we have also provided a **theoretical analysis**, Proposition 4.1, demonstrating the proposed path is near optimal for large batch sizes. \n\n |    | CoMERA |   HEAT     | \n  |----------|-|----------------|\n  |Training|**end-to-end compressed training** |post-training compression|\n  |Tensor format| TT for linear layers and TT matrix for embedding |CP, Tucker, and TT matrix for linear layers|\n  |Contraction path|jointly optimize **d+2 coupled paths** for **forward and back** in TT |one path for forward-propagation in CP|\n  \n***Weakness 3.1**: Does this training take more time in terms of convergence?*\n\nResponse: Thanks! In short, empirically our method is **2-3X** faster in the whole training process than uncompressed training when training transformers on a single GPU, but we do not have theoretical guarantees about the number of epochs although we observed that they used a similar number of epochs. We have provided detailed explanations in the **main author's rebuttal**. Some key points are summarized below:\n * Training neural networks is a highly non-convex optimization problem in both compressed and uncompressed formats, making theoretical convergence analysis very complicated. The overall training time depends on (1) the number of epochs and (2) time per epoch. While we observed consistent **2-3X** speedup in terms of (2), point (1) is highly case-dependent for almost all non-convex optimization solvers.\n * Our CoMERA has a **similar empirical convergence behavior** to the uncompressed training on our tested cases. We observe that on both 6-encoder transformer and DLRM, shown in Figure 6 on paper and **Figure 2 in the attached PDF file**. \n * On all tested transformers, CAMERA generally takes **2-3X less training time** because it has similar convergence curves as the uncompressed model, but each epoch of the tensorized model is 2-3X faster than standard training. \n\nAlthough our method has similar (and sometimes better) convergence behavior compared with standard training, we think that it could be misleading to draw any formal conclusion at this moment without a theoretical proof (which may or may not exist). \n\n***Weakness 3.2**: And could the author discuss more about the training difficulty of tensorized layers and normal layers.*\n\nResponse: From our experiments, we did not observe challenging difficulties of training tensorized layers. Adam is utilized to train the tensor-compressed model. The training hyperparameters, such as learning rate and weight decay, differ from standard training because of the different training parameters and loss landscapes, but the convergence behavior is similar. We will provide more details on training tensor-compressed models in the revised paper and conduct more analysis in the future.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LDgPg3BgSN", "reviewer": "Reviewer_6ExQ", "summary": "This work proposes using low-rank tensor train decomposition to accelerate deep learning model training and save memory usage. In the algorithm, both embedding tables in recommendation systems and large linear weights are written as a tensor train, and rank-adaptive optimization is used to adaptively reduce the rank without sacrificing the model accuracy. Multiple techniques, including TT embedding lookup, contraction path optimization, and CUDA graphs are combined to improve the efficiency. Experimental results show good memory saving and speedup without sacrificing model accuracy.", "questions": "1. line 42: replies -> relies\n\n2. How to choose a reasonable initial TT rank?\n\n3. how does accuracy compare between CoMERA, GaLore, and LTE?", "limitations": "n/a", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "originality: the idea to adaptively change the TT rank during training looks new. In addition to the training algorithm, multiple performance optimization techniques make the paper solid. In particular, the TT embedding lookup algorithm that does sampling and contraction in a interleaved way is also useful.", "weaknesses": "presentation: section 4.2 is hard to understand. For contraction path optimization, it would be good to visualize the process using tensor diagrams.\n\nlimitation of the proposed algorithm: it would be good to discuss the limitation of the algorithm in detail. In particular, I believe low-rank weights can help only under relatively small datasets/tasks, under which the low-rankness is a good regularization. Under large dataset/foundation model setting, I suspect the work can beat the baseline algorithms. In particular, the algorithm is also different from previous algorithms that use low-rank. In previous works, low-rank approximation is only applied on gradients rather than weights. To claim that the work can help foundation model training, we need larger datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Responses to weaknesses:\n\n***Weakness 1**: presentation: section 4.2 is hard to understand. For contraction path optimization, it would be good to visualize the process using tensor diagrams.*\n\nResponse: Thanks a lot for the suggestion! We prepared the tensor diagrams to visualize the contraction paths, but removed them from the paper because of the page limits. **We have included the tensor diagrams in the attached PDF**, see **Figure 3**.  We will include this diagram in the revision.\n\n***Weakness 2**: limitation of the proposed algorithm: it would be good to discuss the limitation of the algorithm in detail. In particular, I believe low-rank weights can help only under relatively small datasets/tasks, under which the low-rankness is a good regularization. Under large dataset/foundation model setting, I suspect the work can beat the baseline algorithms. In particular, the algorithm is also different from previous algorithms that use low-rank. In previous works, low-rank approximation is only applied on gradients rather than weights. To claim that the work can help foundation model training, we need larger datasets.*\n\nResponse: Thank you for the important questions! \n* The tensor-compression approach can achieve a higher compression ratio and better speedup on relatively small datasets. On larger datasets and models, a higher rank may be required to maintain the model performance. We are conducting larger experiments and the preliminary result is shown in the following **table** and **Figure 1 in the attached PDF file**. We pre-train the **CodeBERT-Large** model, released by Microsoft. It has 24 encoder blocks and in total **357 million parameters**, whose architecture is similar to BERT Large. The pre-training is done on the CodeSearchNet, a 20GB dataset widely used for pre-training LLM for automatic code generation. All linear layers in encoders are compressed into TT format. The embedding table and final linear layers are not compressed because CodeBERT enforces them to be the same, but CoMERA uses different tensor formats for TTM and for linear layers. Compared to uncompressed training, the tensor-compressed model shows a similar convergence curve and reaches a similar training loss, while compressing the whole model **4.25 times** and linear layers **9.77 times**. The tensor-compressed training is about **2.3X** faster in sequence length 128 and **1.9X** faster in sequence length 512 than uncompressed training on a single RTX 3090 GPU. The results demonstrate that our approach can scale up to larger models and tasks. We will investigate more in the future and are optimistic about the results on large models and datasets.\n\n|| Pre-training results of CodeBERT-large  |  |   \n  |-|----------|-|\n  |compression ratio|overall| 4.25$\\times$|\n  | |tensorized layers| 9.77$\\times$|\n  |training speedup|sequence length 128| 2.3$\\times$|\n  |                |sequence length 512| 1.9$\\times$|\n\n* The low-rank gradient approximation work, like GaLore, represents gradients by low-rank matrices, reducing the memory costs of first and second momentums for the optimizer. However, GaLore still uses the full model and applies the full back-propagation for gradients. Finding the projector and the projection of gradients into the compact space also bring extra computation overhead. In contrast, our approach directly compresses the weights, and the resulting gradient automatically has a compact low-rank tensorized form. Hence, our tensor-compression approach **does not have the computation overhead** and can reach **better memory savings and speedup**. A comparison is shown in Section 5.3 and Figure 1 in the paper. We will include more details to compare our method with the low-rank gradient approximation method.\n\n\n## Responses to questions:\n\n***Question 1**: line 42: replies -> relies*\n\nResponse: Thanks! We will fix the typo.\n\n***Question 2**: How to choose a reasonable initial TT rank?*\n\nResponse: Thank you for the interesting question! Our work **adaptively** determines the TT ranks during training. Generally, it is a super challenging problem to choose good initial TT ranks maximizing model performance without sacrificing performance prior to training. For our approach, we usually start with relatively larger ranks and let our method gradually prune ranks during training. This choice provides us with good experiment results as shown in Section 5.\n\n***Question 3**: how does accuracy compare between CoMERA, GaLore, and LTE?*\n\nResponse: The CoMERA and GaLore achieve almost the same validation accuracy, 64%, on the MNLI dataset. However, the LTE approach does not converge on the task by using its default setting.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AbbMTfkzik", "reviewer": "Reviewer_hvAM", "summary": "The manuscript presents techniques for efficient training from scratch based on tensor decompositions. The authors propose several modifications to the basic training approach to improve accuracy as well as several optimizations for tensor-compressed training, achieving training speedup. In the experiments, the proposed methodology is used to compress a transformer model and a recommender system model.", "questions": "1) I am a little confused about how to adapt the rank to achieve further compression. It is not clear from the text what the procedure is and when it is applied during training. Specifically, it is not explained how the ranks are chosen and at what point of time the procedure is applied. I would recommend incorporating a new paragraph into the paper that provides a more detailed explanation of the procedure, as well as some practical advice on how to apply it in various scenarios. This could help readers better understand the potential uses and benefits of this approach.\n\n2) As for the transformer model, I am also interested about the model upper compression limit that still maintains training speedup? I would imagine maintaining good performance can only be possible with lower compression ratios (say x2-5).", "limitations": "-", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The authors present a new approach for training a model with tenzorized weights, which allows for adaptation of ranks during training. They also managed to optimize the process and achieve a real-time speed-up, which is not an easy task when working with tensorized models. The described methods are useful and have the potential to significantly contribute to further advancements in efficient operation for tensor decomposition formats.", "weaknesses": "1) There are several places where a comparison with the baseline time (e.g., uncompressed) would be appropriate, such as in Figures 5 and 8.\n\n2) Experiments with the transformers are limited. Training a model on MNLI from scratch may not accurately reflect the behaviour of the method for language models,  as MNLI is a relatively simple classification dataset that might not require a large number of parameters. It would be beneficial if the paper included an experiment or a simulation of some kind in a pre-training setting for language models, such as pretraining Roberta on C4", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Responses to Weaknesses:\n\n***Weakness 1**: Baseline time in Figures 5 and 8.*\n\nResponse: Thanks a lot! We will include baseline time in the figures. For your convenience, we attach results in following **Table 1 and Table 2**. **Table 1** shows time and memory cost for embedding lookup forward-propagation and back-propagation. The uncompressed embedding with sparse gradients is faster than our approach since our TTM embedding table requires extra computation. However, the embedding without sparse gradients is slower than ours because of updating massive gradients. The proposed TTM embedding table uses **much less memory**, **7X less** than embedding with sparse gradients, and **15X less** than embedding without sparse gradients. **Table 2** shows the time and memory cost of training DLRM using CoMERA and standard training.Standard training is faster than CoMERA since DLRM is an embedding-intensive model, and the computation in the embedding table is look-up rather than matrix multiplications. However, CoMERA uses **much less memory**, saving **6.9X, 4.8X, and 3.1X** memory for batch sizes 10000, 2000, and 4000, respectively. \n\n **Table 1. Time and memory for embedding lookups**\n  ||    | CoMERA embedding |   uncompressed w/ sparse gradients     | uncompressed w/o sparse gradients |\n  |-|----------|-|----------------|--------------|\n  | batch 10000|time (s)    | 0.48 | 0.06     | 2.43     |\n  | |memory (MB)    | **670** | 5279     | 10558     |\n  | batch 20000|time (s)    | 0.82 | 0.11     | 2.48     |\n  | |memory (MB)    | **799** | 5284     | 10569     |\n  | batch 40000|time (s)    | 1.42 | 0.22     | 2.47     |\n  | |memory (MB)    | **896** | 5294     | 10592     |\n  \n**Table 2. Time and momery for DLRM training**\n || |CoMERA w/ optimization|CoMERA w/o optimization|uncompressed |\n|-|-|-|-|-|\n|batch size 10000|time (s)|807|1344|420|\n||memory (MB)|**2612**|9259|18261|\n|batch size 20000|time (s)|794|1182|423|\n||memory (MB)|**3947**|18385|19005|\n|batch size 40000|time (s)|791|N/A|424|\n||memory (MB)|**6629**| N/A| 20459|\n  \n***Weakness 2**: Lack of pertaining results on larger models*\n\nResponse: Thank you! Scaling up our approach to larger models and datasets is an important future work. We are conducting larger experiments and the preliminary result is shown in the following **table** and **Figure 1 in the attached PDF**. The details can be found in **main author rebuttal**. Here are some key results. \n* We pre-train the **CodeBERT-Large** model. It has 24 encoder blocks and in total 357 million parameters, whose architecture is similar to BERT Large. The pre-training is done on the CodeSearchNet, a 20GB dataset. Compared to uncompressed training, the tensor-compressed model shows a similar convergence curve and reaches a similar training loss, while compressing the whole model **4.25 times** and linear layers **9.77 times**. The tensor-compressed training is about **2.3X** faster in sequence length 128 and **1.9X** faster in sequence length 512 than uncompressed training on a single RTX 3090 GPU. The results demonstrate that our approach can scale up to larger models. We will investigate more in the future and are optimistic about the results.\n||Pre-training results of CodeBERT-large||   \n|-|---|-|\n|compression ratio|overall| 4.25$\\times$|\n| |tensorized layers| 9.77$\\times$|\n|training speedup|sequence length 128| 2.3$\\times$|\n||sequence length 512| 1.9$\\times$|\n## Responses to Questions:\n***Question 1**:  How to adapt the rank to achieve further compression.*\n\nResponse: Thank you! We will add more details about how to adapt the rank to achieve further compression. Our rank-adaptive approach uses the **multi-objective** formulation and consists of two stages:\n* In the early stage, we solve the multi-objective problem using linear scalarization. The early stage starts from the beginning of the training process and prunes tensor ranks gently without hurting the model performance. \n* After the early stage converges, we may continue training the model in the optional late stage. During the late stage, the multi-objective optimization problem is solved by the achievement scalarization approach to find a model close to our preferred performance and size for specific deployment requirements (e.g., on an FPGA). \n\n***Question 2**: Model upper compression limit that still maintains training speedup?*\n\nResponse: In short, speedup can be achieved by CoMERA even when the compression ratio is close to 1.\nWe consider an m by n linear layer and an input of size b by m and it is represented by TT cores with internal TT rank r. Then, the computation cost for linear layer and TT linear layer is about O(bmn) and O(b(m+n)r) respectively when the batch b is large. TT-compression has computation reduction roughly when $r\\le min(m,n)/2$. However, the compression ratio is close to 1 for such large ranks. We demonstrate the above analysis by testing the training time of the six-encoder transformer on MNLI. The following **Table 3** shows the per-epoch training time of CoMERA on MNLI dataset for different compression ratios. The acceleration is more obvious for larger compression ratios. When the compression ratio is greater than 1, CoMERA always has speedup. When the compression ratio approaches 1, the time of CoMERA approaches that of uncompressed training. We will include the results and discussions in our paper. \n\n**Table 3. Per epoch training time on MNLI for various compression ratios**\n\n  ||    | rank 30 |   rank 120     | rank 240 | rank 360 | rank 480 | uncompressed |\n  |-|----|-|--|--|--|-|--|\n  | compression ratio| | 50$\\times$ | 4.9$\\times$     | 2.2$\\times$     |1.5$\\times$|1.1$\\times$|N/A|\n  | time (min)| batch 32 | 7.2 | 7.79     | 11.16     | 13.94|17.46|18.5|\n  | | batch 64   | 6.4 | 6.73 | 9.7     |12.44 | 15.53|16.6|\n  | | batch 128   | 5.5 | 6.21 | 9.07    | 11.56 | 14.35|16.4|\n  \nIn our pre-training of CodeBERT-large, we still get **2.3X** speedup when the overall compression ratio is 4.25.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H3ZS0kFxTK", "reviewer": "Reviewer_Rz1r", "summary": "The paper introduces CoMERA, a novel training method for large AI models that focuses on optimizing both computing and memory efficiency through rank-adaptive tensor compression. CoMERA aims to reduce training costs and environmental impact by achieving high compression ratios and maintaining accuracy. Key contributions include a multi-objective optimization formulation, performance optimization techniques for tensor-compressed training on GPUs, and empirical validation showing significant speedup and memory savings compared to existing methods.", "questions": "1. For learning TT-ranks, this work imposes diagonal matrices  D  (shown in Eq. 6) with  l_1  norm regularization. This technique for tensor network structure search was also utilized in the recent paper:\n\nZheng, Yu-Bang, et al. “SVDinsTN: A Tensor Network Paradigm for Efficient Structure Search from Regularized Modeling Perspective.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n\n2. In line 122, could you provide more explanation of what the linear scalarization means? Does it refer solely to Eq. (7), or is there additional context?\n\n3. Could you provide more intuition regarding Eq. (11), i.e., the achievement scalarization? While the two scalarizations in the paper may be standard methods in multi-objective optimization, a more intuitive explanation would help general readers unfamiliar with this field grasp the main idea quickly.\n\n4. In lines 171-173, the design of the “tensor-index level” is unclear. Please improve the clarity of this part if possible.\n\n5. In the contraction section, how many core tensors are involved in the contraction? If the number of involved tensors in the contraction is not large, why not use the default contraction path searching algorithms integrated into einsum, e.g., ‘dp’ or ‘optimal’? What benefits can be achieved from handcrafting a new contraction path for the current model compression task?\n\n6. Could you provide more explanation in Sec. 4.3? Why can CUDA graphs improve tensor computation? What specific efforts were made in this work? Is using CUDA graphs straightforward, or does it require advanced programming techniques, such as designing a more GPU-suitable contraction order?\n\n7. In the comparison with GaLore, why was a 3090 used instead of a 4090, which was used in the original GaLore paper?\n\n8. Regarding the training time, only the time per epoch is given. Is it possible to provide the overall running time? Are more epochs required if tensor networks are utilized?\n\n9. Will the codes be released once the paper is accepted?", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experimental results demonstrate significant improvements in training speed and memory efficiency, outperforming recent methods like GaLore and LTE.\n2. By addressing both computing and environmental costs, the method has practical implications for large-scale AI model training, making tensor models more practically useful in machine learning.\n3. In the part of multi-objective optimization, the formulation balances compression ratio and model accuracy, providing a customizable framework for different resource requirements.", "weaknesses": "While the method shows impressive results on tested models, scalability to even larger models and diverse architectures remains an area for further exploration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Responses to Weaknesses:\n\n***Weakness 1**: Scalability to larger models.*\n\nWe are conducting larger experiments. The preliminary result is in **Figure 1 in the attached PDF** and details are in **main author rebuttal**. We pre-train **CodeBERT-Large** with 357 million parameters on CodeSearchNet, a 20GB dataset. Compared to uncompressed training, the tensor-compressed model shows a similar convergence curve and reaches a similar training loss and achieves **training speedup** on a single 3090 GPU. Compression and speedup results are in the following table. \n||Pre-training results of CodeBERT-large||   \n|-|-|-|\n|compression ratio|overall| 4.25$\\times$|\n||tensorized layers| 9.77$\\times$|\n|training speedup|sequence length 128| 2.3$\\times$|\n||sequence length 512| 1.9$\\times$|\n## Responses to Questions:\n***Question 1**: Compare diagonal D with l_1 norm in the recent paper*\n\nWe will cite it and discuss the differences. The referenced paper uses sparsity of diagonal matrices to search for a **compact TT representation of given data**, which is very different from our **end-to-end training**.  We would like to highlight the differences in the following. \n* The referenced paper searches for a TT representation of given data. It is same as **finding a low-rank approximation for a *given* tensor**. In contrast, our work uses TT to compress weights during **end-to-end training** and **we don’t have any prior information on the tensor** (i.e., model parameters). The tensor cores and diagonal entries are determined during training. \n* Our work formulates the problem as a more generic **multi-objective** problem and uses a two-stage algorithm to solve it. The formulation in that paper is similar to linear scalarization approach in our early stage. Our work further uses the achievement scalarization in the late stage to find a model close to our preferred model performance and size.\n\n***Question 2**: Explanation of linear scalarization*\n\nYes, the linear scalarization solely refers to the Eq(7). More details are in the next question.\n\n***Question 3**: Intuition for scalarization?*\n\nWe will include more intuitions. The linear scalarization minimizes the weighted sum of objectives. Solving it provides a Pareto point, but it is hard to control the obtained point. The achievement scalarization finds a Pareto optimal solution that is close to a **given target**. When one objective is too far away from the target, we will mainly optimize this objective by increasing its weight. It gives us the achievement scalarization problem Eq. (11).\n\n***Question 4**: Design of the “tensor-index level”.*\n\nThanks! We will revise that. “Tensor-index level” optimization aims to reduce redundant computation. In the tensor indices, two different rows may have shared values. For instance, (2,3,1,3) and (2,3,2,4) are common in (2,3), so we only need to compute (2,3) entry once.  In this design, we focus on the unique indices required for the lookup and only compute the required unique indices. \n\n***Question 5**: Why not using contraction path search in einsum?*\n\nFor TT with d tensor cores, each contraction has d+1 cores  involved and there are in total d+2 coupled einsums, as in Eq. (19) (20) (21). The algorithm in “einsum” only searches for the path for a **single** einsum operation. However, our path optimization tries to minimize the **total costs of all d+2 coupled contraction** paths, which cannot be handled by the searching options in “einsum”. We also show in Proposition 4.1 that the contraction path for forward is already near-optimal. Similar results can be shown for back-propagation. \n\n***Question 6**: Details about Cuda Graph*\n\nThanks! We will include more explanations about Cuda graphs in Section 4.3. Cuda Graph eliminates overhead of launching lots of kernels sequentially. It is more suitable for CoMERA since tensor-compressed training has much more small kernels than uncompressed training. No dynamic control flows and dynamic shapes are allowed in standard CudaGraph. We specifically revise some codes, making them compatible with Cuda Graph. For instance, the inputs are padded into a fixed length and the late-stage achievement scalarization in Section 3.2 is split into two graphs. \n\n***Question 7**: Why was a 3090 used instead of a 4090 for GaLore comparison?*\n\nThankst! We don’t have the 4090 GPU in our lab. All comparisons are run on the same system and the 3090 GPU. \n\n***Question 8**: Overall training time*\n\nThanks! Empirically CoMERA is **2-3X** faster in the whole training than uncompressed training for transformers on a single GPU, but we do not have theoretical guarantees about the number of epochs although they are similar in our experiments. The detailed explanations are in the **main author's rebuttal**. Some key points are summarized below:\n * Training NNs is a highly non-convex optimization in compressed and uncompressed formats, making theoretical analysis very complicated. Overall training time depends on (1) the number of epochs and (2) time per epoch. While we observed consistent **2-3X** speedup in terms of (2), point (1) is highly case-dependent for almost all non-convex optimization solvers.\n * CoMERA has a **similar empirical convergence behavior** to the uncompressed training on our tested cases. We observe that on both 6-encoder transformer and DLRM, shown in Figure 6 on paper and **Figure 2 in the PDF**. \n * On all tested transformers, CAMERA generally takes 2-3X less training time because it has similar convergence curves as the uncompressed model, but each epoch of the tensorized model is 2-3X faster than standard training. \n\nAlthough our method has similar convergence behavior compared with standard training, we think that it could be misleading to draw any formal conclusion now without theoretical proof (which may or may not exist). \n\n***Question 9**: Code release?*\n\nWe are trying to figure out some potential IP issues. We hope to release the codes if the IP issue can be cleared.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Common Concerns 
We would like to thank the reviewers for their fruitful suggestions and comments. **We have addressed ALL review comments** (see the response to each reviewer).

Here we summarize some concerns arised in the review process. 


## **Scalability of tensor-compressed training to larger models**

A few reviewers are concerned about whether the tensor-compressed training scales up to larger models and tasks. It is also a problem we are now actively investigating. We are conducting larger experiments and the preliminary result is shown in the following **table** and **Figure 1 in the attached PDF file**. We pre-train the **CodeBERT-Large** model, released by Microsoft. It has 24 encoder blocks and in total **357 million parameters**, whose architecture is similar to BERT Large. The pre-training is done on the CodeSearchNet, a 20GB dataset widely used for training LLM for automatic code generation. All linear layers in encoders are compressed into TT format. The embedding table and final linear layers are not compressed because CodeBERT enforces them to be the same, but CoMERA uses different tensor formats for TTM and for linear layers. Compared to uncompressed training, the tensor-compressed model shows a similar convergence curve and reaches a similar training loss, while compressing the whole model **4.25 times** and linear layers **9.77 times**. The tensor-compressed training is about **2.3X** faster in sequence length 128 and **1.9X** faster in sequence length 512 than uncompressed training on a single RTX 3090 GPU. The results demonstrate that our approach can scale up to larger models and tasks. We will investigate more in the future and are optimistic about the results on large models and datasets. 
|| Pre-training results  of CodeBERT |  |   
  |-|----------|-|
  |compression ratio|overall| 4.25$\times$|
  | |tensorized layers| 9.77$\times$|
  |training speedup|sequence length 128| 2.3$\times$|
  |                |sequence length 512| 1.9$\times$|

## **Overall training time and total epochs.** 

The paper mainly compares the per-epoch time of CoMERA and uncompressed training. The reviewers also wonder about the overall training time and the number of epochs. We address the common concern here.

 In short, empirically our method is **2-3X** faster than uncompressed training when training transformers on a single GPU, but we do not have theoretical guarantees about the number of epochs although we observed that they used a similar number of epochs. Detailed explanations are provided below
* Training neural networks is a highly non-convex optimization problem in both compressed and uncompressed formats, making the theoretical convergence analysis very complicated. Even existing theoretical analysis of uncompressed training is done on simplified cases, such as the two-layer neural networks with infinite width. The overall training time depends on (1) the number of epochs and (2) time per epoch. While we observed 2-3X speedup in terms of (2), point (1) is highly case-dependent for almost all non-convex optimization solvers. 
* Our CoMERA has a **similar empirical convergence behavior** to the uncompressed training on our tested cases. For the 6-encoder transformer, CoMERA converges a little slower than the standard training at the beginning as shown in Figure 6 in the paper, but finally CoMERA achieves a higher validation accuracy. For the DLRM task, CAMERA needs fewer iterations than the standard training, as shown in **Figure 2 in the attached PDF file**. We will also add that figure to our revised manuscript. 
* On all tested transformers, CAMERA generally takes 2-3X less training time because it has similar convergence curves as the uncompressed model, but each epoch of the tensorized model is 2-3X faster than standard training. 

Although our method has similar (and sometimes better) convergence behavior compared with standard training. We think that it could be misleading to draw any conclusion at this moment without a theoretical proof (which may or may not exist). 



## **Figures in PDF file** (attached)
**Figure 1.** Training loss of CodeBERT-large. 
* Figure 1 shows the empirical convergence curve of tensor-compressed training on the **CodeBERT-Large** model, a model of 24 encoder blocks and in total **357 million** parameters before compression.

**Figure 2.** The validation normalized cross-entropy loss of training DLRM model. 
* It presents the validation loss of CoMERA and uncompressed training during training DLRM. CoMERA converges slightly faster in terms of iterations than the uncompressed training on this task.

**Figure 3.** Tensor diagrams for contraction paths of TT forward- and back- propagation. 
* The figure visualizes the proposed contraction paths for TT forward- and back- propagation as detailed in Section 4.2.

Meta-review: The reviewers all evaluate as borderline accept / accept, with most improving evaluation after discussions with the authors. The authors had extensive discussion / pushed hard in order to get the evaluations improved. I rate as "less certain" as ultimately it seems none of the reviewers were particularly excited about the paper in initial evaluation. But it seems logical to accept the paper given the evaluation changes.
Decision: Accept (poster)