RotoGrad - Gradient Homogenization in Multitask Learning

Authors: Adrian Javaloy & Isabel Valera from Saarland University

Submitted to ICLR2022, score: 8, 8, 8, 8.

Notations:

$\mathrm{\nabla }$$\nabla$ denotes the gradient operator.

$\underset{―}{x}$$\underline{x}$ underline to denote vectors.

$X$$X$ capital letters denote matrices.

$<,>$$<,>$ dot product operator.

$SO(d)=\{R\in {\mathbb{R}}^{d\times d}|R^{T}R=I,\det \left(R\right)=+1\}$$SO(d)=\{R \in \mathbb{R}^{d\times d} \space|\space R^TR=I, \det(R)=+1\}$ the special orthogonal group.

$\wedge :{\mathbb{R}}^d\to so(d)$$\wedge : \mathbb{R}^{d} \rightarrow so(d)$, i.e., $\underset{―}{x}\to \hat{X}$$\underline{x} \rightarrow \hat{X}$. $\wedge$$\wedge$ denoted as the hat oprator.

$so(d)=\{\hat{X}\in {\mathbb{R}}^{d\times d}|\underset{―}{x}\in {\mathbb{R}}^d\}$$so(d)=\{ \hat{X} \in \mathbb{R}^{d\times d} \space | \space \underline{x}\in\mathbb{R}^d \}$ the associated lie algebra to $SO(d)$$SO(d)$ constitute a set of skew-symmetric matrices.

$\vee$$\vee$ the inverse of the hat operator denoted as the vee operator.

1. What did the authors try to accomplish?

Rotograd tackles the problem of negative transfer in multi-task learning (MTL). In particular, negative transfer in MTL is caused by:

1. Varying gradient magnitude across tasks, e.g., one task with a relatively higher gradient magnitude will dominate the gradient's direction resulting in poor performance of other tasks.
2. Varying gradient direction across tasks, e.g., gradients from different tasks, may cancel out each other, resulting in slow task learning.

1.1 Main Claims

On the one hand, they propose to guide the step size with the least converged task. On the other hand, a task-specific rotation matrix $R_k:{\mathbb{R}}^d\to {\mathbb{R}}^d$$R_k: \mathbb{R}^{d} \rightarrow \mathbb{R}^{d}$ rotates the feature space to align the gradients according to a given direction.

1.2 Contributions

The authors claim that previous work in the literature primarily focused on the first problem. Instead, they propose to tackle both problems. Specifically, they propose a novel way to tackle gradient direction conflict by rotating the feature space such that the gradients align towards a common direction. To the best of my knowledge, they are the first to use rotation transformation in the Lie group $SO(d)$$SO(d)$ to multi-task learning.

2. What were the key element of the approach?

2.1 Idea

Given an input $X$$X$, let $f_{\theta }(X)=\underset{―}{z}$$f_\theta(X) = \underline{z}$ be the shared backbone for all tasks and $h_{{\varphi }_k}$$h_{\phi_k}$ the task-specific head. Moreover, given the linear combination of task gradients ${\mathrm{\nabla }}_{\theta }L=\sum _k{w}_k{\mathrm{\nabla }}_{\theta }{L}_k$$\nabla_{\theta}L= \sum_k w_k \nabla_{\theta}L_k$, we have ${\mathrm{\nabla }}_{\theta }{L}_k={\mathrm{\nabla }}_{\theta }f{\mathrm{\nabla }}_{\underset{―}{z}}{L}_k$$\nabla_{\theta}L_k = \nabla_{\theta}f \space \nabla_{\underline{z}}L_k$ by the chain rule. Since the tasks only compete for ressource on $\underset{―}{z}$$\underline{z}$, we can ignore ${\mathrm{\nabla }}_{\theta }f$$\nabla_{\theta} f$ and focus on ${\mathrm{\nabla }}_{\underset{―}{z}}{L}_k$$\nabla_{\underline{z}}L_k$, where $L_k$$L_k$ is the task-specific loss. Furthermore let $h_k$$h_k$ be the head for each task.

2.1.1 Gradient Magnitude

Goal change the magnitude of each $\Vert G_k\Vert$$\norm{G_k}$.

Let ${\underset{―}{g}}_{n,k}={\mathrm{\nabla }}_{\underset{―}{z}}{L}_k(h_k(X),Y_{n,k})$$\underline{g}_{n,k} = \nabla_{\underline{z}} L_k(h_k(X), Y_{n,k})$ be the gradient for the $k$$k$-th task and the $n$$n$-th data point. Then, its batch version is $G_k^T=[{\underset{―}{g}}_{1,k},...,{\underset{―}{g}}_{b,k}]$$G^{T}_{k} = [\underline{g}_{1,k}, ..., \underline{g}_{b,k}]$.

The authors propose to normalize each gradient $G_k$$G_k$ to a unit vector. Then they rescale it such that the tasks that have converged the least set up the step size. We set step size to be a weighted combination of task-wise gradient combination, where the weights sum up to one. Let $a_k$$a_k$ be this weight for each task, then

A side effect of this is that slow converging tasks will force quick converging tasks to escape from saddle points.

2.1.2 Gradient Direction

Let $R_k$$R_k$ be the task-specific rotation matrix, then instead of optimizing $L_k(\underset{―}{z})$$L_k(\underline{z})$, we optimize $L_k({\underset{―}{r}}_k)$$L_k(\underline{r}_k)$ where ${\underset{―}{r}}_k=R_k\underset{―}{z}$$\underline{r}_k = R_k\underline{z}$. By rotating the feature space we can ensure that gradients direction from different tasks do not conflict. Instead, they align to the same direction given by ${\underset{―}{v}}_n$$\underline{v}_n$. In particular, $R_k$$R_k$ has to minimize the following objective:

where ${\underset{―}{v}}_n=\frac{1}{K}\sum _k{\underset{―}{u}}_n,k$$\underline{v}_n = \frac{1}{K} \sum_k \underline{u}_n,k$ is a sum of normalized gradient direction from the different tasks, and $g_{n,k}^{{}^{\prime }}={\mathrm{\nabla }}_{{\underset{―}{r}}_k}{L}_k(h_k,Y_{n,k})$$g^{'}_{n,k}=\nabla_{\underline{r}_k} L_k(h_k, Y_{n,k})$ is the gradient flowing down from the head before being rotated.

Learning a task-specific rotation matrix requires considering a constrained minimization problem for $R_k:{\mathbb{R}}^d\to {\mathbb{R}}^d$$R_k: \mathbb{R}^{d} \rightarrow \mathbb{R}^{d}$. In practice, $\underset{―}{z}$$\underline{z}$ is of high dimension; hence optimizing this objective is unfeasible as it would require computing the determinant of $R_k$$R_k$, i.e., we want $\det R_k=1$$\det R_k = 1$. Hence, to turn it into an unconstrained minimization problem, we consider $R_k\in SO(d)$$R_k \in SO(d)$ that allows us to work with rotational matrices in a vector space $so(d)$$so(d)$. A vector space of matrices is much nicer to work with since it is closed under vector summation and scalar multiplication. We do not need to keep the constraint $\det R_k=1$$\det R_k =1$ and use regular gradient descent to optimize $R_k$$R_k$.

One might think that we can directly optimize the free parameters of the rotation matrices, however, constructing such a rotation matrix in high dimension is unfeasible.

Lie Group and Lie Algebra

On a high level, the set of all $d$$d$-dimensional rotational matrices (with determinant 1) constitute a group as it respects the four axioms (i.e., identity, associative, inverse, and closure) under matrix multiplication. In addition, if these rotational matrices are differential, it is also a Lie group. For the special orthogonal matrices group denoted as $SO(d)$$SO(d)$, one can smoothly rotate a matrix into another, hence it is a connected Lie group. This group is "special" because these matrices' determinant are 1.

We consider $R_k\in SO(d)$$R_k \in SO(d)$ because every Lie group has an associated Lie algebra. In particular, the Lie algebra denoted $so(d)$$so(d)$ is the set of $d\times d$$d \times d$ skew-symmetric matrices that form a vector space. These skew-symmetric matrices can be mapped back to their corresponding element on the Lie group via exponential maps $\exp :so(d)\to SO(d)$$\exp: so(d) \rightarrow SO(d)$. And this new space $so(d)$$so(d)$ allows us to work with a local approximation of the rotation matrices $SO(d)$$SO(d)$.

More formally, consider a set rotation matrices $R(t):t\to SO(d)$$R(t): t \rightarrow SO(d)$, which continuously transform a point from its original location ($R(0)=I$$R(0)=I$) to a different one:

since $R(t)R(t{)}^T=I$$R(t)R(t)^T =I$, $\mathrm{\forall }t$$\forall t$, we have

then,

We know that, skew-symmetric matrices or anti-symmetric matrices have the two following properties $M=-M^T$$M = -M^{T}$ and $M_{ii}=-M_{ii}$$M_{ii}=-M_{ii}$. The diagonal must be 0 as $R_{ii}=-R_{ii}$$R_{ii}=-R_{ii}$ and $R_{ii}=R_{ii}^T$$R_{ii} = R^{T}_{ii}$.

Let $\hat{W}\in so(d)$$\hat{W}\in so(d)$ be the skew-symmetric matrix of a vector $\underset{―}{w}\in {\mathbb{R}}^{\mathbb{d}}$$\underline{w} \in \mathbb{R^d}$. Then the above results suggest that there exists a vector $\underset{―}{w}\in {\mathbb{R}}^d$$\underline{w} \in \mathbb{R}^{d}$ such that:

and since $R(0)=I$$R(0)=I$, it follows that $\frac{dR(0)}{dt}=\hat{W}(0)$$\frac{dR(0)}{dt} = \hat{W}(0)$. This is a simple ordinary differential equation, for simplicity assume $\hat{W}(t)$$\hat{W}(t)$ constant in $t$$t$, then the solution is the matrix exponential:

yielding the exponential maps $\exp :so(3)\to SO(3)$$\exp: so(3) \rightarrow SO(3)$; $\hat{W}\to e^{\hat{W}}$$\hat{W} \rightarrow e^{\hat{W}}$. In fact, every rotation matrix $R\in SO(d)$$R \in SO(d)$ have infinitely many exponential coordinates $\underset{―}{w}$$\underline{w}$ such that $e^{\hat{W}}=R$$e^{\hat{W}} = R$. The exponential coordinates provide a local parameterization for rotation matrices.

As a sanity check, since $\hat{W}$$\hat{W}$ is skew-symmetric, it yields that $-\hat{W}={\hat{W}}^T$$-\hat{W} = \hat{W}^T$, hence:

Furthermore, we observe that the skew-symmetric $\hat{W}\in so(d)$$\hat{W} \in so(d)$ gives the first-order approximation of a rotation at $t=0$$t=0$:

Hence, the Lie algebra uses the tangent space to the Lie group at the identity element $I_d$$I_d$. We also observe that $\det (I_d+\hat{W}(0)dt)=1$$\det{(I_d + \hat{W}(0)dt)} =1$.

To conclude, we showed that:

Note: The authors have tried using a learned non-linear transformation instead of a rotation matrix. However, such choice results in numerical issues related to scaling of the feature $\underset{―}{z}$$\underline{z}$ in the forward pass, i.e., it affects the effective learning rate of different heads.

2.2 Limitations

The method scales linearly with the number of tasks, i.e., $R_k$$R_k$ induces $K{d}^2$$Kd^2$ additional parameters. And the time complexity induced is $K{d}^3$$Kd^3$.

3. Reference

RotoGrad: Gradient Homogenization in Multitask Learning

An invitation to 3-D Vision, FromImages to Geometric Models. Yi Ma, Stefano Soatto, Jana Kosecka, S. Shankar Sastry.