IMUOptimize: A Data-Driven Approach to Optimal IMU Placement for Human Pose Estimation with Transformer Architecture

In our study, we focus on mapping Inertial Measurement Unit (IMU) data, specifically accelerations and rotations, to full body poses. This approach is crucial in understanding and interpreting the movement dynamics captured by IMU sensors.

Interestingly, the converse of the process is quite easy to calculate. Using a formula derived from TransPose [13], we can calculate the acceleration data between two poses - i.e. we can easily find IMU data given poses. This flexibility allows for the utilization of completely synthetic data in our analysis. The formula is given by:

where $v$ represents an array of frames, and $a[i]$ is the acceleration at frame $i$.

Naturally, the rotation of each joint is equivalent to the rotation of the corresponding synthetic IMU.

The base dataset for our study is AMASS [8], an aggregation of numerous motion capture datasets. AMASS provides data in the SMPL format, a mesh model that demonstrates how a virtual "body" is contorted by a given set of joint rotations. The SMPL model is particularly significant as it contains a $24\times 6890$ matrix that transforms 24 joints into the 6980 vertices of the model. By running an argmax operation on each of the 24 weight sets, we can identify the vertex that most closely corresponds to each joint.

Once we have vertices, which are points on the SMPL mesh, and their corresponding joints, we can compute IMU data on a per-joint basis by applying the above formula to the motion of the vertices between frames.

In addressing this problem, it is crucial to understand that we are dealing with a time series issue where isolated pose frames are insufficient. To synthesize IMU data effectively, one must consider the sequence of frames and their context within the surrounding frames. This calls for a model that ingests a series of $X$ frames of IMU data and outputs $X$ corresponding poses.

Long Short-Term Memory (LSTM) networks are exceptionally suited for this challenge due to their ability to process sequences, retain information about previous frames, and apply this context to new ones. A prediction at any frame is a weighted consideration $f(h,x)$ that considers both the current frame $x$, along with the model’s “memory”, the hidden state $h$.

Below, we have provided our model architecture as a list. Here, $N$ is the number of sensors the model uses - it will either be 24 or 6.

1. 1.

   Linear Layer: Transform from shape $N*(3*3+3)$ to $1024$.

2. 2.

   LSTM Layer: Accept shape $1024$ then feed it through 2 recurrent layers, each with hidden size $1024$. Produce output of shape $2*1024$.

3. 3.

   Linear Layer: Transform from shape $2*1024$ to final output shape $24*(3*3)$.

RNN training was executed over 5 epochs, where each epoch took approximately 20 minutes to run on an RTX3090. Counterintuitively, note that we actually want to overfit the model - we do not implement any sort of regularization, since our desired result is actually the model conditioning heavily on a small set of features, as opposed to using all the features equally.

To figure out which features are most important, we need to leverage a model interpretation technique. Simpler models, like random forests and gradient-boosted trees, [5] are easier to interpret, but since we work with deep learning models, it can be a little harder to understand what goes into the network’s prediction.

We apply Captum, a powerful library for model interpretability in PyTorch. Captum offers various methods for attributing model outputs to input features, aiding in the understanding of the model’s decision-making process. Given that our model input is a tensor of size 288, representing 24 joints with 3-axis acceleration and 3x3 elements for rotations, we employ Captum’s feature masking capabilities to analyze the importance of individual IMUs.

Feature masking involves selectively hiding or masking certain features during model evaluation to observe their impact on predictions. Notably, since only a subset of methods in Captum supports feature masks, we focus on those that apply to our scenario. We tried Shapley value, Shapley value sampling, and feature ablation - we found that feature ablation was the most convenient to use and provided the best results.

Feature ablation is known as a perturbation approach: it systematically replaces given features or groups of features with a baseline, then evaluates how much this influences predictions. Of course, this means that it requires data to operate on - unlike static analysis methods that directly inspect the parameters of the model, feature ablation requires that the dataset be passed in as a parameter. Thus, we reserved a set of test data, separate from the training dataset, that we could use for model evaluation and interpretation.

Transformers, as introduced in [12], have recently brought about significant advancements in time series analysis. Although groundbreaking in the field of large language models, where transformers are arranged in an encoder/decoder architecture to generate coherent text, the fundamental principle of self attention underlying transformers is applicable to a wide variety of time series tasks.

RNNs like LSTMs store context over time in a “hidden state” - after consuming each element $x$ of the input sequence, this hidden state is updated by passing it through a function $f(x,h)\rightarrow(y,h)$, a function that accepts an input $x$ and the current hidden state $h$, then returns an output $y$ and a new hidden state. The output of the RNN is the last $y$ produced - which means that the information that is used to compute the final output of the RNN is limited to the second-to-last hidden state and the last element. Although in theory, RNNs can operate on infinitely long sequences, the actual amount of context that can be used to compute the output of a given sequence element is limited to the amount of information that can be encoded by the hidden state vector. This results in a temporal dependency: the context to compute the output of any sequence element must come from sequence elements close to it.

More damaging is the inherently sequential nature of RNNs. The computation of hidden state $h_{k}$ depends on $h_{k-1}$. This unfortunate fact means that the computation of output $y_{n}$ requires the computation of the $n-1$ hidden states before it, rendering parallelization impossible. We’d like a model that is capable of leveraging modern GPUs.

Transformers address both of these issues through a mechanism known as “self attention”. Rather than sequentially computing hidden states / sequence outputs, self attention applies 3 separate transformations to the input, turning it into “query”, “key”, and “value” vectors (matrices, in practice). The query and key vectors are multiplied and transformed in a pairwise fashion, allowing the computation of a matrix called the “attention weights”. Essentially, this matrix encodes how important each element of the input is to each element of the output. It’s somewhat analogous to computing all of the hidden states in a single matrix multiplication, all at once. It can then be used to transform the “value” vector into a final output sequence.

Since the attention weights are computed by running some transformation $f(a,b)$ on every single pair $(a,b)$, the context usable by each sequence element includes the entire sequence, with equal potential importance given to each element (the actual importance is the attention weight). This successfully breaks temporal dependency. Furthermore, since the attention weights are computed all at once in a single matrix multiplication, this breaks the sequential nature of RNNs and lends itself to massive parallelization.

Therefore, we developed a transformer-based architecture to replace the LSTM. Generative applications of the transformer usually leverage both an encoder and a decoder, wherein previous sequences are used as input to the decoder in order to generate new sequences. However, since each pose sequence in our application is independent, we actually do not need a decoder at all and our model can perform with just an encoder.

Below, we have provided our model architecture as a list. Here, $N$ is the number of sensors the model uses - it will either be 24 or 6.

1. 1.

   Linear Layer: Transform from shape $N*(3*3+3)$ to $512$.

2. 2.

   Positional Encoding: Since transformers have no notion of position, augment each input with information about where it occurs in the sequence as per [12]. Transforms from $512$ to $512$.

3. 3.

   Transformer Encoder: Stack of 6 identical layers, each accepting and returning $512$. Each layer implements multi-headed self attention with 4 heads, then runs the concatenated result through a feedforward network to transform it back to $512$.

4. 4.

   Linear Layer: Transform from $512$ back to $24*(3*3)$.

We trained this for 5 epochs on an RTX3090, where each epoch took around 4 minutes. The transformer architecture yields performance around 5 times as fast as the LSTM.

As a result of our data-driven IMU placement strategy, both our LSTM and transformer-powered models achieve better performance on the TotalCapture [11] dataset than DIP-IMU, the previous work in this domain [7]. See table 2 for DIP-IMU’s results on TotalCapture - the $\mu_{ang}$ captures local rotation errors for each type of model. The important value is $15.85$, which is the best local rotation error that DIP-IMU achieved on TotalCapture.

To evaluate baseline performance - i.e. the best possible rotational error that our models could achieve - we ran evaluations on the models trained on 24 IMUs. Table 4 captures the evaluations for the LSTM, while table 3 captures the evaluations for the transformer. Therefore, the ideal rotational error we hope to achieve would be 13.892 degrees with an LSTM, and 8.695 degrees with a transformer.

That said, tables 5 and 6 reveal the performance of the optimized LSTM and transformer models, respectively. Notice that the optimized LSTM has a local rotation error of 13.018 on TotalCapture, while the optimized transformer has a local rotation error of 12.916. Both models yield significantly better performance than TotalCapture’s solution, while the transformer, as expected, performs slightly better than the LSTM.

Do remember, however, that even though the accuracy delta between the LSTM and transformer isn’t that large, the transformer’s massively parallel nature meant that it only took a fifth of the time to train.

Just looking at model evaluations hides some really interesting parts of the story. When we ran feature ablation, we observed that the most relevant IMUs were quite dataset dependent, in addition to the fact that the transformer and LSTM seemed to prefer vastly different sets of IMUs.

Start off by looking at figures 1 and 2. Immediately, a disparity in preferred IMUs emerges - the BiRNN prefers the pelvis, left shoulder, left wrist, right knee, upper spine, and left knee, in that order. However, the transformer prefers the pelvis, upper spine, middle spine, left foot, right hip, and left knee, in that order. Although there are certainly some overlaps - namely the pelvis, upper spine, and left knee, three of the joints differ.

Another disparity appears on a per-dataset basis. Figures 6, 3, 4, and 5 reveal the feature ablation results on the TotalCapture [11], ACCAD [4], CMU [6], and BioMotionLab_Ntroj [10] datasets. Although each dataset does seem to yield the same set of optimal IMUs as the globally optimal BiRNN IMUs, they occur with different degrees of importance.

1. 1.

   ACCAD: Pelvis, left shoulder, upper spine, left knee, left wrist, right knee

2. 2.

   BioMotionLab_Ntroj: Left shoulder, pelvis, left wrist, right knee, upper spine, left knee.

3. 3.

   CMU: Pelvis, left shoulder, left wrist, right knee, left knee, upper spine

4. 4.

   TotalCapture: Pelvis, left shoulder, left wrist, right knee, upper spine, left knee

A similar sort of disparity occurs with the transformer - take a look at the listing below:

1. 1.

   ACCAD: Pelvis, upper spine, left knee, left foot, right knee, right elbow

2. 2.

   BioMotionLab_NTroj: Upper spine, left foot, middle spine, left knee, right foot, right hip

3. 3.

   CMU: Pelvis, upper spine, left foot, right knee, middle spine, right hip

4. 4.

   TotalCapture: Pelvis, upper spine, middle spine, left foot, left knee, right foot

This time, many individual joints differ as well, rather than just being reordered. This could point to the transformer’s potential for extracting finer nuances of the data - but in either case, reveals that finding globally optimal IMUs in all situations might not be possible. The vastly differing results between datasets and model architectures reveals that IMUs should be carefully chosen in a data-driven manner with regards to individual activity types. Although this paper does try to approximate an optimal set of IMUs for each model architecture, the true result of our study is that the “optimal” set of IMUs for pose estimation should be tailored to the situation.