I'm looking to constrain one layer of my neural network to specifically find the best rotation of its input in order to satisfy an objective. (My end goal, where $R$ is the rotation layer, is of the form $R^T ~f_{objective} \left(Rz \right)$ ).

I am looking to train this (+ other components) via gradient descent. If $z\in\mathbb{R}^2$, then I can just say $R = \left[ \begin{matrix} \cos \theta & -\sin \theta\\ \sin \theta & \cos \theta \end{matrix} \right] $, and have $\theta$ be a learnable parameter.

However, I am lost on how to actually set this up for an $d$-dimensional space (where $d$>10). I've tried looking at resources on how to make a $d$-dimensional rotation matrix and it gets heavy into Linear Algebra and is way over my head. It feels like this should be easier than it seems, so I feel like I'm overlooking something (like maybe $R$ should just be a usual linear layer without any non-linear activations).

Anyone have any ideas? I appreciate you, in advance : )


2 Answers 2


Suppose you have a set of $d$ nodes with linear activation and weights $b_{ij}$ for input $j$ to node $i$. If you can impose the constraints $\sum_j b_{ij}^2=1$ and $\sum_{j} b_{ij}b_{kj}=0$ for $i\neq k$ then the mapping from input to output is multiplying by an orthonormal matrix.

The orthonormal matrices form two connected sets: the rotations (determinant =1) and the rotations with reflection (determinant=-1). If you have a learning rate that isn't too high, your transformation won't be able to jump between these components, so if you start your weights off at rotation they'll stay a rotation.

This assumes you want rotations around the origin. To get rotations around some other point needs non-zero intercept (bias) terms chosen to move that point to the origin, rotate, then move it back.

  • $\begingroup$ I was thinking something along that line! Thank you : ) The question is how can we enforce the orthonormal constraint on a weight matrix? Maybe this is just an implementation detail, and we can have a regulation loss which penalizes and W which is not orthonormal. What do you think? $\endgroup$
    – Sean K
    Sep 28, 2021 at 15:31
  • 1
    $\begingroup$ For future readers, QR decomposition can help with this (since Q is orthogonal) via having W be an unconstrained learnable matrix and solve W = QR, and then actually use Q as your orthonormal. pytorch.org/docs/1.9.0/generated/torch.linalg.qr.html $\endgroup$
    – Sean K
    Sep 28, 2021 at 16:19
  • $\begingroup$ @SeanK That would be one way. I was hoping there was a computationally cheaper way to do it for matrices that are known to be of the form $Q+\epsilon A$ for small $\epsilon$, like you'd get after an update. $\endgroup$ Sep 28, 2021 at 22:29

The most computationally efficient representation of a rotation is as a series of $n$ householder reflections. The concept is pretty simple: your rotation is parameterized by a series of $n$ arbitrary vectors $v_1\dots v_n$ each with $n$ elements. To apply a corresponding rotation to an input $x$ vector, iterate through each of your $n$ parameter vectors $v_i$ and reflect $x$ through that vector

$x \leftarrow x - 2v_i\frac{x \cdot v_i}{v_i \cdot v_i}$

Alternatively, you can construct your entire rotation matrix explicitely from $R = \overset{n}{\underset{i=1}{\prod}}\left(I - 2 \frac{v_i \otimes v_i}{v_i \cdot v_i}\right)$

See here for a publication on this. Autograd frameworks like pytorch, tensorflow, or jax will have no problem backpropagating through this transformation so you can optimize parameters $v$.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.