Quaternion Product Units for Deep Learning on 3D Rotation Groups

Quaternion $q=s+{\mathbf{i}}x+{\mathbf{j}}y+{\mathbf{k}}z\in\mathbb{H}$ is a type of hypercomplex number with 1D real part $s$ and 3D imaginary part $(x,y,z)$. The imaginary parts satisfy ${\mathbf{i}}^{2}={\mathbf{j}}^{2}={\mathbf{k}}^{2}={\mathbf{i}}{\mathbf{j}}{\mathbf{k}}=-1,{\mathbf{i}}{\mathbf{j}}=-{\mathbf{j}}{\mathbf{i}}={\mathbf{k}},{\mathbf{k}}{\mathbf{i}}=-{\mathbf{i}}{\mathbf{k}}={\mathbf{j}},{\mathbf{j}}{\mathbf{k}}=-{\mathbf{k}}{\mathbf{j}}={\mathbf{i}}$. For convenience, we ignore the imaginary symbols and represent a quaternion as the combination of a scalar and a vector as $q=[s,{\bm{v}}]=[s,(x,y,z)]$. When equipped with Hamilton product as multiplication, noted as $\otimes$, and standard vector addition, the quaternion forms an algebra. In particular, the Hamilton product between two quaternions $q_{1}=[s_{1},{\bm{v}}_{1}]$ and $q_{2}=[s_{2},{\bm{v}}_{2}]$ is defined as

$\displaystyle\begin{aligned} q_{1}\otimes q_{2}&=[s_{1}s_{2}-\langle{\bm{v}}_{1},{\bm{v}}_{2}\rangle,{\bm{v}}_{1}\times{\bm{v}}_{2}+s_{1}{\bm{v}}_{2}+s_{2}{\bm{v}}_{1}]\\ &=M_{L}(q_{1})q_{2}=M_{R}(q_{2})q_{1},\end{aligned}$

(2)

where for a quaternion $q=[s,(x,y,z)]$, we have

$\displaystyle M_{L}(q)=\left[\begin{smallmatrix}s&-x&-y&-z\\ x&s&-z&y\\ y&z&s&-x\\ z&-y&x&s\end{smallmatrix}\right],M_{R}(q)=\left[\begin{smallmatrix}s&-x&-y&-z\\ x&s&z&-y\\ y&-z&s&x\\ z&y&-x&s\end{smallmatrix}\right].$

(3)

The second row in Eq. (2) uses matrix-vector multiplications, in which we treat one quaternion as a 4D vector and compute a matrix based on the other quaternion. Note that the Hamilton product is non-commutative — multiplying a quaternion on the left or the right gives different matrices.

Quaternion algebra provides us with a new representation of 3D rotations. Specifically, suppose that we rotate a 3D vector ${\bm{v}}_{1}$ to another 3D vector ${\bm{v}}_{2}$, and the 3D rotation is with an axis ${\bm{u}}$ and an angle $\theta$ shown in Eq. (1). We can represent the 3D rotation as a unit quaternion $q=[s,{\bm{v}}]=[\cos(\frac{\theta}{2}),\sin(\frac{\theta}{2}){\bm{u}}]$, where $\|{\bm{u}}\|_{2}=1$ and $s^{2}+\|{\bm{v}}\|_{2}^{2}=\cos^{2}(\frac{\theta}{2})+\sin^{2}(\frac{\theta}{2})=1$. Representing the two 3D vectors as two pure quaternions, $i.e.$, $[0,{\bm{v}}_{1}]$ and $[0,{\bm{v}}_{2}]$, we can achieve the rotation from ${\bm{v}}_{1}$ to ${\bm{v}}_{2}$ via the Hamilton products of the corresponding quaternions:

$$[0,{\bm{v}}_{2}]=q\otimes[0,{\bm{v}}_{1}]\otimes q^{*},$$

(4)

where $q^{*}=[s,-{\bm{v}}]$ stands for the conjugation of $q$. Additionally, the combination of rotation matrices can also be represented by the Hamilton products of unit quaternions. For example, given two unit quaternions $q_{1}$ and $q_{2}$, which correspond to two rotations, $(q_{2}\otimes q_{1})\otimes[0,{\bm{v}}_{1}]\otimes(q_{1}^{*}\otimes q_{2}^{*})$ means rotating ${\bm{v}}_{1}$ sequentially through the two rotations.

Note that unit quaternion is a double cover of $\mathbb{SO}(3)$ since $q$ and $-q$ represents the same 3D rotation (in opposite directions). As shown in Eq. (1), we use inner product and $\arccos$ to generate unit quaternion and reduce this ambiguity by choosing the quaternion with positive real part. Refer to [5] for more details on unit quaternion and 3D rotation.

In standard deep learning models, each of their neurons can be represented as a weighted summation unit, $i.e.$, $y=\sigma(\sum_{i=1}^{N}w_{i}x_{i}+b)$, where $\{w_{i}\}_{i=1}^{N}$ and $b$ are learnable parameters, $\{x_{i}\}_{i=1}^{N}$ are inputs, and $\sigma(\cdot)$ is a nonlinear activation function. As aforementioned, when the input $x_{i}\in\mathbb{SO}(3)$, such a unit cannot keep the output $y$ on $\mathbb{SO}(3)$ as well. To design a computational unit guaranteeing the closure of $\mathbb{SO}(3)$, we propose an alternate of this unit based on the quaternion algebra introduced above.

Specifically, given $N$ unit quaternions $\{q_{i}=[s_{i},{\bm{v}}_{i}]\}_{i=1}^{N}$ that represent 3D rotations, we can define a weighted chain of Hamilton products as

$$y=\sideset{}{{}_{i=1}^{N}}{\bigotimes}q_{i}^{w_{i}}=q_{1}^{w_{1}}\otimes q_{2}^{w_{2}}\otimes...\otimes q_{N}^{w_{N}},$$

(5)

where the power of a quaternion $q=[s,{\bm{v}}]$ with a scalar $w$ is defined as [5]:

$\displaystyle\begin{aligned} q^{w}:=\begin{cases}[\cos(w\arccos(s)),\frac{{\bm{v}}}{\|{\bm{v}}\|_{2}}\sin(w\arccos(s))],\quad{\bm{v}}\neq\bm{0},\\ [1,\bm{0}],\quad\text{otherwise}.\end{cases}\end{aligned}$

Note that the power of a quaternion only scales the rotation angle and does not change the rotation axis.

Here, we replace the weighted summation with a weighted chain of Hamilton products, which makes $\mathbb{SO}(3)$ closed under the operation. Based on the operation defined in Eq. (5), we proposed our quaternion product unit:

$\displaystyle\begin{aligned} \text{QPU}(\{q_{i}\}_{i=1}^{N};\{w_{i}\}_{i=1}^{N},b)=\sideset{}{{}_{i=1}^{N}}{\bigotimes}\text{qpow}(q_{i};w_{i},b),\end{aligned}$

(6)

where for $q_{i}=[s_{i},{\bm{v}}_{i}]$

$\displaystyle\begin{aligned} &\text{qpow}(q_{i};w_{i},b)\\ =&[\cos(w_{i}(\arccos(s_{i})+b)),\frac{{\bm{v}}_{i}}{\|{\bm{v}}_{i}\|_{2}}\sin(w_{i}(\arccos(s_{i})+b))]\end{aligned}$

represents the weighting function on the rotation angles with a weight $w_{i}$ and a bias $b$.

Compared with Eq. (5), we add a bias to $b$ to shift the origin. The output of $\arccos$ contains an infinite gradient at $\pm 1$. We solve this problem by clamping the input scalar part $s$ between $-1+\epsilon$ and $1-\epsilon$, where $\epsilon$ is a small number.

The following proposition demonstrates the advantage of our QPU on achieving rotation robustness.

The proof above indicates that the intrinsic property and the group law of $\mathbb{SO}(3)$ naturally provided rotation-invariance and rotation-equivariance. Without forced tricks or hand-crafted representation strategies, the principled design of QPU makes it flexible for a wide range of applications with different requirements on rotation robustness. Figure 2 further visualizes the property of our QPU. Given a QPU, we feed it with a human skeleton and its rotated version, respectively. We compare their outputs and find that the imaginary parts of their outputs inherit the rotation discrepancy between them while the real parts of their outputs are the same with each other.

The forward pass of QPU involves $N$ individual weighting functions and a chain of Hamilton products. While it is easy to compute the gradients of the weighting functions [5], the gradient of the chain of Hamilton products is not straightforward.

Given quaternions $\{q_{i}\}_{i=1}^{N}$, we first compute the differential of the $k$-th quaternion $q_{k}$ with respect to chain of Hamilton product $\bigotimes_{i=1}^{N}q_{i}$, denoted as $d(\bigotimes_{i}q_{i})_{q_{k}}(\cdot)$, omitting $N$ for simplicity. As Hamilton product is bilinear, we have

$\displaystyle\begin{aligned} d(\sideset{}{{}_{i}}{\bigotimes}q_{i})_{q_{k}}(h)&=(\sideset{}{{}_{j<k}}{\bigotimes}q_{j})\otimes h\otimes(\sideset{}{{}_{j>k}}{\bigotimes}q_{j})\\ &=M_{L}(\sideset{}{{}_{j<k}}{\bigotimes}q_{j})M_{R}({\sideset{}{{}_{j>k}}{\bigotimes}q_{j}})h\\ &=M_{L}(B_{k})M_{R}(A_{k})h.\end{aligned}$

(9)

where $h$ represents a small variation of the $k$-th input quaternion, $B_{k}=\bigotimes_{j<k}q_{j}$ and $A_{k}=\bigotimes_{j>k}q_{j}$ are the chains of Hamilton product of the quaternions before and after the $k$-th quaternion, respectively.

We then compute the differential of loss $L$ for the $k$-th quaternion. Given the gradient of the scalar loss $L$ with respect to the output quaternion $\frac{\partial L}{\partial(\bigotimes_{i}q_{i})}$ , e.g., computed with autograd [15] of Pytorch, we have

$\displaystyle\begin{aligned} dL_{q_{k}}(h)&=\langle\frac{\partial L}{\partial(\bigotimes_{i}q_{i})},d(\sideset{}{{}_{i}}{\bigotimes}q_{i})_{q_{k}}(h)\rangle\\ &=\langle\frac{\partial L}{\partial(\bigotimes_{i}q_{i})},M_{L}(B_{k})M_{R}(A_{k})h\rangle\\ &=\langle M_{R}^{T}(A_{k})M_{L}^{T}(B_{k})\frac{\partial L}{\partial(\bigotimes_{i}q_{i})},h\rangle,\end{aligned}$

(10)

where $\langle\cdot,\cdot\rangle$ stands for the inner product between vectors and $M^{T}$ stands for the transpose of matrix $M$. Thus the gradient of loss $L$ for $q_{k}$ is given as

$$\frac{\partial L}{\partial q_{k}}=M_{R}^{T}(A_{k})M_{L}^{T}(B_{k})\frac{\partial L}{\partial(\bigotimes_{i}q_{i})}.$$

(11)

To compute $B_{k}$ and $A_{k}$, we first compute the cumulative Hamilton product $C=(c_{1},c_{2},...,c_{N})$, where $c_{k}=\bigotimes_{j\leq k}q_{j}$. Then $B_{k}=c_{k}\otimes q_{k}^{*}$ and $A_{k}=c_{k}^{*}\otimes c_{N}$. We note that the gradient depends on a cumulative Hamilton product of all quaternions other than the $k$-th quaternion. This computation process is different from that of standard weighted summation units, in which the backpropagation through addition only involves the differential variable instead of other inputs. In contrast, it is interesting to see the gradient is a joint result of all other inputs to QPU except the items of the given differential variable.

As shown in Figure 1(b), multiple QPUs receiving the same input quaternions form a fully-connected (FC) layer. The parameters of the QPU-based FC layer include the weights and the bias used in the QPUs. We initialize them using the Xavier uniform initialization [8]. Different from real-valued weighted summation, the QPU itself is nonlinear. Hence there is no need to introduce additional nonlinear activation functions. As aforementioned, the output of QPU is still a unit quaternion so that we can connect multiple QPU-based FC layers sequentially. Accordingly, the stack of multiple QPU-based FC layers establishes a quaternion multi-layer perceptron (QMLP) model for 3D rotation data. Note that the Hamilton product is not commutative, $i.e.$, $q_{1}\otimes q_{2}\neq q_{2}\otimes q_{1}$ in general, the stacking of two QPU-based FC layers is not equivalent to one QPU-based FC layer.

Besides the QPU-based FC layer and the corresponding QMLP model, our QPU can also be used to implement a quaternion-based graph convolution layer. Specifically, given a graph with an adjacency matrix $A=[a_{ij}]\in\mathbb{R}^{N\times N}$, the aggregation of its node embeddings is achieved by the multiplication between its adjacency matrix and the embeddings [11]. When the node embeddings correspond to 3D rotations and are represented as unit quaternions, we can implement a new aggregation layer using the chain of Hamilton products shown in Eq. (5), $i.e.$, given input quaternions $\{q_{i}\}_{i=1}^{N}$ and the adjacency matrix $A$, the $i$-th output $q_{i}^{agg}=\bigotimes_{j=1}^{N}q_{j}^{a_{ij}}$. Stacking such a QPU-based aggregation layer with a QPU-based FC layer, we achieve a quaternion-based graph convolution layer accordingly.

Besides building pure quaternion-valued models ($i.e.$, the QMLP mentioned above), we can plug our QPU-based layers into existing real-valued models easily. Suppose that we have one QPU-based layer, which takes $N$ unit quaternions as inputs and derives $M$ unit quaternions accordingly. When the QPU-based layer receives rotations from a real-valued model, we merely need to reformulate the inputs as unit quaternions. When a real-valued layer follows the QPU-based layer, we can treat the output of the QPU-based layer as a real-valued matrix with size $M\times 4$ and directly feed the output to the subsequent real-valued layer. Moreover, when rotation-invariant (rotation-equivariant) features are particularly required, we can feed only the real part (the imaginary part) to the real-valued layer accordingly.

Besides the straightforward strategy above, we further propose an angle-axis map to connect QPU-based layers with real-valued ones. Specifically, given a unit quaternion $q=[s,{\bm{v}}]$, the proposed mapping function is defined as:

$$\text{AngleAxisMap}(q)=[\arccos(s),\frac{{\bm{v}}}{\|{\bm{v}}\|_{2}}].$$

(12)

We use this function to map the output of a QPU, which lies on a non-Euclidean manifold [10], to the Euclidean space. Using this function to connect QPU-based layers with real-valued ones makes the learning of the downstream real-valued models efficient, which helps us improve learning results. Our angle-axis mapping function is better than the logarithmic map used in [28, 27, 9, 24]. The logarithmic map $\log(q)=[0,\arccos(s)\frac{{\bm{v}}}{\|{\bm{v}}\|_{2}}]$ mixes the angular and axial information, while our map keeps this disentanglement, which is essential to preserve the rotation-invariance and rotation-equivariance of the output.

As shown in Section 2.4, our QPU supports gradient-based update, so both QPU-based layers and the models combining QPU-based layers with real-valued ones can be trained using backpropagation. Since each QPU handles a quaternion as one single element, QPU is efficient in terms of parameter numbers. A real-valued FC with $N$ inputs and $M$ outputs requires $NM$ parameters, while a QPU with the same input and output dimensions requires only $\frac{1}{16}NM$ parameters.

Due to the usage of Hamilton products, however, QPU requires more computational time than the real-valued weighted summation unit. Take a QPU-based FC layer with $N$ input quaternions and $M$ output quaternions as an example. In the forward pass, the weighting step requires $NM$ multiplications and $NM$ sine and cosine computations. Each Hamilton product requires $16$ multiplications and $12$ additions. So the chain of $N-1$ Hamilton products requires $16(N-1)M$ multiplications and $12(N-1)M$ additions. In total, a QPU-based FC layer requires $(17N-16)M$ multiplications, $(13N-12)M$ additions and $NM$ cosine and sine computations. As a comparison, the real-valued FC layer with the same input and output size ($i.e.$, $4N$ inputs and $4M$ outputs in terms of real numbers) requires $16NM$ multiplications, $16(N-1)M$ additions.

In the backward pass, the gradient computation of a chain of Hamilton products requires the cumulative Hamilton product of the power weighted input quaternions. In particular, when computing the gradient of a chain of Hamilton products for an input quaternion, we need to do two Hamilton products and one matrix-matrix multiplication. As a result, the computational time is at least doubled compared with the forward pass and tripled if we recompute the cumulative Hamilton product. To deal with this challenge, we have two strategies: (a) store the result of cumulative Hamilton product in the forward pass, (b) recompute the same quantity during the backward pass. The strategy (a) saves computational time but requires to save a potentially large feature map ($i.e.$, $N$ times the size of the feature map). The strategy (b) requires more computation time but no additional memory space. We tested both options in our experiments. When the QPU-based layer is with small feature maps ($i.e.$, at the beginning of a neural network with fewer input channels), both these two strategies work well. When we stack multiple QPU-based layers and apply the model to large feature maps, we need powerful GPUs with large memory spaces to implement the first strategy.

According to the above analysis, the computational bottleneck of our QPU is the chain of Hamilton products. Fortunately, we can accelerate this step by breaking down the computation like a tree. Specifically, as Hamilton product satisfies the combination law, $i.e.$, $(q_{1}\otimes q_{2})\otimes q_{3}=q_{1}\otimes(q_{2}\otimes q_{3})$, we can multiply quaternions at odd positions with quaternions at even positions in a parallel way and compute the whole chain by repeating this step recursively. For a chain of $N-1$ Hamilton products, this parallel strategy reduces the time complexity from $\mathcal{O}(N)$ to $\mathcal{O}(\log(N))$. We implemented this strategy in the forward pass and witnessed a significant acceleration. In summary, although the backpropagation of our QPU-based model requires much more computational resources ($i.e.$, time, or memory) than that of real-valued models, its forward pass merely requires slightly more computations than that of real-valued models and can be accelerated efficiently via a simple parallel strategy.

We first design a synthetic dataset called “CubeEdge” and propose a simple task to demonstrate the rotation robustness of our QPU-based model. The dataset consists of skeletons composed of consecutive edges taken from a 3D cube (a chain of edges). These skeletons are categorized into classes according to their shapes, $i.e.$, the topology of consecutive edges. For each skeleton, we use the rotations between consecutive edges as its feature.

For each skeleton in the dataset, we generate it via the following three steps, as shown in Figure 3(a): $i$) We initialize the first two adjacent edges deterministically. $ii$) We ensure that the consecutive edges cannot be the same edge, so given the ending vertex of the previous edge, we generate the next edge randomly along either of the two valid directions. Accordingly, each added edge doubles the number of possible shapes. We repeat this step several times to generate a skeleton. $iii$) We apply random shear transformations to the skeletons and augment the dataset, and split the dataset for training and testing. We generate each skeleton with seven edges so that our dataset consists of 32 different shapes (classes) of skeletons. We use 2,000 samples for training and testing, respectively. For the skeletons in the testing set, we add Gaussian noise $\mathcal{N}(0,\sigma^{2})$ to their vertices before shearing, where the standard deviation $\sigma$ controls the level of noise. Figure 3(b) gives four testing samples. Additional random rotation is applied to testing set when validating the robustness to rotations.

To classify the synthetic skeletons, we apply our QMLP model and its variant QMLP-RInv and compared them with a standard real-valued MLP (RMLP). For fairness, all three models are implemented with three layers, and the size of output feature maps are the same. The RMLP is composed of three fully-connected layers connected by two ReLU activation functions. Our QMLP is composed of two stacked QPU layers, followed by a fully-connected layer. Our QMLP-RInv is composed of two stacked QPU layers with only the real part of the output retained, and a fully-connected layer is followed. The architectures of these three models are shown in Figure 4.

After training these three models, we test them on the following two scenarios: $i$) comparing them on the testing set directly; and $ii$) adding random rotations to the testing skeletons and then comparing the models on the randomly-rotated testing set. The first scenario aims at evaluating the feature extraction ability of different models, while the second one aims at validating their rotation-invariance. For each scenario, we compare these three models on their classification accuracy. Their results with respect to different levels of noise ($i.e.$, $\sigma$) are shown in Figure 5. In both scenarios, our QMLP and QMLP-RInv outperform the RMLP consistently, which demonstrates the superiority of our QPU model. Furthermore, we find that the QMLP-RInv is robust to the random rotations imposed on the testing data — it achieves nearly the same accuracy in both scenarios and works better than QMLP in the scenario with random rotations. This phenomenon verifies our claims: $i$) the real part of our QPU’s output is rotation-invariant indeed; $ii$) the disentangled representation of rotation-invariant feature and rotation-equivariant feature could inherently make the network robust to rotations.

Besides testing on synthetic data, we also consider two real-world skeleton datasets: the NTU dataset [20] for human action recognition and the FPHA dataset [7] for hand action recognition. The NTU dataset provides 56,578 human skeleton sequences performed by 40 actors belonging to 60 classes. The cross-view protocol is used in our experiments. We use 37,646 sequences for training and the remaining 18,932 sequences for testing. FPHA dataset provides hand skeletons recorded with magnetic sensors and computed by inverse kinematics, which consists of 1,175 action videos belonging to 45 categories and performed by 6 actors. Following the setting in [7], we use 600 sequences for training and 575 sequences for testing.

For these two datasets, we implement three real-valued models as our baselines.

RMLP-LSTM is composed of a two-layer MLP and a one-layer LSTM. The MLP merges each frame into a feature vector and is shared between frames. Features of each frame are then fed into the LSTM layer. We average the outputs of the LSTM at different steps and feed it into a classifier.

AGC-LSTM [23] first uses a FC layer on each joint to increase the feature dimension. The features are fed into a 3-layer graph convolutional network.

DGNN [21] is composed of 10 directed graph network (DGN) blocks that merge both node features and edge features. Temporal information is extracted using temporal convolution between the DGN blocks.

To demonstrate the usefulness of our QPU, we substitute some of their layers with our QPU-based FC layers, and propose the following three QPU-based models:

QMLP-LSTM: The two FC layers of the RMLP-LSTM are replaced with two QPU-based FC layers.

QAGC-LSTM: The input FC layer of the AGC-LSTM replaced with a QPU-based FC layer. In the original AGC-LSTM, the first FC layer only receives a feature from one joint. As a result, the input would be a single quaternion, and the QPU layer would be unable to capture interactions between inputs. We thus augment the input by stacking the feature of the parent joint so that the input channel to QPU is two quaternions.

QDGNN: The FC layers in the first DGN block of QDGNN are replaced with QPU-based FC layers. The original DGNN uses 3D coordinates of joints as node features and vectors from parent joint to child joint as edge features. In our QDGNN, the node features are the unit quaternions corresponding to the joint rotations and the edge features are the differences between joint rotations, $i.e.$, $q_{1}\otimes q_{2}^{*}$ computes the edge features for joint rotations $q_{1}$ and $q_{2}$. Moreover, the original DGNN aggregates node features and edge features through the normalized adjacent matrix. In our QDGNN, we implement this aggregation as the QPU-based aggregation layer proposed in Section 3.1.

Similar to the experiments on synthetic data, we consider for the three QPU-based models their variants that only keep the real part of each QPU’s output, denoted as QMLP-LSTM-RInv, QAGC-LSTM-RInv, and QDGNN-RInv, respectively. For these variants, we multiply the output channels of the last QPU-based FC layer by $4$ and keep only the real parts of its outputs. In all real-world experiments, we use the angle-axis map (Eq. (12)) to connect the QPU-based layers and real-valued neural networks. The number of parameters of the QPU-based models is almost the same with that of the corresponding real-valued models.

We compare our QPU-based models with corresponding baselines under two configurations: $i$) both training data and test data have no additional rotations (NR); $ii$) the training data are without additional rotations while the testing ones are with arbitrary rotations (AR). The comparison results on NTU and FPHA are shown in Table 1, respectively. We can find that in the first configuration, the performance of our QPU-based models is at least comparable to the baselines. In the second and more challenging setting, which requires rotation-invariance, our QPU-based models that use only outputs’ real parts as features retain high accuracy consistently in most situations, while the performance of the baselines degrades a lot. Additionally, we quantitatively analyze the impact of our angle-axis map on testing accuracy. In Table 2, the results on the FPHA dataset verify our claim in Section 3.2: for our QPU-based models, their performance boosts a lot when we connect their QPU-based layers with the following real-valued layers through the angle-axis map.