SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anastasia Batsheva*
a.batsheva@skoltech.ru
Gleb Ryzhakov*
g. ryzhakov@skoltech.ru
Andrei Chertkov*
and AIRI, Moscow, Russia
a.chertkov@skoltech.ru
i. oseledets@skoltech.ru
We developed a new method PROTES for black-box optimization, which is based on the probabilistic sampling from a probability density function given in the low-parametric tensor train format. We tested it on complex multidimensional arrays and discretized multivariable functions taken, among others, from real-world applications, including unconstrained binary optimization and optimal control problems, for which the possible number of elements is up to  $2^{1000}$ . In numerical experiments, both on analytic model functions and on complex problems, PROTES outperforms popular discrete optimization methods (Particle Swarm Optimization, Covariance Matrix Adaptation, Differential Evolution, and others).
The multidimensional optimization problem is one of the most common in machine learning. It includes the important case of gradient-free discrete optimization [30, 20, 21, 41]:
$$
\boldsymbol {x} _ {\text {m i n}} = \min  _ {\boldsymbol {x}} \mathbf {f} (\boldsymbol {x}), \quad \text {s . t .} \boldsymbol {x} = \left[ n _ {1}, n _ {2}, \dots , n _ {d} \right], \quad n _ {i} \in \{1, 2, \dots , N _ {i} \}, \tag {1}
$$
where  $d$  is the dimensionality of the problem, and  $N_{1}, N_{2}, \ldots, N_{d}$  are the numbers of items for each dimension. Such settings arise when searching for the minimum or maximum element in an implicitly given multidimensional array (tensor $^{1}$ ), including when considering the discretization of functions from a continuous argument. Multidimensional discrete optimization problems are still computationally difficult in the case of complex target functions or large dimensions, and efficient direct gradient-free optimization procedures are highly needed.
The development of methods for low-rank tensor approximations has made it possible to introduce fundamentally new approaches for the approximation, storage, and operation with multidimensional
<sup>1</sup>A tensor is a multidimensional array with several dimensions  $d$  ( $d \geq 1$ ). A two-dimensional tensor ( $d = 2$ ) is a matrix, one-dimensional ( $d = 1$ ) - a vector. We use lower case and upper case (where this does not lead to confusion) letters for scalars, bold letters ( $\mathbf{a}, \mathbf{b}, \mathbf{c}, \ldots$ ) for vectors, upper case letters ( $A, B, C, \ldots$ ) for matrices, and calligraphic upper case letters ( $\mathcal{A}, \mathcal{B}, \mathcal{C}, \ldots$ ) for tensors with  $d > 2$ . The  $(n_1, n_2, \ldots, n_d)$ -th entry of a  $d$ -dimensional tensor  $\mathcal{Y} \in \mathbb{R}^{N_1 \times N_2 \times \ldots \times N_d}$  is denoted by  $y = \mathcal{Y}[n_1, n_2, \ldots, n_d]$ , where  $n_i = 1, 2, \ldots, N_i$  ( $i = 1, 2, \ldots, d$ ), and  $N_i$  is a size of the  $i$ -th mode. The mode- $i$  slice of such tensor is denoted by  $\mathcal{Y}[n_1, \ldots, n_{i-1}, :, n_{i+1}, \ldots, n_d]$ .
tensors [13, 6, 7, 40]. One of the common methods for low-rank approximation is the tensor train (TT) decomposition [26], which allows bypassing the curse of dimensionality. Many useful algorithms (e.g., element-wise addition, multiplication, solution of linear systems, convolution, integration, etc.) have effective implementations for tensors given in the TT-format. The complexity of these algorithms turns out to be polynomial in the dimension  $d$  and the mode size  $N$ . It makes the TT-decomposition extremely popular in a wide range of applications, including computational mathematics [29, 2] and machine learning [31, 19].
In the last few years, new discrete optimization algorithms based on the TT-format have been proposed: TTOpt [37], Optima-TT [5], and several modifications [34, 35, 24, 36]. However, the development of new, more accurate, and robust TT-based methods for multidimensional discrete optimization is possible. In this work, we extend recent approaches for working with probability distributions and sampling in the TT-format [10, 25] to the optimization area. That is, we specify a multidimensional discrete probability distribution in the TT-format, followed by efficient sampling from it and updating its parameters to approximate the optimum in a better way. This makes it possible to build an effective optimization method, and the contributions of our work are as follows:
- We develop a new method PROTES for optimization (finding the minimum or maximum<sup>3</sup> value) of multidimensional data arrays and discretized multivariable functions based on a sampling method from a probability distribution in the TT-format;
- We apply<sup>4</sup> PROTES for various analytic model functions and for several multidimensional QUBO and optimal control problems to demonstrate its performance and compare it with popular discrete optimization algorithms (Particle Swarm Optimization, Covariance Matrix Adaptation, Differential Evolution, and NoisyBandit) as well as TT-based methods (TTOpt and Optima-TT). We used the same set of hyperparameters of our algorithm for all experiments and obtained the best result for 19 of the 20 problems considered.
2 Optimization with probabilistic sampling
Our problem is to minimize the given multivariable discrete black-box function  $f(1)$ . It can be formulated in terms of the multi-index search in an implicitly given  $d$ -dimensional tensor
$$
\mathcal {Y} \in \mathbb {R} ^ {N _ {1} \times N _ {2} \times \dots \times N _ {d}}, \quad \mathcal {Y} [ n _ {1}, n _ {2}, \dots , n _ {d} ] = \mathbf {f} (\boldsymbol {x}), \quad \boldsymbol {x} = [ n _ {1}, n _ {2}, \dots , n _ {d} ],
$$
for all  $n_i = 1,2,\ldots ,N_i$ $(i = 1,2,\dots ,d)$ . The essence of our idea is to use a probabilistic method to find the minimum  $\pmb{x}_{min}$ . We propose establishing a discrete distribution  $\mathsf{p}(\pmb {x})$  from which the minimum could be sampled with high probability. This distribution can be specified as a tensor  $\mathcal{P}_{\theta}\in \mathbb{R}^{N_1\times N_2\times \ldots \times N_d}$  in some low-parametric representation with a set of parameters  $\theta$  having the same shape as the target tensor  $\mathcal{Y}$ .
We start from a random non-negative tensor  $\mathcal{P}_{\theta}$  and iteratively perform the following steps until the budget is exhausted or until convergence (see graphic illustration in Figure 1):
1. Sample  $K$  candidates of  $\pmb{x}_{min}$  from the current distribution  $\mathcal{P}_{\theta} \colon \mathcal{X}_K = \{\pmb{x}_1, \pmb{x}_2, \dots, \pmb{x}_K\}$ ;
2. Compute the corresponding function values:  $y_{1} = \mathsf{f}(\pmb{x}_{1}), y_{2} = \mathsf{f}(\pmb{x}_{2}), \dots, y_{K} = \mathsf{f}(\pmb{x}_{K})$
3. Select  $k$  best candidates with indices  $\mathcal{S} = \{s_1, s_2, \ldots, s_k\}$  from  $\mathcal{X}_K$  with the minimal objective value, i.e.,  $y_j \leq y_J$  for all  $j \in \mathcal{S}$  and  $J \in \{1, 2, \ldots, K\} \setminus \mathcal{S}$ ;
4. Update the probability distribution  $\mathcal{P}_{\theta}$  ( $\theta \to \theta^{(new)}$ ) to increase the likelihood of selected candidates  $\mathcal{S}$ . We make several  $(k_{gd})$  gradient ascent steps with the learning rate  $\lambda$  for the tensor  $\mathcal{P}_{\theta}$ , using the following loss function
$$
\widehat {L} _ {\theta} \left(\left\{x _ {s _ {1}}, x _ {s _ {2}}, \dots , x _ {s _ {k}} \right\}\right) = \sum_ {i = 1} ^ {k} \log \left(\mathcal {P} _ {\theta} \left[ \boldsymbol {x} _ {s _ {i}} \right]\right). \tag {2}
$$
As will be shown in the work, we use several gradient ascent iterations for the likelihood to update the parameters of the TT-decomposition. However, it is important to note that the proposed optimization method is gradient-free, that is, it does not use the internal structure and gradient of the objective function.
Further, for concreteness, we will consider the minimization problem in this paper, while the proposed method can be applied to the discrete maximization problem without any significant modifications.
4The program code with numerical examples is available in the repository https://github.com/anabatsh/PROTES.
Figure 1: Schematic representation of the proposed optimization method PROTES.
After a sufficient number of iterations, we expect the tensor  $\mathcal{P}_{\theta}$  to represent an almost Kronecker delta-function with a pronounced peak in the value of the minimum of the target function (or several peaks if the minimum is not unique). Therefore, this value will be sampled during the steps of our algorithm since the probability of sampling other values will be sufficiently small.
From a low-parameter representation  $\mathcal{P}_{\theta}$  we expect an efficient sampling algorithm and efficient calculation procedure for the logarithms in (2) with automatic differentiation capability to enable gradient ascent methods. As will be shown below, the TT-representation of tensors satisfies these requirements. Further, we will omit the index  $\theta$ , assuming that the parameterized representation of the tensor  $\mathcal{P}$  corresponds to the TT-format. Note that the values  $K$ ,  $k$ ,  $k_{gd}$ ,  $\lambda$  and the number of parameters in  $\theta$  (i.e., rank of the TT-decomposition) are the hyperparameters of our algorithm.
3 Basic properties of the tensor train format
Let us dwell on the concept of the TT-format. A  $d$ -dimensional tensor  $\mathcal{P} \in \mathbb{R}^{N_1 \times N_2 \times \ldots \times N_d}$  is said to be in the TT-format [26] if its elements are represented by the following formula
$$
\mathcal {P} [ n _ {1}, n _ {2}, \dots , n _ {d} ] = \sum_ {r _ {1} = 1} ^ {R _ {1}} \sum_ {r _ {2} = 1} ^ {R _ {2}} \dots \sum_ {r _ {d - 1} = 1} ^ {R _ {d - 1}} \mathcal {G} _ {1} [ 1, n _ {1}, r _ {1} ] \mathcal {G} _ {2} [ r _ {1}, n _ {2}, r _ {2} ] \dots \mathcal {G} _ {d} [ r _ {d - 1}, n _ {d}, 1 ], \tag {3}
$$
where  $(n_{1}, n_{2}, \ldots, n_{d})$  is a multi-index ( $n_{i} = 1, 2, \ldots, N_{i}$  for  $i = 1, 2, \ldots, d$ ), integers  $R_{0}, R_{1}, \ldots, R_{d}$  (with convention  $R_{0} = R_{d} = 1$ ) are named TT-ranks, and three-dimensional tensors  $\mathcal{G}_{i} \in \mathbb{R}^{R_{i-1} \times N_{i} \times R_{i}}$  ( $i = 1, 2, \ldots, d$ ) are named TT-cores. The TT-decomposition (3) (see also an illustration in Figure 2) allows to represent a tensor or a discretized multivariable function in a compact and descriptive low-parameter form, which is linear in dimension  $d$ , i.e., it has less than  $d \cdot \max_{i=1,\ldots,d}(N_{i}R_{i}^{2}) \sim d \cdot \overline{N} \cdot \overline{R}^{2}$  parameters, where  $\overline{N}$  and  $\overline{R}$  are effective ("average") mode size and TT-rank respectively.
Linear algebra operations (e.g., element-wise addition, solution of linear systems, convolution, integration, etc.) on tensors in the TT-format, respectively, also have complexity linear in dimension if the TT-ranks are bounded. The TT-approximation for a tensor or discretized multivariable function may be built by efficient numerical methods, e.g., TT-SVD [27], TT-ALS [4], and TT-cross [28]. A detailed description of the TT-format and related algorithms are given in the works [26, 6, 7]. Below, we discuss only three operations in the TT-format, which will be used later in the work.
Construction of the random TT-tensor. In order to build a random non-negative TT-tensor of a given size  $(N_{1},N_{2},\ldots ,N_{d})$  with a constant TT-rank  $R$ , it is enough to generate  $d$  TT-cores  $\mathcal{G}_1,\mathcal{G}_2,\dots ,\mathcal{G}_d$  (3-dimensional tensors) with random elements from the uniform distribution on the interval  $(0,1)$ . We will refer to this method as tt_random  $(R,[N_1,N_2,\dots ,N_d])$ .
Figure 2: Schematic representation of the TT-decomposition. The top picture demonstrates the calculation of the specific tensor element  $\pmb{x} = [n_{1}, n_{2}, \dots, n_{d}]$  from its TT-representation, and the bottom picture presents the related tensor network diagram.
Computation of the log-likelihood in the TT-format. To calculate the logarithm  $\log \mathcal{P}[\pmb{x}]$  in a given multi-index  $\pmb{x} = (n_{1}, n_{2}, \dots, n_{d})$ , we can use the basic formula (3) and then take the logarithm of the result. It can be shown that this operation has complexity  $\mathcal{O}\left(d \cdot \overline{R}^{2}\right)$ , because, roughly speaking, we  $(d - 1)$  times multiply a vector of length  $\overline{R}$  by a matrix of size  $\overline{R} \times \overline{R}$  to get the result. The corresponding method will be called tt_log( $\mathcal{P}$ ,  $\pmb{x}$ ).
Sampling from the tensor in the TT-format. We use the approach proposed in the work [10] to generate a multi-index  $\pmb{x}$  with a probability proportional to the corresponding value  $p = \mathcal{P}[x]$  of the TT-tensor  $\mathcal{P}$ . The method is based on the sequential calculation of univariate conditional densities with efficient integration in the TT-format. The estimate for its complexity turns out to be the following:  $\mathcal{O}\Big(K\cdot d\cdot (\overline{N} +\overline{R})\cdot \overline{R} +K\cdot d\cdot \alpha (\overline{N})\Big)$ , where  $K$  is a number of requested samples, and  $\alpha (n)$  is a complexity of sampling from generalized Bernoulli distribution with  $n$  outcomes. Note that the algorithm allows sampling in the case of the initially non-normalized tensor, so we do not have to calculate the normalization factor. We will refer to this method as tt_sample(  $\mathcal{P},K$  ).
4 Optimization method PROTES
With the formal scheme of the proposed approach given in Section 2 and the description of operations tt_random, tt_log, and tt_sample given in Section 3, we can formulate our method PROTES for gradient-free discrete optimization in the TT-format, as presented in Algorithm 1. We denote as adam, a function that performs  $k_{gd}$  steps of gradient ascent for the TT-tensor  $\mathcal{P}$  at multi-indices  $\mathcal{X}$  by the well-known Adam method [18]. In this case, the learning rate is  $\lambda$ , the loss function is given in (2), and tt_log with automatic differentiation support is used for the log-likelihood computation.
Computational complexity of the method. Let us estimate the complexity of the proposed algorithm, assuming that the number of requests to the target function (black-box)  $M$  is fixed. With the known estimate for the complexity of the tt_sample function, we can obtain the complexity of the sampling operations:  $\mathcal{O}\left(\frac{M}{K} \cdot K \cdot d \cdot ((\overline{N} + \overline{R}) \cdot \overline{R} + \alpha(\overline{N}))\right)$ . Assuming that the complexity of one gradient step coincides with the complexity of calculating the differentiated function and using the estimate for the tt_sample function, we can estimate the total complexity of the tensor updates:  $\mathcal{O}\left(\frac{M}{K} \cdot k \cdot k_{gd} \cdot d \cdot \overline{R}^2\right)$ . Combining the two above estimates we obtain the complexity of the method
$$
\circ \left(M \cdot d \cdot \left(\frac {k}{K} \cdot k _ {g d} \cdot \bar {R} ^ {2} + (\bar {N} + \bar {R}) \cdot \bar {R} + \alpha (N)\right)\right). \tag {4}
$$
Algorithm 1 Method PROTES in the TT-format for multidimensional discrete black-box optimization
Data: the function  $\mathbf{f}(\pmb{x})$ , that computes the value of the target tensor  $\pmb{y} \in \mathbb{R}^{N_1 \times N_2 \times \dots \times N_d}$  at the multi-index  $\pmb{x} = [n_1, n_2, \dots, n_d]$ ; the maximum number of requests  $M$ ; the number of generated samples per iteration  $K$ ; the number of selected candidates per iteration  $k$ ; the number of gradient ascent steps  $k_{gd}$ ; the gradient ascent learning rate  $\lambda$ ; the TT-rank of the probability tensor  $R$ .
Result:  $d$ -dimensional multi-index  $\pmb{x}_{min}$ , which relates to the minimum value of the tensor  $\mathcal{Y}$ .
1 Initialize target multi-index and tensor value:  $\pmb{x}_{min} =$  None,  $y_{min} = \infty$
2 Generate random non-negative rank-  $R$  TT-tensor:  $\mathcal{P} = \mathrm{tt\_random}$  ( $R, [N_1, N_2, \dots, N_d]$ )
3 for iter  $= 1$  to  $M / K$  do
4 Generate  $K$  samples from  $\mathcal{P}$  ..  $\pmb {x}_1,\pmb {x}_2,\dots ,\pmb {x}_K = \mathrm{tt\_sample}(\mathcal{P},K)$
5 Compute related tensor values:  $y_{1} = \mathsf{f}(\pmb{x}_{1}), y_{2} = \mathsf{f}(\pmb{x}_{2}), \dots, y_{K} = \mathsf{f}(\pmb{x}_{K})$
6 Find indices  $\mathcal{S} = \{s_1,s_2,\dots ,s_k\}$  for top-  $k$  minimum items in the list  $[y_{1},y_{2},\ldots ,y_{K}]$
7 Collect k-top candidates:  $\mathcal{X} = \{\pmb{x}_{s_1},\pmb{x}_{s_2},\dots,\pmb{x}_{s_k}\}$ ,  $\mathcal{Y} = \{y_{s_1},y_{s_2},\dots,y_{s_k}\}$
8 If  $\mathcal{Y}$  contains a value less than  $y_{min}$ , then update  $\pmb{x}_{min}$  and  $y_{min}$ .
9 Update the TT-tensor:  $\mathcal{P}\gets \mathsf{adam}(\mathcal{P},\hat{L},\mathfrak{X},k_{gd},\lambda)$  // with the loss function (2) and the method tt_log
10 return  $x_{min}$
However, it is important to note that this estimate does not consider the complexity of calculating  $M$  times the objective function  $f$ , which in practical applications can be significant and many times greater than the estimate (4).
The intuition behind the method. The proposed method PROTES, like most gradient-free optimization approaches, is empirical; however, we can establish its connection with a well-known REINFORCE trick algorithm [43]. Let make a monotonic transformation  $\mathsf{F}[\mathsf{f}](\pmb{x})$  of the target function  $f$  to be minimized that transforms minimum to maximum. A reasonable choice for  $\mathsf{F}[\cdot ]$  is the Fermi-Dirac function
$$
\mathsf {F} [ \mathsf {f} ] (\boldsymbol {x}) = \frac {1}{\exp \left((\mathsf {f} (\boldsymbol {x}) - y _ {\min } - E) / T\right) + 1},
$$
where  $y_{\mathrm{min}}$  is an exact or approximate minimum of  $\mathsf{f}$ ,  $T > 0$  is a parameter and  $E > 0$  is some small threshold. With the function  $\mathsf{F}[\mathsf{f}]$  we can find a maximum of the expectation  $\max_{\theta} \mathsf{E}_{\xi_{\theta}} \mathsf{F}[\mathsf{f}](\xi_{\theta})$ , where a family of random variables  $\xi_{\theta}$  has a parameterised distribution function  $\mathsf{p}_{\theta}(\boldsymbol{x})$ . Using REINFORCE trick, we can estimate the gradient of the expectation by the following Monte-Carlo-like expression
$$
\nabla_ {\theta} \mathsf {E} _ {\xi_ {\theta}} \mathsf {F} [ \mathbf {f} ] (\xi_ {\theta}) \approx \frac {1}{M} \sum_ {i = 1} ^ {M} \mathsf {F} [ \mathbf {f} ] (\boldsymbol {x} _ {i}) \nabla_ {\theta} \log \mathsf {p} _ {\theta} (\boldsymbol {x} _ {i}), \tag {5}
$$
where  $\{\pmb{x}_i\}_{1}^{M}$  are independent realizations of the random variable  $\xi_{\theta}$ . If we find the optimal values of  $\theta$ , then we expect the optimal distribution  $\mathsf{p}_{\theta}$  to peak at the maximum point for function  $\mathsf{F}[\mathsf{f}]$ . Thus, we can obtain the argument of its maximum by sampling from this distribution. For very small values of  $T$ , only a few terms contribute to the sum (5), namely those  $\pmb{x}_i$  for which  $\mathsf{f}(\pmb{x}_i) - y_{min} < E$  is hold. For these values of  $\pmb{x}$ ,  $\mathsf{F}[\mathsf{f}]$  is close to 1, while for the other samples its value is 0. Hence, we can discard all other samples and keep a few with the best values. So, we come to the loss function (2), where instead of the parameter  $E$  we use a fixed number  $k$  of the best samples, i.e., the samples for which the value of the target function  $\mathsf{f}$  is the smallest.
Application of the method to constrained optimization. A very nice property of the proposed method is that it can be adapted to efficiently handle constraints such as a specified set of admissible multi-indices. One option is just to remove invalid samples from the top-k values, but in some cases, the probability of sampling multi-indices that are admissible is very low, so this approach will not work. Instead, if the constraint permits, we use the algorithm from the work [32] for the constructive building of tensors in the TT-format by a known analytic function, which defines the constraints. Once the indicator tensor (1 if the index is admissible and 0 if it is not) is built in the TT-format, we can just initialize the starting distribution  $\mathcal{P}$  by it, and it will be guaranteed that the samples almost always belong to the admissible set.
Below, we briefly analyze classical approaches for discrete optimization and then discuss the methods based on the low-rank tensor approximations, which have become popular in recent years.
Classical methods for gradient-free optimization. In many situations, the problem-specific target function is not differentiable, too complex, or its gradients are not helpful due to the non-convex nature of the problem [1], and standard well-known gradient-based methods cannot be applied directly. The examples include hyper-parameter selection, training neural networks with discrete weights, and policy optimization in reinforcement learning. In all these contexts, efficient direct gradient-free optimization procedures are highly needed. As for high-dimensional black-box optimization, evolutionary strategies (ES) [9] are one of the most advanced methods. This approach aims to optimize the parameters of the search distribution, typically a multidimensional Gaussian, to maximize the objective function. Finite difference schemes are commonly used to approximate gradients of the search distribution. Numerous works proposed techniques to improve the convergence of ES [23]; for example, second-order natural gradient [42] or the history of recent updates (Covariance Matrix Adaptation Evolution Strategy; CMA-ES) [15] may be used to generate updates. There is also a large variety of other heuristic methods for finding the global extremum. In particular, we note such popular approaches as NoisyBandit [33], Particle Swarm Optimization (PSO) [17], Simultaneous Perturbation Stochastic Approximation (SPSA) [22], Differential Evolution (DE) [38] and scrambled-Hammersley (scr-Hammersley) [14].
Tensor-based methods for gradient-free optimization. Recently, the TT-decomposition has been actively used for multidimensional optimization. An iterative method TTOpt based on the maximum volume approach is proposed in the work [37]. TTOpt utilizes the theorem of sufficient proximity of the maximum modulo element of the submatrix having the maximum modulus of the determinant to the maximum modulo element of the tensor. Based on this observation, tensor elements are sampled from specially selected successive unfoldings of the tensor. Dynamic mapping of the tensor elements is carried out to find the minimum element, which converts the minimum values into maximum ones. The authors applied this approach to the problem of optimizing the weights of neural networks in the framework of reinforcement learning problems in [37] and to the QUBO problem in [24]. A similar optimization approach was also considered in [34] and [35]. One more promising algorithm, Optima-TT, was proposed in recent work [5]. This approach is based on the probabilistic sampling from the TT-tensor and makes it possible to obtain a very accurate approximation for the optimum of the given TT-tensor. However, this method is intended for directly optimizing the TT-tensors, which means that its success strongly depends on the quality of the TT-approximation for the original multidimensional data array. Therefore, one of the related methods in the TT-format (TT-SVD, TT-ALS, TT-cross, etc.) should be additionally used for approximation.