SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hannah Dröge
University of Siegen 57076 Siegen, Germany hannah.droege@uni-siegen.de
Zorah Lahner
University of Siegen 57076 Siegen, Germany zorah.laehner@uni-siegen.de
Yuval Bahat
Princeton University  
Princeton, NJ 08544, United States  
yuval.bahat@gmail.com
Onofre Martorell
University of Balearic Islands  
Investigador ForInDoc del Govern de les Illes Balears  
07122 Palma, Illes Balears, Spain  
o.martorell@uib.cat
Felix Heide
Princeton University  
Princeton, NJ 08544, United States  
fheide@cs.princeton.edu
Michael Möller  
University of Siegen  
57076 Siegen, Germany  
michael.moeller@uni-siegen.de
Permutation matrices play a key role in matching and assignment problems across the fields, especially in computer vision and robotics. However, memory for explicitly representing permutation matrices grows quadratically with the size of the problem, prohibiting large problem instances. In this work, we propose to tackle the curse of dimensionality of large permutation matrices by approximating them using low-rank matrix factorization, followed by a nonlinearity. To this end, we rely on the Kissing number theory to infer the minimal rank required for representing a permutation matrix of a given size, which is significantly smaller than the problem size. This leads to a drastic reduction in computation and memory costs, e.g., up to 3 orders of magnitude less memory for a problem of size  $n = 20000$ , represented using  $8.4 \times 10^{5}$  elements in two small matrices instead of using a single huge matrix with  $4 \times 10^{8}$  elements. The proposed representation allows for accurate representations of large permutation matrices, which in turn enables handling large problems that would have been infeasible otherwise. We demonstrate the applicability and merits of the proposed approach through a series of experiments on a range of problems that involve predicting permutation matrices, from linear and quadratic assignment to shape matching problems.
Permutation matrices, which encode the reordering of elements, arise naturally in any problem that can be phrased as a bijection between two equally sized sets. As such, they are fundamental to many important computer vision applications, including matching semantically identical key points in images [50, 47, 48, 49], matching 3D shapes or point clouds [18, 46, 27], estimating scene flow on point clouds [35] and solving jigsaw puzzles [28], as well as to various sorting tasks [1, 16].
Figure 1: Geometric intuition behind our approach on a 2D unit sphere. For well-distributed vectors  $V \in \mathbb{R}^{\mathrm{Kiss}(2) \times 2}$ , where the number of vectors is determined by the Kissing number  $(\mathrm{Kiss}(2) = 6)$ , the cosine angle between different vectors  $V_{i,:}$  and  $V_{j,:}$ ,  $i \neq j$ , is  $\langle V_{i,:}, V_{j,:} \rangle = \cos(\alpha) \leq 0.5$ , while  $\langle V_{i,:}, V_{i,:} \rangle = 1$  for the same vector. Thus, for any permutation  $P$ , the matrix-matrix product of  $V$  and  $(PV)^T$  merely has to be thresholded suitably to represent the permutation  $P$ , i.e.  $P = 2\max(V(PV)^T - 0.5, 0)$ .
A permutation  $p$ , corresponding to the bijection from the set  $\{1, \dots, n\}$  onto itself, can be represented efficiently without a permutation matrix by merely enumerating the  $n$  elements
$$
(p (1), p (2), \dots , p (n)) \in \mathbb {N} ^ {n}. \tag {1}
$$
However, this representation is unsuitable for most computer vision problems that involve estimating  $p$  through optimization since this representation (i) is inherently discrete, yielding combinatorial problems for which no natural relaxation exists, and (ii) induces a solution space with a meaningless distance metric, as element  $i$  in the set generally is not 'closer' to element  $i + 1$  than it is to any other element  $j$ .
As a result, almost all methods for predicting permutations, including learning-based methods, favor a permutation matrix representation instead, i.e., formulating a permutation as an element in the set
$$
\mathcal {P} _ {n} = \left\{P \in \{0, 1 \} ^ {n \times n} \mid \sum_ {i} P _ {i j} = 1, \sum_ {j} P _ {i j} = 1 \forall i, j \right\}, \tag {2}
$$
with  $p(i) = j$  in representation (1) corresponding to  $P_{ij} = 1$  in the matrix representation form (2), which allows predicting  $p$  via optimization methods. Yet, the advantages of the matrix form representation (2) come at the cost of a prohibitive increase in memory, as it requires storing  $n^2$  binary numbers  $P_{ij} \in \{0,1\}$ , or - after commonly used relaxations - even  $n^2$ -many floating point numbers instead of the  $n$  integers in (1). This renders matching problems with  $n > \sim 10^4$  largely infeasible as their corresponding permutation matrix  $P$  constitutes over one hundred million entries.
To handle large matrices whose size prohibits explicit processing and storage, existing approaches typically either turn to sparse representations, i.e., storing only a small portion of matrix values in the form of  $(i,j,P_{i,j})$  triplets, where  $P_{i,j} \neq 0$ , or employ low rank representations, i.e., forming a large matrix  $P$  as a product of matrices
$$
P = V W ^ {T}, \tag {3}
$$
with  $V, W \in \mathbb{R}^{n \times m}$  and  $m < n$ . Unfortunately, neither of these approaches is applicable to permutation matrices: sparse representation cannot be used as the sparsity pattern is not only unknown a-priori but actually the sought-after solution to the problem. On top of that, since permutation matrices are by definition full rank, a low-rank representation (3) can yield only a crude approximation at best.
Contributions. In this work, we alleviate the limitation on problem size by harnessing the well-studied problem of (bounds for) the so-called Kissing number, which, in practice, translates to introducing a simple adaptation to the matrix factorization approach (3). In particular, we exploit the fact that for row-normalized matrices  $V$  and  $W$ , the entries of  $V W^T$  correspond to the cosines of the angles between the matrix rows. We then apply a pointwise non-linearity on the product of the matrices in (3), which allows representing any permutation while using  $m \ll n$ . We use the Kissing number theory to provide an estimate for how small an  $m$  we can use. We elaborate on these theoretical considerations in Section 3 and provide an illustration of the geometric intuition for our approach in Fig. 1.
We then demonstrate the applicability of the proposed approach through several numerical experiments tackling various problems that involve estimating permutations, including a study on point alignment, linear assignment problems (LAPs), quadratic assignment problems (QAPs), and a real-world shape matching application. We find that the proposed approach trades off only little accuracy to offer significant memory saving, thus enabling handling bijection mapping problems that are larger than was previously possible, when full permutation matrices had to be stored.
2.1 Permutation Learning and Representation
Permutation learning aims to develop a model capable of predicting the optimal permutation matrix that matches a given input data or label. Previous studies have suggested relaxing permutation matrices to continuous domains, such as the Birkhoff polytope [2], to approximate solutions for these optimization problems but still faced the problem of enforcing the sum-to-one constraints, which are commonly approximated using Sinkhorn layers [34]. In 2011, Adams et al. [34] proposed the Sinkhorn propagation method, with the goal of iteratively learning doubly-stochastic matrices that allow the prediction of permutation matrices. Following the same research line, Cruz et al. [38] proposed Sinkhorn networks, followed by Mena et al. [28] proposing Gumble-Sinkhorn networks, adding Gumbel noise to the Sinkhorn operator to learn latent permutations. Grover et al. [16] made further efforts in the relaxation of a learned permutation matrix and propose to relax the matrix to become unimodal row-stochastic, meaning that the rows must maintain a total sum of one while having a distinct argmax. More recent studies suggest to circumvent the constraint of row and column sums being one by predicting permutations in Lehmer code, whose matrix form results to be row-stochastic [9]. Other works propose alternative representations of permutations in different domains, as to work in the Fourier domain to give a compact representations of distribution over permutations [17, 20], or embed them to the surface of a hypersphere [33].
2.2 Nonlinear Matrix Decomposition
In previous work on nonlinear matrix decomposition, Saul [40] aims to find a low-rank matrix  $L$  that approximates a sparse nonnegative matrix  $P$  by applying an elementwise nonlinearity  $\sigma$  to  $L$ , i.e.  $P \approx \sigma(L)$ , and proposed an alternating minimization method to solve for  $L$  directly. In further work, [39], Saul analyzed the geometric relationship between the sparse and low-rank matrices, specifically for  $\sigma$  as a rectified linear unit (ReLU). For ReLU as nonlinearity, Seraghiti et al. [42] subsequently proposed an accelerated method of [40] to solve this problem, and a method, working with a version on  $L = VW$ , derived from the product of two matrices. In our work, we exploit this problem of low-rank approximation for permutation matrices and show its relationship to the Kissing number.
2.3 Assignment Problems
The goal of assignment problems is to find a permutation between two ordered sets while minimizing the assignment costs. The two most common versions are the linear and quadratic assignment problem (LAP and QAP) which are based on element-wise and pair-wise costs, respectively. An LAP can be solved in cubic time using the Hungarian algorithm [22]. The QAP was first introduced by Koopmans and Beckmann [21] and can be written as  $\min_{P\in \mathcal{P}_n}\mathrm{tr}(APBP^\top) + \mathrm{tr}(C^\top P)$ , for  $A,B\in \mathbb{R}^{n\times n}$  and  $C\in \mathbb{R}^{n\times n}$  being a problem where  $A\in \mathbb{R}^{n\times n}$  is the cost function between elements of the first object to match,  $B\in \mathbb{R}^{n\times n}$  is the distance function between elements in the second object and  $C\in \mathbb{R}^{n\times n}$  is a point-wise energy. The QAP has been proved to be NP-hard so no polynomial time solution can be expected for general cases. As a result, many relaxations of the problem exist, for example by relaxing the permutation constraint [14, 36], or by lifting the problem to higher dimensions [18, 51]. A survey on various relaxation approaches can be found in [24]. While the relaxations do ease some aspects of the problems, they normally do not decrease the dimensionality of the problem which remains demanding for large  $n$ .
2.4 3D Shape Correspondence
3D shape correspondence is also often posed as an assignment problem between the sets of vertices of a pair of shapes, for example through point descriptors matched by an LAP or in an QAP aiming to
preserve distances between all point pairs. However, 3D shapes are often discretized with thousands of vertices which makes optimization for a permutation computationally challenging. Hence, the permutation constraint is often relaxed [36] and, even though the tightness of relaxation might be known [1, 10], the optimization variables still scale quadratically with the problem size. In [13] and [46] the QAP is deconstructed into smaller problems and then each of them is optimized with a series of LAPs, while [41] solve for permutations as a series of cycles that gradually improve the solution.
Because permutation constraints for large resolution become infeasible, and, hence, the restriction to cases with the same number of vertices, recent methods often do not impose these constraints at all. The functional maps framework [30] converts the output to a mapping between functions on the shapes instead of vertices and can reduce the dimensionality drastically by using the isometry invariance property of the eigenfunctions of the Laplace-Beltrami operator [32]. Other lines of work rely on a given template to constrain the solution [15, 44], impose physical models that regularizes the deformation between inputs [11, 12, 31], or learn a solution space through training [8, 23, 26]. However, with the exception of template-based models, these cannot guarantee permutations.
3 Low-Rank Permutation Matrix Representation
A common approach to solve optimization problems with costs  $E$  over the set of permutation matrices  $\mathcal{P}_n$  (including those arising from training neural networks for predicting assignments) is to relax the problem by replacing  $\mathcal{P}_n$  by its convex hull  $\mathrm{conv}(\mathcal{P}_n)$ , i.e., the set of doubly-stochastic matrices:
$$
\min  _ {P \in \operatorname {c o n v} \left(\mathcal {P} _ {n}\right)} E (P). \tag {4}
$$
Since  $P$  grows quadratically in  $n$ , has an unknown sparsity pattern, and the true solution is always full rank, such problems pose significant challenges for large  $n$ . In this work, we make the interesting observation that a non-linearity as simple as a rectified linear unit (ReLU, denoted by  $\sigma$ ) is sufficient not only to restore a full rank, but to represent any permutation matrix exactly. More precisely, we propose to replace the set  $\mathrm{conv}(\mathcal{P}_n)$  in (4) with the set  $\mathcal{K}_m(\mathcal{P}_n) = \{\sigma(2VW^T - 1) \mid V, W \in \mathbb{R}^{n \times m}\}$  and use the so-called Kissing number [4, 29, 52] to show that  $\mathcal{P}_n \subset \mathcal{K}_m(\mathcal{P}_n)$  for a surprisingly small  $m$ . Let us first formalize our approach by defining the Kissing number:
Definition 1. For a given  $m \in \mathbb{N}$ , we define the Kissing number  $\operatorname{Kiss}(m)$  as
$$
K i s s (m) := \max  _ {n} \left\{n \in \mathbb {N} \mid \exists A \in \mathbb {R} ^ {n \times m}, \| A _ {i,:} \| _ {2} = 1, 2 \langle A _ {i,:}, A _ {j,:} \rangle \leq 1, i \neq j \right\}. \tag {5}
$$
Note that the Kissing number can be interpreted geometrically as the maximum number of points that can be distributed on an  $m$ -dimensional unit sphere such that the angle formed between each pair of different points is at least  $\arccos(0.5)$ . This property quickly establishes  $\mathcal{P}_n \subset \mathcal{K}_m(\mathcal{P}_n)$ :
Proposition 1. Let  $P \in \mathcal{P}_n$  be an arbitrary permutation matrix, and let  $\sigma, \sigma(x) = \max(x, 0)$  denote a rectified linear unit (ReLU). Then for every  $m$  such that  $n \leq \text{Kiss}(m)$  there exist  $V, W \in \mathbb{R}^{n \times m}$  such that
$$
P = \sigma \left(2 V W ^ {T} - 1\right). \tag {6}
$$
Proof. Let  $V \in \mathbb{R}^{n \times m}$  be a matrix that satisfies the equalities and inequalities of (5), and let  $W = PV$ . Then it holds that
$$
2 \left\langle V _ {i,:}, W _ {j,:} \right\rangle \left\{ \begin{array}{l l} \leq 1 & \text {i f} P _ {i, j} \neq 1 \\ = 2 & \text {o t h e r w i s e} \end{array} . \right. \tag {7}
$$
Consequently
$$
\sigma \left(2 \left\langle V _ {i,:}, W _ {j,:} \right\rangle - 1\right) = \left\{ \begin{array}{l l} 0 & \text {i f} P _ {i, j} \neq 1 \\ 1 & \text {o t h e r w i s e} \end{array} , \right. \tag {8}
$$
which proves the assertion.
To determine the minimal rank  $m$  that is required for representing a permutation of  $n$  elements, we rely on extensive studies in the past few decades which computed either exact values or lower and upper bounds for different values of  $m$  [7].
Using  $P = \sigma(2VW^T - 1)$  for relaxing (4) yields a relaxation that requires only  $2mn$  instead of  $n^2$  parameters, with  $m << n = \mathrm{Kiss}(m)$ . For instance,  $\mathrm{Kiss}(24) = 196560$  implies that matrices of rank  $m = 24$  are sufficient for representing any arbitrary permutation matrix of up to  $n = 196560$  elements, thus requiring  $\sim 4000$  times less storage memory:  $2 \cdot 24 \cdot 196560$  instead of  $196560^2$  parameters. Furthermore,  $\mathcal{P}_n \subset \mathcal{K}_m(\mathcal{P}_n)$  ensures that – in stark contrast to direct low-rank factorization – any permutation matrix can still be represented exactly. Empirically, the optimization over parametrizations  $\sigma(2VW^T - 1)$  turned out to cause significant challenges, likely due to the non-convexity and non-smoothness of the problem. To alleviate this problem, we resort to a smoother version of (6) which can still approximate permutations to an arbitrary desired accuracy:
Proposition 2. Let  $P \in \mathcal{P}_n$  and  $g$  denote an arbitrary permutation matrix and an arbitrary entrywise strictly monotonically increasing function, respectively, and let  $s$  denote the row-wise Softmax function  $s(A)_{i,j} = \frac{\exp A_{i,j}}{\sum_k \exp A_{i,k}}$ . Then  $\forall n \leq \text{Kiss}(m)$  and  $\forall \epsilon > 0$  there exist  $V, W \in \mathbb{R}^{n \times m}$  and  $\alpha > 0$ , such that
$$
\left\| P - s \left(\alpha g \left(V W ^ {T}\right)\right) \right\| \leq \epsilon .
$$
Proof. Similar to the proof in Proposition 1, we start by choosing  $V$  satisfying (5) and setting  $W = PV$  to obtain
$$
(V W ^ {T}) _ {i j} = \langle V _ {i,:}, W _ {j,:} \rangle \left\{ \begin{array}{l l} = 1 & \text {i f} P _ {i, j} = 1 \\ \leq 0. 5 & \text {o t h e r w i s e} \end{array} \right..
$$
Then  $\forall i, j, k$  s.t.  $P_{ij} = 1$  and  $k \neq j$  (i.e.,  $P_{ik} = 0$ ) it holds that  $g(VW^T)_{ij} > g(VW^T)_{ik}$ . Finally, to yield the assertion we use the Softmax property of converging to the unit vector in the limit  $s(\alpha A_{i,:}) \stackrel{\alpha \to \infty}{\rightarrow} e_j$  (with  $j = \arg \max A_{i,:}$ ), by taking  $\alpha > 0$  to be large enough.
In practice, we use  $g(x) = 2x$ , in accordance with the representation in (6). We use this smoother version to validate the proposed low-rank representation for handling large matching problems in the experiments we report next. Please notice that this is only a proof of existence of a decomposition into  $V, W$ , not every  $V, W$  lead to a permutation and an optimisation is not guaranteed to find a correct decomposition.
The following experiments validate our efficient permutation estimation method for different applications, and they confirm the ability to scale to very large problem sizes. First, as a proof of concept, we demonstrate our approach on the application of point cloud alignment for the two non-linearities proposed in Section 3 and introduce our sparse training technique. We then validate the effectiveness of our approach in the context of linear assignment problems and show how to handle sparse cost matrices. We perform further experiments in the context of generic NP-hard quadratic assignment problems, and integrate our approach into a state-of-the-art shape matching pipeline, thus providing the same level of accuracy while enabling a higher spatial resolution.
We use the PyTorch Adam optimizer [19] with its default hyperparameters in all our experiments. The estimation of  $V$  and  $W$  is performed in parallel.
Stochastic Optimization. Fully benefiting from our proposed compact representation requires the costs  $E$  (or an approximation thereof) to be evaluated without ever forming the full (approximate) permutation matrix, as this step would inherently return to necessitate  $n^2$  many entries. To this end, we introduce the concept of stochastic optimization, which – for our softmax-based representation  $s(2\alpha VW^T)$  arising from Proposition 2 – is not a stochastic training in a classical sense: we propose to fix all but two entries in each row of our approximate permutation. Specifically, in any supervised (learning-based) scenario where it is known that the  $y_i$ -entry of the  $i$ -th row of the final permutation  $P$  ought to be equal to one, each step of our optimizers merely computes the  $y_i$ -th and one randomly chosen  $(r_i$ -th entry) of each row, and computes the softmax  $s$  on these two entries only while implicitly assuming  $P_{i,j} = 0$  for  $j \notin \{y_i, r_i\}$ , i.e.,
$$
P _ {i, \left[ y _ {i}, r _ {i} \right]} = s \left(2 \alpha V _ {i,:} \left(W _ {\left[ y _ {i}, r _ {i} \right], :}\right) ^ {T}\right). \tag {9}
$$
In the above, we used  $W_{[y_i,r_i]}$ : to denote the  $2 \times m$  matrix consisting of the  $y_i$ -th and the  $r_i$ -th row of  $W$ . Our stochastic approach requires the computation of  $2n$  entries per gradient descent iteration only and - by randomly choosing the  $r_i$  - manages to still approximate the desired objective well.
Normalization of  $V$  and  $W$ . Since propositions 1 and 2 rely on row-normalized matrices, we explicitly enforce this constraint whenever we compute  $P$ , by using  $V_{i,:} \gets \frac{1}{\|V_{i,:}\|} V_{i,:}$ ,  $W_{i,:} \gets \frac{1}{\|W_{i,:}\|} W_{i,:}$ . We omit this step from the presentations below for the sake of readability.
Softmax Temperature. Since the values of  $\langle V_{i,:}, W_{j,:} \rangle$  are bounded by one following the aforementioned normalization, the temperature parameter  $\alpha$  determines the trade-off between approximating a hard maximum (as required for accurately representing permutations, see Prop.2) and favorable optimization properties (i.e., meaningful gradients). We specify the schedule (constant or monotonically increasing) in each of the experiments below.
4.2 Point Cloud Alignment
As a proof of concept, we demonstrate that our proposition is correct and the optimization process converges. We explore the different choices of non-linearity, starting with ReLU and continuing with Softmax, using the task of predicting a linear transformation over point clouds. In this task we aim to match a point cloud  $X_{1} \in \mathbb{R}^{n \times m}$  consisting of  $n$ $m$ -dimensional points, uniformly distributed on the unit hyper-sphere, to its linearly transformed and randomly permuted version  $X_{2} \in \mathbb{R}^{n \times m}$ . To obtain  $X_{2}$  we multiply  $X_{1}$  by a randomly drawn matrix  $\Theta_{\mathrm{GT}} \in \mathbb{R}^{m \times m}$  and apply a random permutation. Then, we optimize over the estimated transformation matrix  $\Theta$  which in this experiment defines our permutation matrix  $P(\Theta)$ :
$$
P (\Theta) = \sigma (2 V W (\Theta) ^ {T} - 1). \tag {10}
$$
Note that our representation in (6) is fully parameterized by  $\Theta$ , with  $V = X_{1}$  and  $W(\Theta) = X_{2}\Theta$ , and  $V$  and  $W$  are row-wise normalized in each iteration. Here  $P(\Theta)$  is equal to the correct permutation if  $\Theta$  correctly aligns the point clouds, i.e. minimizes the angle between corresponding points in  $V$  and  $W(\Theta)$  while maximizing the angles between non-corresponding points.
We solve for the permutation by performing 20000 minimizing steps with a learning rate set to 0.01 over the negative log-likelihood loss
$$
\hat {\Theta} = \underset {\Theta} {\arg \min } - \frac {1}{n} \sum_ {i = 1} ^ {n} \log \left(P (\Theta) _ {i, y _ {i}}\right), \tag {11}
$$
where  $y_{i}$  is the index of the point in  $X_{2}$  which corresponds to the  $i^{\mathrm{th}}$  point in  $X_{1}$ . We experiment with different numbers of points  $n$ , each time choosing the dimension  $m$  to be just big enough to satisfy the Kissing number constraint from Proposition 1, i.e.,  $\mathrm{Kiss}(m)\geq n > \mathrm{Kiss}(m - 1)$ . To check that we were able to find the correct transformation matrix  $\Theta$  and therefore the correct permutation matrix  $P$  through optimization, we verify that the nearest neighbor (closest point) for each row  $i$  in  $V$  is located in row  $j$  of matrix  $W$  that satisfies  $P_{i,j} = 1$ . We find that this is indeed the case in all experiments with different numbers of points  $n\in \{10,100,1000,10000\}$ , thus establishing that we could reach the correct permutation through optimization. We achieve equally good results when replacing the point-wise non-linearity ReLU with Softmax  $P(\Theta) = s\left(2\alpha VW(\Theta)^T\right)$ .