ASFormer: Transformer for Action Segmentation

The input for the encoder is the pre-extracted feature sequences of size $T\times D$, where $T$ is the video length and $D$ is the feature dimension. The first layer of the encoder is a fully connected layer that adjusts the dimension of the input feature. Then, this layer is followed by serials of encoder blocks. After that, a fully connected layer will output predictions $y_{e}\in{R}^{T\times C}$ from the last encoder block, where $C$ denotes the number of action classes.

Each encoder block contains two sub-layers. The first is a feed-forward layer, and the second is a single-head self-attention layer. We employ a residual connection around each of the two sub-layers, followed by instance normalization and ReLU activation, as shown in Fig. 1(a). Different from the vanilla Transformer, we use a dilated temporal convolution as the feed-forward layer instead of point-wise fully connected layer. This design is inspired by the properties of the action segmentation task, which are a) lack of the large training dataset for completely free inductive bias. b) high locality of features since every action occupies continued timestamps in the input video. In contrast to the fully connected layer, the temporal convolution layer can bring beneficial local inductive bias to our model.

The self-attention layer is hard to be learned to focus on meaningful locations over thousands of frames. For an input video, those self-attention layers are difficult to cooperate with each other to form an effective representation. To mitigate this issue, we pre-define a hierarchical representation pattern. Such a hierarchical pattern is inspired from modern neural network design: one would first focus on the local feature and then gradually enlarge the receptive field to capture the global information. For example, CNNs achieve such pattern with serials of pooling layers to enlarge the receptive field in higher layers; or use dilated convolutions with gradually increasing dilation rate. Motivated by the success of such a hierarchical pattern, we constraint the receptive fields of each self-attention layer within a local window with size $w$ (e.g for a frame $t$, we only calculate the attention weights with the frames that within its local window). The size of the local window is then doubled at each layer $i$ (i.e., $w=2^{i},i=1,2...$). Meanwhile, we also double the dilation rate of the temporal convolution layer with the encoder depth increasing, keeping consistent with the self-attention layer.

For an encoder with ${J}$ blocks, the whole approximate memory usage for a vanilla Transformer is $({J}\cdot T\cdot T)$, where $T$ is the video length. With the hierarchical representation pattern, We reduce the total space complexity to $((2-\epsilon)\cdot 2^{J}\cdot T)$, where $\epsilon$ is a small number. In our settings, we use ${J}=9$, where $2^{J}=512$ is almost 10 times smaller than $T$. Compared to the vanilla Transformer, our ASFormer is applicable to receiving long input sequence.

Temporal relations among multiple action segments play an important role in the action segmentation task. There are some prior relationship between actions segments, e.gthe action after take bottle and pour water usually to be drink water. As previous works showed, applying additional TCNs [Farha and Gall(2019), Wang et al.(2020b)Wang, Gao, Wang, Li, and Wu] or GCNs [Huang et al.(2020)Huang, Sugano, and Sato, Wang et al.(2021a)Wang, Hu, Li, and Dou] over the initial prediction to perform a refinement process can boost the performance. In this section, we illustrate how the new designed decoder carries out the refinement task for the initial predictions output by the encoder in one-forward-pass. For better explanation, we first introduce a single decoder, and naturally extend it to a multiple version to perform an iterative refinement.

Compared to the self-attention layer, the cross-attention has the following difference: The query $Q$ and key $K$ are obtained from the concatenation of the output from the encoder and previous layer, while the value $V$ is only obtained from the output of the previous layer. The cross-attention mechanism allows every position in the encoder to attend over all positions in the refinement process by generating the attention weights. The feature space $V$ is completely transformed from the input predictions, and will not be disturbed with the participant of the encoder, because the generated attention weights are only used to perform a linear combination within $V$. Such design is inspired by the previous work [Farha and Gall(2019)], where they show the refinement process is very sensitive to the disturbance of the learned feature space from the predictions.

Multiple Decoders. One would naturally expand the single decoder to a multiple version to perform iterative refinement. In the multiple version, the input of each decoder are from the previous one, as shown in Fig. 1(b).

The cross-attention mechanism allows to bring in external information to guide the refinement process. We hope to gradually reduce the weight of external information to avoid the problem of error accumulation. For the input $x$ in each decoder block, we use a weighted residual connection for the output of the feed-forward layer and the cross-attention layer:

We set the $\alpha=1$ for the first decoder and then exponentially decrease $\alpha$ for rest decoders.

The loss function is a combination of classification loss $\mathcal{L}_{cls}$ for each frame and smooth loss $\mathcal{L}_{smo}$ [Farha and Gall(2019)]. The classification loss is a cross-entropy loss, while the smooth loss calculates the mean squared error over the frame-wise probabilities. The final loss function $\mathcal{L}$ is,

where $y_{t,\hat{c}}$ is the the predicted probability for the ground truth label $\hat{c}$ at time t. $\lambda$ is a balance weight that is set to 0.25 in our experiments. Finally to train the complete model, the sum of the losses over the encoder and all decoders is minimized.

The final ASFormer consists of one encoder and three decoders, while each encoder and the decoder contains nine blocks. The dimension of the first fully connected layer in the encoder and decoder is set to 64, as well as the feature dimension in each encoder and decoder block. Moreover, a special dropout layer is applied to the input feature of the encoder, which randomly drops the entire feature channel with a dropout rate of 0.3. In all experiments, we train the model for 120 epochs through Adam optimizer with a learning rate 0.0005.

In terms of implementation details, our ASFormer has two differences compared to the vanilla Transformer. First, the position encoding is not applied for the input of both encoder and decoders. Second, vanilla Transformer usually adapts multi-head attention that concatenates the output from multiple single-head attention to enhance the feature transformation ability. Instead, we only use a single-head attention in each encoder/decoder block. We conduct the following two experiments to demonstrate the introduction of temporal convolution makes our model free from position encoding and multi-head version of self-attention.

The first experiment studies whether position encoding is still needed for ASFormer. The results on 50Salads dataset are shown in Table 5.1. For the temporal convolution, the position encoding is even counterproductive. When the position encoding is only applied for the encoder, the performance already drops a lot. Further, when the position encoding is applied for both encoder and decoders, the performance is even worse. The possible reason is that the temporal convolution already has the ability to model the relative positional relationships. The redundant absolute position encoding might be harmful to the temporal convolutions to learn the feature embedding. Thus, we drop the position encoding from our model. The second experiment studies the impact of multi-head self-attention. Table 5.1 shows the comparison of multi-head self-attention with the different numbers of heads on 50Salads. For each head in the multi-head attention, the feature dimension is the same as the single-head attention. We can observe that the single head self-attention achieves comparable performance with the multiple ones. This is because the convolution operation has a similar form with the multi-head self-attention operation (e.g a 64-filters convolution operator also concatenates the output from 64 1-filter convolutions). Considering the extra computation and memory budgets brought by the multi-head self-attention, we use the single-head self-attention by default.

In this section, we study the effect of bringing in local connective inductive bias with the temporal convolution. For comparison, we use an MLP as the feed-forward layer as the same as the vanilla Transformer. It worth noting that when removing the temporal convolution from our model, the position encoding is needed to bring in position information as the vanilla transformer does. In order to directly compare the the two solutions, we only use the encoder part for the experiments to avoid the impact of the refinement process. Results on the 50Salads dataset are shown in Table 3. We can observe that when MLP is used as the feed-forward layer, the performance drops greatly, especially on F1 scores and Edit score, which denotes that the model fails to model the temporal relationship among frames to produce smoother and consistent predictions. This also demonstrates that the fusion of neighboring local information plays an important role in the performance.

The hierarchical representation pattern plays an important role in our ASFormer. To demonstrate its effectiveness, we conduct a non-hierarchical version by setting the size of the local window in all attention layers to 512, which is the largest window size in the last encoder/decoder block (out of memory if we directly set the size of the local window in all attention layers to the video length). The non-hierarchical version allows each self-attention layer to span their attention weights ‘freely’. Experiments on 50Salads dataset are shown in Table 4. The performance of the non-hierarchical version drops significantly.

To better understand why there is such a huge performance gap, we show some visualization results in Fig. 2. For an anchor query frame, we plot its attention weights in each self-attention layer in the encoder (The attention weights are normalized with min-max¹¹ 1 $x=(x-min)/(max-min)$.). Fig. 2(a) shows the non-hierarchical version, while Fig. 2(b) shows the case that with the pre-defined hierarchical pattern. We have two observations. First, for the high-level blocks (e.gblock #9) that has the same window size, there are much more activation (blue ribbon) in the non-hierarchical version after the min-max normalization. This means that the attention weights in many locations have close values and more similar to a trivial uniform distribution. In contrast, self-attention trained with the hierarchical pattern tend to focus on several meaningful locations. Secondly, the ‘freely’ attention cannot automatically learn a hierarchical pattern from data. As shown in Fig. 2(a), the low-level blocks do not cast attention on its neighbors when the self-attention layer is not constrained.

To demonstrate that our decoder takes advantage of temporal relations among multiple action segments for refinement, we conduct an ablation study about stacking the different numbers of decoders. Results on 50Salads dataset are shown in Table. 5.5. The decoders largely boost the performance compared to the encoder, and we achieve the best results when stacking three decoders. To show the iterative refinement process, we further plot the predictions in the encoder and all decoders, as shown in the supplementary material.

The number of blocks ${J}$ in the encoder/decoder is an important hyper-parameter, where more blocks bring larger receptive fields but also lead to higher memory costs as introduced in Sec. 3.1. We conduct an ablation study about the different number of blocks in the encoder/decoder on 50salads dataset and report the peak value of the GPU memory cost for batch size 1 during training. Results are shown in Table 5.5. We achieve higher performance with the increasing number of blocks until ${J}=9$. Using more than 9 blocks ($J=10$) slightly improves the frame-wise accuracy but does not increases the F1 scores. Meanwhile, the GPU memory cost increases dramatically with the large $J$. Thus, we choose ${J}=9$ for all our experiments.

This section compares the proposed ASFormer to the state-of-the-art methods on three datasets: 50Salads, GTEA, and Breakfast datasets. The results are presented in Table 7. Our model achieves the state-of-the-art methods on the three datasets compared to previous work. The results on three metrics highlight the ability of our ASFormer to obtain accurate and smooth predictions.

It is worth noting that, our ASFormer can serve as a strong backbone model and thus can be easily adapted by other works. To demonstrate that, we replace the original TCN-based backbone model MS-TCN [Farha and Gall(2019)] in ASRF [Ishikawa et al.(2020)Ishikawa, Kasai, Aoki, and Kataoka] with our ASFormer. The new model, denoted as ASRF*, achieves even higher results on the 50salads dataset than the original ASRF [Ishikawa et al.(2020)Ishikawa, Kasai, Aoki, and Kataoka], as shown in Table 5.6. Besides the performance, we further compare with the popular TCN-based backbone MS-TCN [Farha and Gall(2019)] in terms of the computation cost, as shown in Table 5.6. The additional computational burden of our ASFormer is acceptable with the modern chips. Thus, our ASFormer provides an alternative backbone choice for the community to build their models.