[Paper Review] Seeing the Pose in the Pixels: Learning Pose-Aware Representations in Vision Transformers

Dominick Reilly, Aman Chadha+1 more|arXiv (Cornell University)|Jun 15, 2023

Human Pose and Action Recognition4 citations

TL;DR

This paper proposes two novel methods—Pose-aware Attention Block (PAAB) and Pose-Aware Auxiliary Task (PAAT)—to learn pose-aware representations in Vision Transformers for video understanding. By integrating 2D pose keypoints into the ViT architecture, both methods enhance performance in action recognition and multi-view video alignment, with PAAT achieving state-of-the-art results across seven datasets, improving baseline models by up to 21.8% in robotic video alignment and 9.8% in real-world action recognition.

ABSTRACT

Human perception of surroundings is often guided by the various poses present within the environment. Many computer vision tasks, such as human action recognition and robot imitation learning, rely on pose-based entities like human skeletons or robotic arms. However, conventional Vision Transformer (ViT) models uniformly process all patches, neglecting valuable pose priors in input videos. We argue that incorporating poses into RGB data is advantageous for learning fine-grained and viewpoint-agnostic representations. Consequently, we introduce two strategies for learning pose-aware representations in ViTs. The first method, called Pose-aware Attention Block (PAAB), is a plug-and-play ViT block that performs localized attention on pose regions within videos. The second method, dubbed Pose-Aware Auxiliary Task (PAAT), presents an auxiliary pose prediction task optimized jointly with the primary ViT task. Although their functionalities differ, both methods succeed in learning pose-aware representations, enhancing performance in multiple diverse downstream tasks. Our experiments, conducted across seven datasets, reveal the efficacy of both pose-aware methods on three video analysis tasks, with PAAT holding a slight edge over PAAB. Both PAAT and PAAB surpass their respective backbone Transformers by up to 9.8% in real-world action recognition and 21.8% in multi-view robotic video alignment. Code is available at https://github.com/dominickrei/PoseAwareVT.

Motivation & Objective

• Address the limitation of conventional Vision Transformers in handling subtle, viewpoint-varied actions in daily-living videos by incorporating pose priors.
• Overcome the uniform patch processing in standard ViTs that neglects pose-based entities like human skeletons or robotic arms.
• Develop efficient, plug-and-play methods that integrate 2D pose keypoint information into ViTs without requiring depth sensors or complex 3D pose estimation.
• Enable viewpoint-agnostic, fine-grained video representation learning by leveraging pose information during training.
• Demonstrate generalization across diverse downstream tasks, including action recognition, video retrieval, and multi-view robotic video alignment.

Proposed method

• Introduce the Pose-aware Attention Block (PAAB), a plug-and-play ViT block that applies localized self-attention specifically on pose-related patches identified via 2D keypoint heatmaps.
• Design the Pose-Aware Auxiliary Task (PAAT), a multi-task learning strategy where a separate head predicts 2D keypoint locations during training, jointly optimized with the main video classification or alignment task.
• Utilize readily available 2D pose estimation models (e.g., MediaPipe) to extract accurate, low-cost keypoint supervision for both PAAB and PAAT.
• Maintain the original ViT architecture and training pipeline, with only minor modifications: adding a pose-aware attention mechanism (PAAB) or an auxiliary loss (PAAT).
• Ensure that the pose supervision is applied at the patch level, allowing the model to learn disentangled representations between pose and non-pose features.
• Leverage feed-forward layers to learn pose-aware representations, as shown by ablation studies indicating minimal changes in attention distributions compared to the baseline.

Experimental results

Research questions

• RQ1 Can integrating 2D pose keypoint information into Vision Transformers improve performance on video understanding tasks involving subtle or viewpoint-varied actions?
• RQ2 Does the use of pose priors in ViTs lead to more disentangled and viewpoint-agnostic feature representations compared to standard ViTs?
• RQ3 To what extent do attention mechanisms or feed-forward layers contribute to learning pose-aware representations in ViTs?
• RQ4 How does the performance of pose-aware ViTs generalize across diverse datasets and downstream tasks such as action recognition and multi-view video alignment?
• RQ5 Is the auxiliary pose prediction task (PAAT) more effective than architectural modifications (PAAB) for learning pose-aware features in ViTs?

Key findings

• PAAT achieves state-of-the-art performance across seven datasets, outperforming the baseline ViT by up to 21.8% in multi-view robotic video alignment and 9.8% in real-world action recognition.
• Both PAAB and PAAT significantly improve feature separability between pose and non-pose tokens in the feature space, with PAAT showing superior disentanglement due to its keypoint-specific prediction objective.
• The attention distributions of PAAB and PAAT remain largely unchanged compared to the baseline TimeSformer, indicating that pose-aware representations are primarily learned through the feed-forward layers, not attention mechanisms.
• PAAT demonstrates higher effectiveness in early network layers, suggesting that pose priors are most beneficial during initial feature abstraction.
• PAAT requires no additional parameters or inference computation beyond the main ViT, making it more efficient than PAAB, which adds a dedicated attention block.
• The proposed methods generalize well across tasks and datasets, including cross-data evaluation for video retrieval, confirming their robustness and transferability.