# Permutohedral Attention Module for Efficient Non-Local Neural Networks

## TL;DR

This paper introduces the Permutohedral Attention Module (PAM), a memory- and computation-efficient attention mechanism designed to capture non-local information in 3D medical images, improving segmentation accuracy.

## Contribution

The paper presents a novel, efficient attention module for neural networks that effectively captures non-local context in 3D medical imaging tasks.

## Key findings

- Demonstrates improved vertebrae segmentation accuracy.
- Shows scalability and efficiency of PAM in 3D medical imaging.
- Provides GPU implementation suitable for real-world applications.

## Abstract

Medical image processing tasks such as segmentation often require capturing non-local information. As organs, bones, and tissues share common characteristics such as intensity, shape, and texture, the contextual information plays a critical role in correctly labeling them. Segmentation and labeling is now typically done with convolutional neural networks (CNNs) but the context of the CNN is limited by the receptive field which itself is limited by memory requirements and other properties. In this paper, we propose a new attention module, that we call Permutohedral Attention Module (PAM), to efficiently capture non-local characteristics of the image. The proposed method is both memory and computationally efficient. We provide a GPU implementation of this module suitable for 3D medical imaging problems. We demonstrate the efficiency and scalability of our module with the challenging task of vertebrae segmentation and labeling where context plays a crucial role because of the very similar appearance of different vertebrae.