TTOpt: A Maximum Volume Quantized Tensor Train-based Optimization and its Application to Reinforcement Learning
Abstract
We present a novel procedure for optimization based on the combination of efficientquantized tensor train representation and a generalized maximum matrix volumeprinciple.
We demonstrate the applicability of the new Tensor Train Optimizer(TTOpt) method for various tasks, ranging from minimization of multidimensionalfunctions to reinforcement learning.
Our algorithm compares favorably to populargradient-free methods and outperforms them by the number of function evaluationsor execution time, often by a significant margin.
Introduction
In recent years learning-based algorithms achieved impressive results in various applications, rangingfrom image and text analysis and generation [52] to sequential decision making and control [41]and even quantum physics simulations [49].
The vital part of every learning-based algorithm isan optimization procedure, e.g., Stochastic Gradient Descent.
In many situations, however, theproblem-specific target function is not differentiable, too complex, or its gradients are not helpful dueto the non-convex nature of the problem [30, 2].
The examples include hyper-parameter selectionduring the training of neural models, policy optimization in reinforcement learning (RL), trainingneural networks with discrete (quantized) weights [57] or with non-differentiable loss functions [20].
In all these contexts, efficient direct gradient-free optimization procedures are highly needed.
Recently, [53] showed that an essential class of gradient-free methods, namely the evolutionarystrategies (ES) [22, 26, 56], are competitive in reinforcement learning problems.
In RL, the goal is tofind the agent’s action distribution π (the policy), maximizing some cumulative reward function J .
Thepolicy is usually parameterized with a set of parameters θ .
It follows that the reward is a function ofthe parameters of the policy: J ( π ( θ ))
=
J ( θ ) .
The idea of [53] and similar works [37, 15, 33, 10, 11]is to directly optimize the cumulative reward function J ( θ ) =
J ( θ 1 , θ 2 , . .
, θ d ) with respect tothe parameters of the policy.
We pursued a similar approach to transform a traditional Markov
Decision Process into an optimization problem (we provide the details in Appendix B.3).
Althoughthe agents trained with ES often demonstrate more rich behavior and better generalization comparedto traditional gradient-based policy optimization, the convergence of ES is often slow [11].
As an alternative to previous works, we present a tensor-based 1 gradient-free optimization approachand apply it to advanced continuous control RL benchmarks.
The algorithm that we called TensorTrain (TT) Optimizer ( TTOpt ), works for multivariable functions with discrete parameters byreformulating the optimization problem in terms of tensor networks.
Consider a function J ( θ ) :
N d −→ R of a d -dimensional argument θ , where each entry θ k of the vector θ takes a value in thediscrete set { ω i } Ni =1 .
The function J may be viewed as an implicitly defined d -dimensional tensor J .
Each entry in J is a value of J for some argument.
Maximizing J is equivalent to finding the sets ofindices { θ ( m ) max } Mm =1 of maximal entries { j ( m ) max } Mm =1 of J .
By using only a tiny fraction of adaptively selected tensor elements (e.g., a small number of functionevaluations), our method builds a representation of J in the TT-format and finds a set of the largestelements.
Although the algorithm works only with functions of discrete arguments, the grid size foreach parameter can be huge thanks to the efficiency of the TT-format.
Fine discretization makes itpossible to almost reach a continuous limit and obtain large precision with TTOpt.
The strength ofour approach is, however, the direct handling of discrete (quantized) parameters.
Contributions.
We propose an efficient gradient-free optimization algorithm for multivariablefunctions based on the low-rank TT-format and the generalized maximum matrix volume principle 2 .
We demonstrate that our approach is competitive with a set of popular gradient-free methods foroptimizing benchmark functions and neural network-based agents in RL problems.
We empiricallyshow that agents with discrete (heavily quantized) weights perform well in continuous control tasks.
Our algorithm can directly train quantized neural networks, producing policies suitable for low-powerdevices.
Optimization with tensor train
In this section, we introduce a novel optimization algorithm.
We show how to represent the optimization problems in the discrete domain efficiently and then formulate optimization as a sampling of theobjective function guided by the maximum volume principle.
2.1 Discrete formulation of optimization problems
We first need to transfer the problem to the discrete domain to apply our method.
It may seem thatdiscretizing an optimization problem will make it harder.
However, it will allow us to use powerfultechniques for tensor network representation to motivate the algorithm.
For each continuous parameterθ k ( k = 1 , 2 , . . .
, d ) of the objective function J ( θ ) we introduce a grid { θ ( n k ) k } N k n k =1 .
At each point( θ ( n 1 ), θ ( n 2 ), . . .
, θ ( n d ) d ) of this grid with index ( n 1 , n 2 , . . . , n d ) the objective function takes a value
J ( θ ( n 1 ), θ ( n 2 ), . . .
, θ ( n d ) d ) ≡ J [ n 1 , n 2 , . . . , n d ] .
We thus can regard the objective function as animplicit d -dimensional tensor J with sizes of the modes N 1 , N 2 , . .
. , N d .
Finding the maximum ofthe function J ( θ ) translates into finding the maximal element of the tensor J in the discrete setting.
Notice that the number of elements of J equals: |J | = N 1 · N 2 · . . .
· N d ∼ (max 1 ≤ k ≤ d N k ) d .
The size of J is exponential in the number of dimensions d .
This tensor can not be evaluated orstored for sufficiently large d .
Fortunately, efficient approximations were developed to work withmultidimensional arrays in recent years.
Notable formats include Tensor Train (TT) [47, 46, 50],
Tensor Chain/Tensor Ring [29, 64] and Hierarchical Tucker [21].
We use the most studied TTformat [13], but the extensions of our method to other tensor decompositions are possible.
2.2 Tensor Train decomposition
Definition 2.1.
A tensor J ∈ R N 1 × N 2 ×···× N d is said to be in the TT-format [46] if its elements arerepresented by the following expression
[Equation]
where n k = 1 , 2 , . . . , N k for k = 1 , 2 , . .
, d .
In TT-format the d -dimensional tensor J is approximated as a product of three-dimensional tensors
G k ∈ R R k − 1 × N k × R k , called TT-cores.
The sizes of the internal indices R 0 , R 1 , · · · , R d (withconvention R 0 = R d = 1 ) are known as TT-ranks.
These ranks control the accuracy of theapproximation.
The storage of the TT-cores, G 1 , G 2 , . .
. , G d , requires at most d · max 1 ≤ k ≤ d N k · (max 0 ≤ k ≤ d R k ) 2memory cells, and hence the TT-approximation is free from the curse of dimensionality 3 if the
TT-ranks are bounded.
The basic linear algebra operations (such as finding a norm, differentiation,integration, and others) can also be implemented in the TT-format with polynomial complexity indimensionality and mode size.
Building TT-approximation.
Several efficient schemes were proposed to find TT-approximation ifall or some of the elements of the initial tensor are known or may be generated by the function’s call.
Examples include TT-SVD [47, 46], TT-ALS [27] and TT-CAM [48] (Cross Approximation Methodin the TT-format).
We build upon the TT-CAM but modify it not to compute the approximation for the entire tensor, butrather to find a small subset of its maximal entries.
The original algorithm builds a TT-approximationby adaptively requesting elements of the input tensor.
As we will show below, these elements withhigh probability will have large absolute values.
Based on this observation, we formulate a robustoptimization algorithm for multivariate functions (either discrete or continuous).
To simplify theunderstanding, we outline the approach for the two-dimensional case, and after that, we describe ourgradient-free optimization method for the multidimensional case.
2.3 Maximal element in a matrix
The Cross Approximation Method (CAM) for matrices [18, 9, 1] is a well-established algorithm forbuilding a rank- R approximation ˜
J of an implicitly given matrix J :
[Equation]
where J C consists of R columns of J , J R is composed of R rows of J , and ˆ J is a submatrix attheir intersection.
Such approximation (also called cross or skeleton decomposition) may be builtiteratively using a well known alternating directions method and a maximum volume ( maxvol )algorithm 4 [18], as we will sketch below.
Intuition behind TTOpt.
The main interest in optimization problems is not the approximation (2)itself, but the following property of the resulting maximum volume submatrix ˆ J .
[18] proved that ifˆ J is an R × R submatrix of maximal volume (in selected rows and columns)
then the maximal (bymodulus) element ˆ J max ∈
ˆ J bounds the absolute maximal element J max in the full matrix J :
[Equation]
algorithm optimum the dimensional
J
This statement is evident for R = 1 , and for the case R > 1 it gives an upper bound for the element.
By using elementwise transformations of J , this upper bound can be used to obtain a sequence thatconverges to the global optimum.
The main idea of the maxvol -based methods is that it is easier tofind a submatrix with a large volume rather than the element with the largest absolute value.
Moreover,our numerical experiments show that this bound is pessimistic, and in practice, the maximal-volumesubmatrix contains the element which is very close to the optimal one.
TTOpt algorithm for matrices.
The idea of the TTOpt algorithm for matrices is to iterativelysearch for the maximal volume submatrices in the column and row space of the implicitly given 5 inputmatrix J ∈ R N 1 × N 2 .
After T iterations a series of “intersection” matrices { ˆ J ( t ) }
Tt =1 ∈ R R × R isproduced.
The maximal element is searched in these small submatrices.
We schematically representthe TTOpt algorithm in Figure 1, and a description is given below:
At the initial stage, we set the expected rank of the approximation, R , and select R random columns I ( C, 1) .
We then generate the corresponding column submatrix J
(1) C =
J [:
, I ( C, 1) ] ∈ R N 1 × R .
Using the maxvol algorithm, we find the maximal-volume submatrixˆ J (1) ∈ R R × R in J (1) C and store its row indices in the list
I ( R, 1) .
The indices I ( R
1) are used to generate
a row
submatrix
J (2) R =
J [ I ( R, 1) , :] ∈ R R × N 2 .
Then, using the maxvol algorithm, we find the maximal-volume submatrix ˆ J (2) in thematrix J (2) R and store the corresponding column indices in the list I ( C, 2) .
We generate the related columns J (3) C =
J [: , I ( C, 2) ] , apply again the maxvol algorithm tothe column submatrix, and iterate the process until convergence.4.
The approximate value of the maximum modulus element of the matrix J is found as
[Equation]
2.4 Optimization in the multidimensional case
As we explained previously, the target tensor, J , is defined implicitly, e.g., by a multivariable function
J .
We propose a novel method to find the optimum in this implicit tensor.
We outline the approachbelow and provide detailed algorithms in Appendix A.2.
As shown in Figure 2, we begin by considering the first unfolding 6 J 1 ∈ R N 1 × N 2 ...N d of the tensor
J and select R 1 random columns I ( C ) 1 .
Precisely, I ( C )here is a list of R 1 random multi-indices ofsize d − 1 , which specify positions along modes k = 2 to k = d .
We then generate the submatrix
J ( C )∈ R N 1 × R 1 for all positions along the first mode (shown in green in Figure 2).
Like in matrixcase, we apply the maxvol algorithm to find the maximal-volume submatrix ˆ J 1 ∈ R R 1 × R 1 and storethe corresponding indices of R 1 rows in the list
I ( R ) 1 .
In contrast with matrix case, we cannot generate the row submatrix J ( R )∈ R R 1 × N 2 N 3 ...N d forthe selected row indices
I ( R ) 1 , since it contains an exponential number of elements.
The followingtrick is used instead.
We consider the implicit matrix J ( R )and reshape it to a new matrix J 2 ∈
R R 1 N 2 × N 3 ...N d .
We sample R 2 random columns
I ( C )in the matrix J 2 and generate the entiresmall submatrix J ( C )∈ R R 1 N 2 × R 2 .
Next we find the maximal-volume submatrix ˆ J 2 ∈ R R 2 × R 2in J ( C )and store the corresponding R 2 row multi-indices in the list
I ( R ) 2 .
The resulting submatrix
J ( R )∈ R R 2 × N 3 N 4 ...
N d is then transformed (without explicitly evaluating its elements) into thematrix J 3 ∈ R R 2 N 3 × N 4 ...N d .
We continue the described operations, called sweeps, until the last mode of the initial tensor, J , isreached.
After that, we repeat the process in the opposite direction, sampling now the row indicesinstead of the column indices.
These sequences of forward and backward sweeps are continued untilthe algorithm converges to some row and column indices for all unfolding matrices 7 or until theuser-specified limit on the number of requests to the objective function, J , is exceeded.
Finally, after
T sweeps, the approximate value of the maximum modulo element can be found by the formula (4),as in the two-dimensional case.
Note that currently there is no analog of Eq.
(3) in the multidimensional case, and hence there areno formal guarantees of the convergence of the sweeps to the global minimum, nor the rate of thisconvergence.
The only guarantee is that the result will monotonically improve with iterations.
2.5 Complexity of the algorithm
It can be easily shown that the described algorithm requires to evaluate only
O (cid:0) d · max
[Equation]
N
[Equation]
k (cid:1)(cid:1) elements of the implicit tensor in one sweep.
Thus, with a totalnumber of sweeps T , we will have
O
[Equation]
(5) calls to the objective function J .
In practice, it turns out to be more convenient to limit the maximumnumber of function calls, M , according to the computational budget.
[Equation]
If the time of a single call to J is significant, then the effort spent on the algorithm’s operation willbe negligible.
Otherwise, the bottleneck will be the calculation of the maximal-volume submatricesby the maxvol algorithm.
Taking into account the estimate of the maxvol complexity, given in
Appendix A.1, it can be shown that in this case the complexity of our algorithm is
[Equation]
2.6 Implementation details
For the effective implementation of the TTOpt algorithm, the following important points should betaken into account (see also the detailed pseudocode in Appendix A.2).
Stability.
Submatrices J ( C ) k
(or J ( k ) C and J ( k ) R for the two-dimensional case) that arise during theiterations may degenerate, and in this case it is impossible to apply the maxvol algorithm.
To solvethis problem, we first calculate the QR decomposition for these matrices and then apply the maxvolto the corresponding Q factors 8 .
Rank selection.
We do not know in advance the exact ranks R 1 , R 2 , . . .
, R d of the unfoldingmatrices (or rank R of the matrix J for the two-dimensional case).
Therefore, instead of the maxvolalgorithm, we use its modification, i.e., the rect_maxvol algorithm 9 [40], within the framework ofwhich several (“most important”) rows are added to the maximal-volume submatrix.
In this case,we have ˆ J ( C )
[Equation]
∈ R ( R 2 +∆ R 2 ) ×
R 2 , etc.
Note that the final approximationobtained in this case may have an overestimated rank, which, if necessary, can be reduced byappropriate rounding, for example, by truncated SVD decomposition.
Mapping function.
Maximal-volume submatrices contain the maximum modulus element butnot the minimum or maximum element of the tensor (i.e., the sign is not taken into account).
Weintroduce a dynamic mapping function to find the global minimum (or maximum).
Instead of J ateach step of the algorithm, we evaluate 10 :
[Equation]
where J min is the current best approximation for the minimum element of the tensor.
Quantization.
To reach high accuracy, we often need fine grids.
In the case when the sizes of thetensor modes N 1 , N 2 , . . .
, N d are large, the sizes of unfolding matrices become large, which leads toa significant increase in computational complexity of the maxvol algorithm.
To solve this problem,we apply additional compression based on quantization of the tensor modes [45].
Assume without theloss of generality that the size of each mode is N k = P q ( k = 1 , 2 , . . .
, d ; P ≥ 2 ; q ≥ 2 ).
Then wecan reshape the original d -dimensional tensor J ∈ R N 1 × N 2 ×···× N d into the tensor ˜ J ∈ R P × P ×···× Pof a higher dimension d · q , but with smaller modes of size P .
The TTOpt algorithm can be appliedfor this “long” tensor instead of the original one 11 .
We found that this idea significantly boosts theaccuracy of our algorithm and reduces the complexity and execution time.
Typically, P is taken assmall as possible, e.g. P = 2 .
Experiments
To demonstrate the advantage of the proposed optimization method, we tested TTOpt on severalnumerical problems.
First, we consider analytical benchmark functions and then the practically
The mapping function
[Table]
[Table]
significant problem of optimizing the parameters of the RL agent.
We select GA (Genetic Algorithm [26, 59]), openES (basic version of OpenAI Evolution Strategies [53]) and cmaES (the
Covariance Matrix Adaptation Evolution Strategy [22]) as baselines for both experiments.
Additionally we used DE (Differential Evolution [58]), NB (NoisyBandit method from Nevergrad [5])and PSO (Particle Swarm Optimization [24, 34]) for benchmark functions 12 .
We also comparedthe proposed approach with gradient-based methods applied for all benchmark functions 13 :
BFGS(Broyden–Fletcher–Goldfarb–Shanno algorithm), L-BFGS (Limited-memory BFGS), CG (Conjugate Gradient algorithm),
NCG (Newton CG algorithm), Newton (Newton Exact algorithm), TR NCG (Trust-Region NCG algorithm) and TR (Trust-Region Exact algorithm). According to our approach, the TTOpt solver has the following configurable parameters: a and b arelower and upper grid bounds (for simplicity, we use the same value for all dimensions); R is a rank(for simplicity, we use the same value for all unfolding matrices); P is a submode size (mode size ofthe quantized tensor; for simplicity, we use the same value for all dimensions); q is the number ofsubmodes in the quantized tensor (each mode of the original tensor has size N = P q ); M is a limiton the number of requests to the objective function. [Table]
3.1 Benchmark functions minimization
[Table] To analyze the effectiveness of the TTOpt , we applied it to 10 -dimensional benchmark functions withknown global minimums; see Table 1 in Appendix B.1 with the list of functions and their properties(note that some of the considered benchmarks are multimodal non-separable functions).
Also, in
Appendix B.1, we present a more detailed study of the TTOpt solver and the dependence of theaccuracy on the value of its parameters ( R , q and M ).
In all experiments with baselines ( GA , openES , cmaES , DE , NB , PSO , BFGS , L-BFGS , CG , NCG , Newton , TR NCG and TR ), we used default parameter values.
In Appendix B.2 we alsopresent the additional experiments with Bayesian Optimization [38].
For TTOpt we selected rank {+
+} R = 4 , submode size P = 2 and the number of submodes q = 25 .
For all methods, a limit on thenumber of requests to the objective function is chosen as M = 10 5 .
All calculations are performedon a standard laptop.
The results are demonstrated in Table 1.
For each method we list the absolute deviation of theresult ˆ J min from the global minimum J min , i.e., ϵ = | ˆ J min − J min | .
We also present the totalrunning time, τ .
Compared to other benchmarks, TTOpt is consistently fast, accurate and avoidsrandom failures to converge seen in other algorithms.
Additionally, TTOpt turns out one of thefastest gradient-free algorithms ( GA , openES , cmaES , DE , NB , PSO ), despite a simple Pythonimplementation.
One of the advantages of the proposed TTOpt approach is the possibility of its application toessentially multidimensional functions.
In Table 2 we present the result of TTOpt for the F1benchmark function of various dimensionality (results for other benchmarks are in Appendix B.1).
Note that as a limit on the number of requests to the objective function we chose M = 10 4 · d , andthe values of the remaining parameters were chosen the same as above.
As can be seen, even for500 -dimensional functions, the TTOpt method gives a fairly accurate result.
3.2 Application of TTOpt to Reinforcement Learning
We used several continuous RL tasks implemented in Mujoco [60] and OpenAI-GYM [8]: Swimmerv
[16], LunarLanderContinuous-v2, InvertedPendulum-v2 and HalfCheetah-v3 [62].
In all experiments, the policy π is represented by a neural network with three hidden layers and with tanh and
ReLU activations.
Each layer is a convolution layer.
See additional details about hyperparameters in
Table 6 in Appendix B.
We discretize (quantize) the values of agent’s weights.
The TTOpt method is used to optimize discreteagent’s weights in order to maximize the cumulative reward of the episode.
This corresponds toon-policy learning.
[Equation]
[Figure] HalfCheetah-v [Figure] To properly compare TTOpt with other methods, we propose modified evolutionary baselines thatenforce constrained parameter domain.
We adapt penalty term and projection techniques from [31, 6]to introduce constraints, see Appendix B.4 for details.
First, we run benchmarks with small mode size N = 3 1 with lower and upper grid bounds ± 1 .
Another series of experiments was done with finer mode of size N = 2 q with the same bounds.
Theseexperiments model the case of neural networks with discrete (quantized) weights which use q -bitsquantization.
Finally, we provide the results of using TTOpt as a fine-tuning procedure for linearpolicies from [37].
We present characteristic training curves based on the number of environment interactions andexecution time for the HalfCheetah-v3 experiment ( N = 3 ) in Figure 3.
Training curves for otherenvironments can be found in Appendix B.5, in Figure 3 for N = 3 and in Figure 4 for N = 256 .
TTOpt consistently outperforms all other baselines on the coarse grid with mode size N = 3 .
Our method is still best for finer grids with mode size N = 256 on InvertedPendulum-v2 and
LunarLanderContinuous-v2, and second-best on HalfCheetah-v3.
Moreover, TTOpt has significantlylower execution time compared to evolutionary baselines.
Another interesting observation is that thetraining curves of TTOpt have low dispersion, e.g., the algorithm performs more consistently thanthe baselines (see Appendix B.5).
Table 3 summarizes the experiments for the coarse and fine grids.
Results for fine-tuning of linear policies are presented in Table 5 in Appendix B.5.
We also did rankand reward dependency study in Appendix B.
Related work
In the case of high dimensional optimization, evolutionary strategies (ES) [56, 43] are one of themost advanced methods of black-box optimization.
This approach aims to optimize the parameters ofthe search distribution, typically a multidimensional Gaussian, to maximize the objective function.
Finite difference schemes are commonly used to approximate gradients of the parameters of thesearch distribution.
Numerous works proposed techniques to improve the convergence of ES [43].[63] proposed to use second-order natural gradient of [12] to generate updates, while [22] suggestedto include the history of recent updates to generate next ones.
[36] presented the concept of surrogategradients for faster convergence.
Another series of works aimed to reduce the high sampling complexity of ES.
In [11] the authors described how to use active subspaces [14] to reduce the number of objective function callsdynamically.
Plenty of the already mentioned works in gradient-free optimization specifically applied thesemethods to RL tasks [53, 10, 11, 37, 15, 33, 23].
Overall, the performance of ES-based methods is comparable to conventional policy gradients, especially if the number of model parameters issmall [54, 55].
Another advantage of ES over policy gradient is that it produces more robust anddiverse policies [32, 33] by eliminating the problem of delayed rewards and short length time horizons.
Finally, evolutionary approaches are suitable for the problems with non-Markovian properties [23].
Other metaheuristic [25, 39] and classical optimization [51] techniques are also studied within RLscope.
The examples include simulated annealing [4], particle swarm optimization [24, 34] and evenclassical Nelder-Mead algorithm [42, 44].
These methods, however, are not tested on common RLtask sets.
Several other works combined evolutionary methods with RL to achieve better performancein complex scenarios [28, 17], e.g. AlphaStar [61, 3].
Finally, low-rank tensor approximations have been applied to RL problems in settings different fromours, including multi-agent scenarios [35], and improving dynamic programming approaches [19, 7].
The idea of our method is quite different from the presented works, especially in the RL area.
Conclusion
We proposed a new discrete optimization method based on quantized tensor-train representationand maximal volume principle.
We demonstrate its performance for analytical benchmark functionsand reinforcement learning problems.
Our algorithm is more efficient under a fixed computationalbudget than baselines, especially on discrete domains.
Moreover, the execution time for TTOpt islower by a significant margin compared with other baselines.
Finally, we show that the agents withdiscrete parameters in RL can be as efficient as their continuous parameter versions.
This observationsupports the broad adoption of quantization in machine learning.
We hope that our approach willserve as a bridge between continuous and discrete optimization methods.