Posts | 🤖

Stereo Transformer & Sinkhorn Iterations

Fri, 08 Mar 2024 00:00:00 +0000

Revisiting Stereo Depth Estimation from a Sequence-to-Sequence Perspective with Transformers¹.

This may seem like an outdated paper from prehistoric 2021, especially in the contemporary high-speed deep learning race, but it’s full of original ideas:

Using transformer’s cross-attention for stereo matching without restriction to a disparity range and without the need for constructing highly memory-demanding discrete cost volumes.
Clever use of entropy-regularised optimal transport to impose a uniqueness constraint on disparity with the elegant differentiable Sinkhorn algorithm.
No finite and fixed discrete disparity values, instead a dense continuous disparity function allowing for a sub-pixel resolution.
Naturally handles occlusions by using a learnable parameter to assign the cost of an unmatched pixel.

Let’s scratch a little bit on the surface of the theory of optimal transport and see how it came to be useful in the modern deep learning workflow.

The optimal transport problem originated from Gaspard Monge’s 1781 paper on the optimality of moving a pile of dirt from one place to another. It minimises the total work required, taking into account the transportation cost, which is determined by the amount of mass moved and the distance taken:

This formulation is a special case of the optimal transport problem where the cost of transportation is given by Euclidean distance. It is commonly referred as Earth Mover’s Distance.

Sometimes we need to measure the distance between two distributions ( $ \pi $ and $ p $ ), for example in GANs or variational autoencoders.

For this purpose we already have a well known family of $ f $-divergencies $ D_f(\pi || p) = \mathbb{E}_{p(x)} f\left ( \frac{\pi(x)}{p(x)} \right )$, Kullback–Leibler or Jensen-Shannon to name a few. But for disjoint manifolds or distributions with disjoint supports they no longer represent a feasible distance as $ KL(\pi ||p ) = \infty $ and $ JSD(\pi || p) = \log(2) $, let alone proper gradient propagation.

Another application might be interpolation between distributions, where we can find a smooth transition expressed in barycentric coordinates:

Suppose we have two descrete distributions of size $ n $ and a cost matrix $ \mathbf{C} \in \mathbb{R}^{n \times n} $. Gaspard Monge’s optimal assignment problem looks for a mapping that associates to each point $ x_i $ a single point $ y_i $ or in other words it attempts to find a bijective function $ \sigma $ in the set of permutations of $ n $ elements to minimize:

$$ \min_{ \sigma \in \mathrm{Perm}(n) } \frac{1}{n} \sum_{i=1}^{n} \mathbf{C}_{i, \sigma(i)} . $$

This formulation of optimal assignment has some drawbacks: the problem itself is nonconvex, it may have several optimal solutions and sometimes a feasible solution does not exist. To overcome these inconveniences Kantorovich proposed to relax this deterministic nature of assignment. Originally a source point $ x_i $ can only be assigned to another point $ y_{\sigma_i} $, but Kantorovich suggested instead that the mass at any point $ x_i $ can be potentially dispatched across several locations or it can be splitted from a source towards several targets. This relaxation allows, in place of permutations $ \sigma $, to use a transportation matrix $ \mathbf{P} \in \mathbb{R}^{n \times m}_{+} $:

$$ \begin{split} \mu = \sum_{i=1}^{n} a_i \delta_{x_i} \; \nu = \sum_{j=1}^{m} b_j \delta_{y_j}. \\ \\ \mathbf{C} = [D(x_i, y_i)^p]_{ij} \\ \\ \mathbf{U}( \mathbf{a}, \mathbf{b} )= \{ \mathbf{P} \in \mathbb{R}^{n \times m}_{+} \; \| \; \mathbf{P} \mathbb{1}_m = \mathbf{a}, \mathbf{P}^{ \mathbf{T} } \mathbb{1}_n = \mathbf{b} \} \\ \\ \mathbf{L_{\mathbf{C}}}( \mathbf{a}, \mathbf{b} ) = \min_{ \mathbf{P} \in \mathbf{U}( \mathbf{a}, \mathbf{b})} \langle \mathbf{C}, \mathbf{P} \rangle \end{split} $$

where $ \mu, \nu $ — discrete measures with weights $ a_i, b_j $ at locations $ x_i, y_j $; $ \delta_x, \delta_y $ — Dirac functions; $ \mathbf{a}, \mathbf{b} $ — marginal source and target distrubutions; $ \mathbf{C} $ — a transporatation cost matrix or a matrix of pairwise distances; $ \mathbf{P} $ — a transportation plan matrix we are looking for, with a constraint that its columns and rows sums should match marginals.
$ P_{i,j} $ describes the amount of mass flowing from location $ i $ to location $ j $.

It looks like a linear program, so there must exist a dual which was formulated by Kantorovich (probably furnished him the Nobel Prize in economics):

$$ \begin{split} \mathbf{L}_{ \mathbf{C} } ( \mathbf{a} , \mathbf{b} ) = \max_{ \begin{split} \alpha \in \mathbb{R}^n, \beta \in \mathbb{R}^m \\ \alpha_i + \beta_j \le \mathbf{C}_{i,j} \end{split} } \alpha^T \mathbf{a} + \beta^T \mathbf{b} \end{split} $$

Here we are looking for optimal dual variables $ \alpha, \beta $ (Kantorovich potentials). We can write the Lagrangian of the primal problem:

$$ \min_{ \mathbf{P} \ge 0 } \max_{ \alpha \in \mathbb{R}^n, \beta \in \mathbb{R}^m } \langle \mathbf{C}, \mathbf{P} \rangle + \langle \mathbf{a} - \mathbf{P} \mathbb{1}_m, \alpha \rangle + \langle \mathbf{b} - \mathbf{P}^{\mathrm{T}} \mathbb{1}_n, \beta \rangle. $$

After exchanging the min and the max we obtain:

$$ \max_{ \alpha \in \mathbb{R}^n, \beta \in \mathbb{R}^m } \langle \mathbf{a}, \alpha \rangle + \langle \mathbf{b}, \beta \rangle + \min_{ \mathbf{P} \ge 0 } \langle \mathbf{C} - \alpha \mathbb{1}_m^{\mathrm{T}} - \mathbb{1}_n \beta^{\mathrm{T}}, \mathbf{P} \rangle. $$

We can observe that for any $ \mathbf{Q} $:

$$ \min_{\mathbf{P} \ge 0} \langle \mathbf{Q}, \mathbf{P} \rangle = \begin{cases} 0 & \text{ if } \mathbf{Q} \ge 0, \\ -\infty & \text{ otherwise } \end{cases}, $$

so that the constraints definetely are $ \mathbf{C} - \alpha \mathbb{1}_m^{\mathrm{T}} - \mathbb{1}_n \beta^{\mathrm{T}} \ge 0 $ or $ \alpha_i + \beta_j \le \mathbf{C}_{i, j} $.

Returning to the original problem, we can see that the set of matrices $ \mathbf{U} ( \mathbf{a}, \mathbf{b}) $ is defined by $ n+m $ equality constraints that form a convex polytope of feasible solutions with a cost vector $ \mathbf{C} $ and the optimal solution $ \mathbf{P}^{\star} $ located at a vertex:

The problem could be solved with a network flow solver with $ O(n^3log(n)) $ complexity². But unfortunately, this solution is not differentiable on discrete inputs, making it unsuitable for use in a deep learning pipeline. To complicate things more, the solution $ \mathbf{ P } ^{\star} $ is not stable upon small changes in the input distrubutions and not always unique when the cost vector $ \mathbf{C} $ is orthogonal to a face of the polytope:

To address these inconveniences, sacrifices must be made and the particular one is giving up on the exact solution.

The idea of the entropic regularization of optimal transport is to use the entropy of the transport matrix as a regularisation function to obtaine an approximate solution to the original problem³.

$$ \mathbf{L}_\mathbf{C}^{ \gamma } = \min_{ \mathbf{P} \in \mathbf{U}( \mathbf{a}, \mathbf{b} )} \langle \mathbf{P}, \mathbf{C} \rangle - \gamma \mathbf{H}(\mathbf{P}) , $$

where

$$ \mathbf{H}(\mathbf{P}) = -\sum_{i,j}P_{i,j}( \log P_{i,j} - 1 ) $$

is the discrete entropy of the transport matrix, $ \gamma $ — determines the extent to which entropy influences the solution.

This augmented objective function has a solution which is no longer located on a vertex of the polytope. The entropy term pushes the original solution away from the boundary to the regularized entropic minimum. You can see the impact of increasing $ \gamma $ on the cost landscape of the linear problem:

Formally the unique solution converges to the optimal one with the maximum entropy:

$$ \mathbf{P}_{\gamma} \xrightarrow{\gamma \rightarrow 0 } \mathop{\mathrm{argmin}}\limits_{\mathbf{P}} \{ -\mathbf{H}(\mathbf{P}) : \mathbf{P} \in \mathbf{U}(\mathbf{a}, \mathbf{b}), \langle \mathbf{P}, \mathbf{C} \rangle = \mathbf{L}_{\mathbf{C}}( \mathbf{a}, \mathbf{b} ) \}. $$

In the edge cases, when $ \gamma \rightarrow 0 $, it converges to the linear program solution:

$$ \mathbf{P}_{\gamma} \xrightarrow{\gamma \rightarrow 0 } \mathbf{L}_{\mathbf{C}}( \mathbf{a}, \mathbf{b} ), $$

and for a large regularization $ \gamma \rightarrow \infty $, the solution converges to the joint probability between two independent random variables distributed according to marginals $ \mathbf{a} $ and $ \mathbf{b} $:

$$ \mathbf{P}_{\gamma} \xrightarrow{\gamma \rightarrow \infty } \mathbf{a} \mathbf{b}^{\mathrm{T}}. $$

A natural way to solve a constrained minimization problem is to make a Lagrangian:

$$ \mathcal{L}(P, \alpha, \beta) = \sum_{i,j} P_{ij} C_{ij} + \gamma P_{i,j}( \log P_{i,j} - 1 ) + \alpha^{\mathrm{T}} ( \mathbf{P} \mathbb{1}_m - \mathbf{a} ) + \beta^{\mathrm{T}} ( \mathbf{P}^{\mathrm{T}} \mathbb{1}_n - \mathbf{b}). $$

After having assigned derivatives over components $ P_{i,j} $ to zero, we can express the solution for $ \mathbf{P} $:

$$ \begin{split} \frac{ \partial{ \mathcal{L} } }{ \partial{ P_{i,j} } } = C_{i,j} + \gamma (\log P_{i,j} - 1) + \gamma P_{i,j} \times \frac{1}{P_{i,j}} + \alpha_i + \beta_j = 0. \\ \log P_{i,j} = -\frac{\alpha_i}{\gamma} - \frac{C_{i,j}}{\gamma} - \frac{\beta_j}{\gamma}. \\ P_{i,j} = \exp( -\frac{\alpha_i}{\gamma} ) \exp( -\frac{C_{i,j}}{\gamma} ) \exp( -\frac{ \beta_j }{\gamma} ) \\ \end{split} $$

If we assume:

$$ u_i = \exp( -\frac{\alpha_i}{\gamma} ), \; K_{i,j} = \exp( -\frac{C_{i,j}}{\gamma} ), \; v_j = \exp( -\frac{ \beta_j }{\gamma} ) $$

we can rewrite the expression in a vector form:

$$ \mathbf{P} = \mathrm{diag}(\mathbf{u}) \mathbf{K} \mathrm{diag}(\mathbf{v}), $$

$ \mathbf{K} $ is called Gibbs kernel and $ \mathbf{u}, \mathbf{v} $ are Sinkhorn scalings.

Let’s use the normalization constraints $ \mathbf{P} \mathbb{1}_m = \mathbf{a}, \mathbf{P}^{\mathrm{T}} \mathbb{1}_n = \mathbf{b} $:

$$ \begin{split} \mathbf{P} \mathbb{1}_m = \mathrm{diag}(\mathbf{u}) \mathbf{K} \mathrm{diag}(\mathbf{v}) \mathbb{1}_m = \mathrm{diag}(\mathbf{u}) \mathbf{K} \mathbf{v} = \mathbf{a} \\ \mathbf{P}^{\mathrm{T}} \mathbb{1}_n = \mathrm{diag}(\mathbf{v}) \mathbf{K}^{\mathrm{T}} \mathrm{diag}(\mathbf{u}) \mathbb{1}_n = \mathrm{diag}(\mathbf{v}) \mathbf{K}^{\mathrm{T}} \mathbf{u} = \mathbf{b} \end{split} $$

or simpler:

$$ \begin{split} \mathbf{u} \odot \mathbf{K} \mathbf{v} = \mathbf{a} \\ \mathbf{v} \odot \mathbf{K}^{\mathrm{T}} \mathbf{u} = \mathbf{b} \end{split} $$

where $ \odot $ is Hadamard product. Solving this for $ \mathbf{u} $ and $ \mathbf{v} $ gives us an iterative solution know as famous Sinkhorn algorithm⁴ with $ O(nm) $ complexity⁵:

$$ \begin{split} \mathbf{u} = \mathbf{a} \oslash (\mathbf{K} \mathbf{v}) \\ \mathbf{v} = \mathbf{b} \oslash (\mathbf{K}^{\mathrm{T}} \mathbf{u}) \\ \end{split} $$

where $ \oslash $ stands for an element-wise division.

We can see the impact of the regularzation parameter $ \gamma $ on the solution, wich behaves like a temperature:

If we let the temperature approach zero, we approximate the exact solution. Conversely, if we let $ \gamma $ go to infinity, the solution is just an outer product of the two marginal distributions.
So Sinkhorn procedure involves a sequence of vector matrix multiplication that is perfectly differentiable. This allows it to be used within a neural network, taking advantage of GPU acceleration and PyTorch auto-differentiation.

And this is a suitable place to return to Stereo Transformer paper.
The achitecture of the network consists of a feature extractor backbone, a bunch of self- and cross-attention Transformer blocks and finally, a regression head with the Context Adjustment Layer:

Image credit: Li et al. Revisiting Stereo Depth Estimation From a Sequence-to-Sequence Perspective with Transformers

The optimal transport finds its place in the last transformer’s cross-attention block and it is replacing the Softmax operator usually used in the Transformer achitecture.

The optimal transport problem is formulated as follows:

The cost matrix $ \mathbf{C} $ is set to the negative cross-attention. Softmax is not used as optimal transport normalizes the attention values.
The two marginal distributions $ \mathbf{a} $ and $ \mathbf{b} $ are uniform, which imposes a soft uniquness constraint and mitigates ambiguity.

If the two marginal distribution have equal length and $ \mathbf{a} = \mathbf{b} = \mathbb{1}_n / n $, then the optimal assignment problem becomes the optimal matching problem and there exists an optimal solution $ \mathbf{P}_{\sigma^{\star} } $, which is a permutation matrix associated to an optimal permutation $ \sigma^{\star} \in \mathrm{Perm}(n) $³.

Additional boilerplate is done to add a learnable dustbin $ \phi $ responsible for unmached occluded pixels. The intuition behind the optimal transport solution $ \mathbf{P} $ is that entries $ P_{i,j} $ could be interpreted as pair-wise pixel matching probabilities.

Let’s examine the code and analyse how it was implemented:

def _optimal_transport(self, attn: Tensor, iters: int):
 """
 Perform Differentiable Optimal Transport in Log-space for stability
 Adapted from SuperGlue (https://github.com/magicleap/SuperGluePretrainedNetwork)

 :param attn: raw attention weight, [N,H,W,W]
 :param iters: number of iterations to run sinkhorn
 :return: updated attention weight, [N,H,W+1,W+1]
 """
 bs, h, w, _ = attn.shape

 # set marginal to be uniform distribution
 marginal = torch.cat([torch.ones([w]), torch.tensor([w]).float()]) / (2 * w)
 log_mu = marginal.log().to(attn.device).expand(bs, h, w + 1)
 log_nu = marginal.log().to(attn.device).expand(bs, h, w + 1)

 # add dustbins
 similarity_matrix = torch.cat([attn, self.phi.expand(bs, h, w, 1).to(attn.device)], -1)
 similarity_matrix = torch.cat([similarity_matrix, self.phi.expand(bs, h, 1, w + 1).to(attn.device)], -2)

 # sinkhorn
 attn_ot = self._sinkhorn(similarity_matrix, log_mu, log_nu, iters)

 # convert back from log space, recover probabilities by normalization 2W
 attn_ot = (attn_ot + torch.log(torch.tensor([2.0 * w]).to(attn.device))).exp()

 return attn_ot

The iterative process is done in logarithmic space for better numerical stability.
The two marginals $ \nu $ and $ \mu $ are initialized as log of the uniform distribution.
The shape of the attention tensor is [bs, n, h, w, w], the shape of marginals is [bs, h, w + 1], the last value here is for the dustbin $ \phi $.
All the values in the range [1:w] are initialized to 1/(2*w) and the last value is 0.5, so the whole distribution sums up to 1. Increased weight on the dustbin here is for balancing out its influence with the disparity bins.

 # set marginal to be uniform distribution
 marginal = torch.cat([torch.ones([w]), torch.tensor([w]).float()]) / (2 * w)
 log_mu = marginal.log().to(attn.device).expand(bs, h, w + 1)
 log_nu = marginal.log().to(attn.device).expand(bs, h, w + 1)

Then a similarity_matrix (optimal transport matrix) is constructed by adding the learnable dustbin phi to the last row and column of the attention tensor:

 # add dustbins
 similarity_matrix = torch.cat([attn, self.phi.expand(bs, h, w, 1).to(attn.device)], -1)
 similarity_matrix = torch.cat([similarity_matrix, self.phi.expand(bs, h, 1, w + 1).to(attn.device)], -2)

phi is a learnable scalar parameter previously defined in the constructor:

 self.phi = nn.Parameter(torch.tensor(0.0, requires_grad=True)) # dustbin cost

The last step is to run a few Sinkhorn interations and recover the transport matrix by exponentiation. Number of iterations is set to 10:

 # sinkhorn
 attn_ot = self._sinkhorn(similarity_matrix, log_mu, log_nu, iters)

 # convert back from log space, recover probabilities by normalization 2W
 attn_ot = (attn_ot + torch.log(torch.tensor([2.0 * w]).to(attn.device))).exp()

Let’s look into _sinkhorn method:

def _sinkhorn(self, attn: Tensor, log_mu: Tensor, log_nu: Tensor, iters: int):
 """
 Sinkhorn Normalization in Log-space as matrix scaling problem.
 Regularization strength is set to 1 to avoid manual checking for numerical issues
 Adapted from SuperGlue (https://github.com/magicleap/SuperGluePretrainedNetwork)

 :param attn: input attention weight, [N,H,W+1,W+1]
 :param log_mu: marginal distribution of left image, [N,H,W+1]
 :param log_nu: marginal distribution of right image, [N,H,W+1]
 :param iters: number of iterations
 :return: updated attention weight
 """

 u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
 for idx in range(iters):
 # scale v first then u to ensure row sum is 1, col sum slightly larger than 1
 v = log_nu - torch.logsumexp(attn + u.unsqueeze(3), dim=2)
 u = log_mu - torch.logsumexp(attn + v.unsqueeze(2), dim=3)

 return attn + u.unsqueeze(3) + v.unsqueeze(2)

Firstly vectors $ \mathbf{u} $ and $ \mathbf{v} $ are zero initialized. Then the cost matrix $ \mathbf{C} $ is set to the negative attention attn, which is intuitively understandable: the most attended pixels have the least transportation cost. The regularization parameter $ \gamma = 1 $.

With these assumptions in mind we can rewrite the kernel $ \mathbf{K} $ as:

$$ \mathbf{K} = \exp( -\frac{\mathbf{C}}{ \gamma} ) = \exp( \mathrm{attn} ). $$

Plugging this kernel into the Sinkhorn update and taking log of the equality, we achieve the original expression:

$$ \begin{align} \log(\mathbf{v}) &= \log(\nu) - \log(\sum_i \mathbf{K}_{i,j} u_i ) \\ &= \log(\nu) - \log( \sum_i \exp( \mathrm{attn}_{i,j} + \log( u_i ) ) ), \\ \log(\mathbf{u}) &= \log(\mu) - \log(\sum_j \mathbf{K}_{i,j} v_j ) \\ &= \log(\mu) - \log( \sum_j \exp( \mathrm{attn}_{i,j} + \log( v_j ) ) ). \end{align} $$

This conсludes the matter. It is worth mentioning that in the ablation studies section, it was demonstrated that the soft uniqueness constraint imposed by optimal transport improved all the metrics: end-point-error (EPE), 3px error and occlusion IOU.

Finally, I would like to express my gratitude to the authors for their outstanding work and for providing a clear example of how Optimal Transport can effectively complement existing deep learning techniques.

Homography from the Correspondences of Lines

Thu, 08 Feb 2024 00:00:00 +0000

Recently I stumbled upon a task where I had to find a mapping between a camera and a target plane, but with a little twist. I could define line equations in the plane, and once I had captured the image and found the corresponding lines, the problem boiled down to calculating the homography based on line correspondences. Following the famous solution from the Multiple View Geometry book ¹, I arrived at a simple result.

Let’s assume we have a line in the target plane $ L_i $ given by

$$ \begin{equation} L_i = ( a_i^p, b_i^p, c_i^p )^T :\ L_i^T \mathbf{X}^p = 0, \qquad \qquad (1) \label{ref:eq1} \end{equation} $$

where $ \mathbf{X}^p = \begin{pmatrix} x^p \\ y^p \\ 1 \end{pmatrix} $ — point in the target plane. It’s projection in the image we denote as $ l_i $:

$$ l_i = (a_i^i, b_i^i, c_i^i)^T :\ l_i^T \mathbf{ X}^i = 0, $$

$ \mathbf{X}^i = \begin{pmatrix} x^i \\ y^i \\ 1 \end{pmatrix} $ — being a point in the image.
A homography $ \mathbf{H} $ takes a point in the projector plane $ \mathbf{X}^p $ and sends it to the image plane: $$ \mathbf{X}^i = \mathbf{H} \mathbf{X}^p. $$

So for any point in the laser plane satisfying $ (1) $ we can write:

$$ l_i^T \mathbf{H} \mathbf{X}^p = 0. $$

Line $ L_i $ may be represented by any two points $ \mathbf{X}_0^p $ and $ \mathbf{X}_1^p $, so we have two equations for $ L_i \rightarrow l_i $ correspondence:

$$ \begin{equation} \begin{split} l_i^T \mathbf{H} \mathbf{X}_0^p &= 0, \qquad \qquad (2) \\ l_i^T \mathbf{H} \mathbf{X}_1^p &= 0. \end{split} \label{ref:eq2} \end{equation} $$

If the $ j $-th row of the matrix H is denoted by $ \mathbf{h}^{iT} $, then we may write:

$$ \mathbf{H} \mathbf{X}^p = \begin{pmatrix} \mathbf{h}^{1T} \mathbf{X}^p \\ \mathbf{h}^{2T} \mathbf{X}^p \\ \mathbf{h}^{3T} \mathbf{X}^p \end{pmatrix} $$

We can rewrite $ (2) $ in terms of $ \mathbf{h}^{iT} $:

$$ \begin{equation} \begin{split} a_i^i \mathbf{h}^{1T} \mathbf{X}_0^p + b_i^i \mathbf{h}^{2T} \mathbf{X}_0^p + c_i^i \mathbf{h}^{3T} \mathbf{X}_0^p &= 0, \qquad (3) \\ a_i^i \mathbf{h}^{1T} \mathbf{X}_1^p + b_i^i \mathbf{h}^{2T} \mathbf{X}_1^p + c_i^i \mathbf{h}^{3T} \mathbf{X}_1^p &= 0. \end{split} \label{ref:eq3} \end{equation} $$

We can take any two points on $ L_i $, but for the simplicity we may choose these ones:

$$ \begin{equation} \begin{split} \mathbf{X}_0^p = \begin{pmatrix} -\dfrac{c_i^p}{a_i^p} \\ 0\\ 1 \end{pmatrix}, \quad \mathbf{X}_1^p = \begin{pmatrix} 0 \\ -\dfrac{c_i^p}{b_i^p} \\ 1 \end{pmatrix}. \end{split} \qquad \qquad (4) \label{ref:eq4} \end{equation} $$

Plugging $ (4) $ into $ (3) $ and having components of $ \mathbf{H} = \begin{pmatrix} h_{00} & h_{01} & h_{02} \\ h_{10} & h_{11} & h_{12} \\ h_{20} & h_{21} & h_{22} \end{pmatrix} $ we are obtaining:

$$ \begin{equation} \begin{split} a_i^i ( -\dfrac{c_i^p}{a_i^p} h_{00} + h_{02} ) + b_i^i ( -\dfrac{c_i^p}{a_i^p} h_{10} + h_{12} ) + c_i^i ( -\dfrac{c_i^p}{a_i^p} h_{20} + h_{22} ) &= 0, \\ a_i^i ( -\dfrac{c_i^p}{b_i^p} h_{01} + h_{02} ) + b_i^i ( -\dfrac{c_i^p}{b_i^p} h_{11} + h_{12} ) + c_i^i ( -\dfrac{c_i^p}{b_i^p} h_{21} + h_{22} ) &= 0. \end{split} \end{equation} $$

After collecting terms around $ h $ we can write:

$$ \begin{equation} \begin{split} \begin{pmatrix} -a_i^i \dfrac{c_i^p}{a_i^p} & 0 & a_i^i & -b_i^i \dfrac{c_i^p}{a_i^p} & 0 & b_i^i & -c_i^i \dfrac{c_i^p}{a_i^p} & 0 & c_i^i \\ 0 & -a_i^i \dfrac{c_i^p}{b_i^p} & a_i^i & 0 & -b_i^i \dfrac{c_i^p}{b_i^p} & b_i^i & 0 & -c_i^i \dfrac{c_i^p}{b_i^p} & c_i^i \end{pmatrix} \times \\ \times \begin{pmatrix} h_{00} & h_{01} & h_{02} & h_{10} & h_{11} & h_{12} & h_{20} & h_{21} & h_{22} \end{pmatrix}^T = 0 \end{split} \end{equation} $$

or simply:

$$ \begin{equation} \begin{split} \mathbf{A}_i \mathbf{h} = 0. \end{split} \end{equation} $$

Each line correspondence gives two constraints on $ \mathbf{H} $. Given a set of such line correspondences, we obtain a set of equations $ \mathbf{A} \mathbf{h} = 0 $. To solve this, one might use SVD of $ \mathbf{A} $ and the unit singular vector corresponding to the smallest singular value is the solution $ \mathbf{h} $. Note, that lines shouldn’t be all mutually collinear, otherwise we would have an underdetermined system.

Richard Hartley, Andrew Zisserman. Multiple View Geometry in Computer Vision. 2003. ↩︎

LoRa and Projection to a Matrix Subspace

Fri, 05 Jan 2024 00:00:00 +0000

Low-Rank Adaptation technique was introduced to adapt large language models to a specific downstream task while keeping the original model intact and, by doing that, dramatically reducing the number of trainbale parameters. No only it lowers the hardware requirments to implement a task-specific adaptation but also it adds no inference latency when deployed.

Based on the assumption that pretrained language models have a low intrinsic dimension, authors constrain the update of a pre-trained weight matrix $ W_q \in \mathbb{R}^{d \times k } $ with a low-rank decomposition: $$ h = W_q x + \Delta W_q x = W_q x + B A x, $$ where $ B \in \mathbb{R}^{d \times r}, A \in \mathbb{R}^{r \times k}$ , and the adapted rank $ r \ll min\{ d, k \} $. In the paper they applied LoRa only to the attention weights of the Transformer architecture. This is why we see in the adaptation here a Transformer’s query matrix $W_q$.

To gain a deeper understanding of the relationship between the increment $ \Delta W $ and the original $ W $ matrices, it seems reasonable to project $W$, onto a subspace of $ \Delta W $ and to compute the Frobenius norm $ \| U^T W_q V^T \|_F $, where $ U $ and $ V $ are the left and the right singular-vector matrices of $ \Delta W $. Let’s see the reason behind this.

Suppose we have a matrix $ A $ with linearly independent comlumns. Then the projection matrix onto subspace of $ A $ would be:

$$ P = A(A^T A)^{-1}A^T. $$

In our case, if we want to project on the subspace of $ \Delta W_q $, we can express the projection matrix as (bearing in mind that SVD of $ \Delta W_q $ is $ U \Sigma V^T $): $$ \begin{split} P = \Delta W_q ( \Delta W_q^T \Delta W_q )^{-1} \Delta W_q^T = \\ U \Sigma V^T ( (U \Sigma V^T)^T U \Sigma V^T )^{-1} (U \Sigma V^T)^T = \\ U \Sigma V^T ( V \Sigma U^T U \Sigma V^T )^{-1} V \Sigma U^T = \\ U \Sigma V^T ( V \Sigma^2 V^T )^{-1} V \Sigma U^T = \\ U \Sigma V^T V^{-T} \Sigma^{-2} V^{-1} V \Sigma U^T = \\ U \Sigma \Sigma^{-2} \Sigma U^T = U U^T. \end{split} $$

And finally, the projection itself of matrix $ W_q $ is $ P W_q = U U^T W_q $ or $ U U^T W_q V^T V $, as matrix $ V $ is unitary, $ V^TV = I $.
As Frobenius norm is invariant under unitary operations $ U $ and $ V $ (follows from the cyclic nature of the trace), we can state that $ \| U U^T W_q V^T V \|_F = \| U^T W_q V^T \|_F $. From where comes the Feature amplificaiton factor mentioned in the paper:

$$ \frac{ \left \| \Delta W_q \right \|_F }{ \left \| U^T W_q V^T \right \|_F } $$ This factor demonstrates the interaction between subspaces and helps to understand what the adaptation is doing. For instance, experimental evaluation revealed that the matrix $ \Delta W $ amplifies singular directions that are not emphasized in the original matrix $W$ and the amplification factor is quite significant. It can be deduced from this that

the low-rank adaptation matrix potentially amplifies the important features for specific downstream tasks that were learned but not emphasized in the general pre-training model ¹.

Another tool that can shed light on the mechanics behind the adaptation is the measurement of similarity between subspaces with different adaptation ranks $ r $. This quantity can be calculated by using a normalized subspace similarity based on the Grassmann distance.

Suppose we have two adaptation matrices $ A_{r=8} $ and $ A_{r=64} $ with different ranks. Then we compute SVD for each of them and obtain the right singular matrices $ U_{r=8} $ and $ U_{r=64} $. The proposed in the paper similarity metric $ \phi $ helps to understand how much of the subspace spanned by the top $ i $ singular vectors in $ U_{r=8} $ is contained in the subspace spanned by top $ j $ singular vectors of $ U_{r=64} $:

$$ \phi(A_{r=8}, A_{r=64}, i, j) = \frac{ \| U_{A_{r=8}}^{iT} U_{A_{r=64}}^{j} \|_F^2 }{ min(i,j) }, $$

where $ U_{A_{r=8}}^i $ represents the columns of $ U_{A_{r=8}} $ corresponding to the top $ i $ singular vectors. This similarity metric resembles very much the well-known Projection metric which measures the distance between subspaces.

According to the definition, the Grassmann manifold $ \mathcal{G}(m,D) $ is a set of $ m $-dimensional linear subspaces of the Euclidean space $ \mathbb{R}^D $. Let’s take any two elements of the manifold $ \mathcal{G}(m,D) $ — the orthonormal matrices $ A $ and $ B $ of size $ (D \times m) $. In the same time $ span(A) $ and $ span(B) $ represent two subspaces in $ \mathbb{R}^D $. The distance between two subspaces $ span(A) $ and $ span(B) $ can be measured using the pricipal angles $ [\theta_1, \dots, \theta_m] $ as the usual dot-product of vectors. In the Grassmann manifold however, these two subspaces are represented simply by two points and the Riemannian distance between them is simply $ d(A, B) = \| \theta \|_2 $.
The set of principal angles $ 0 \le \theta_1 \le \dots \le \theta_m \le \pi/2 $ between two subspaces $ span(A) $ and $ span(B) $ are defined by $$ cos \theta_k = \max_{u_k \in span(A)} \max_{v_k \in span(B)} u_k^Tv_k. $$

where $ u_k $ spans $ A $, $ v_k $ spans $ B $ and $$ \begin{split} u_k^T u_k = 1, v_k^T V_k = 1, \\ u_k^T u_i = 0, v_k^T v_i = 0, ( i = 1, \dots, k-1 ). \end{split} $$

In other words, the first principal angle $ \theta_1 $ is the smallest angle between all pairs of unit vectors in the first and the second subspaces ².

To compute these principal angles we need to find the SVD of $ A^TB $:

$$ A^TB = U diag(\sigma_1, \dots, \sigma_m) V^T = U(cos\Theta)V^T, $$

where $ U = [u_1, \dots, u_m], V = [v_1, \dots, v_m] $ and $ cos\Theta $ is the diagonal matrix $ cos \Theta = diag(cos\theta_1, \dots, cos\theta_m ) $.
Another popular expression of a distance betwen subspaces is the Projection Metric (2-norm of the sine of principal angles): $$ d_P(A,B) = \left ( \sum_{i=1}^m sin^2 \theta_i \right )^{1/2} = \left ( m - \sum_{i=1}^m cos^2 \theta_i \right )^{1/2} = \left ( m - \sum_{i=1}^m \sigma_i^2 \right )^{1/2}. $$

Let’s return to the similarity metric $ \phi $. If the singular values of the product $ U_{A_{r=8}}^{iT} U_{A_{r=64}}^j $ are $ \sigma_1, \sigma_2, \dots, \sigma_p $, where $ p = min(i,j) $, then the similarity is defined as:

$$ \begin{split} \phi(A_{r=8}, A_{r=64}, i, j) = \frac{ \| U_{A_{r=8}}^{iT} U_{A_{r=64}}^{j} \|_F^2 }{ min(i,j) } = \\ = \frac{ \sum_{i=1}^{p} \sigma_i^2 }{p} = \frac{1}{p} \left ( 1 - d_P( U_{A_{r=8}}^{iT}, U_{A_{r=64}}^{j} ) \right ), \end{split} $$ which is exactly a reverse of the Projection Metric defined earlier.

Using this similarity metric and having carried out multiple experiments the authors concluded

that the top singular-vector directions of $ A_{r=8} $ and $ A_{r=64} $ are the most useful, while other directions potentially contain mostly random noises accumulated during training. Hence, the adaptation matrix can indeed have a very low rank ¹.

Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. 2021. ↩︎ ↩︎
Ham and Daniel D. Lee. Grassmann discriminant analysis: a unifying view on subspace-based learning. In ICML, pp. 376–383, 2008. ↩︎