Deformable Attention

La atención deformable es un mecanismo de atención para visión artificial en el que cada consulta atiende a un número reducido y fijo de puntos de muestreo alrededor de una ubicación de referencia, en lugar de a todas las posiciones de un mapa de características. La propia consulta predice dónde se encuentran esos puntos, como desplazamientos respecto a la referencia, y cuánto pesa cada uno; así, el coste aumenta con el número de consultas y puntos, no con la resolución de la imagen. Xizhou Zhu, Jifeng Dai y sus colegas de SenseTime, la University of Science and Technology of China y la Chinese University of Hong Kong la presentaron en Deformable DETR, un artículo presentado oralmente en ICLR 2021. Resolvía los dos puntos débiles del Transformer de detección DETR original: la atención completa lo limitaba a características de baja resolución y el entrenamiento necesitaba 500 épocas. Deformable DETR necesitó 10 veces menos épocas y detectó mejor los objetos pequeños.

Cómo funciona la atención deformable#

La idea amplía la convolución deformable, también desarrollada en parte por Dai, que aprende desplazamientos para los puntos de muestreo de un núcleo de convolución. Para cada consulta con un punto de referencia, cada cabezal de atención aplica dos capas lineales a las características de la consulta: una predice K desplazamientos de muestreo y la otra K pesos de atención, normalizados con una softmax. Las características de valor de cada punto desplazado se leen mediante interpolación bilineal, ya que los desplazamientos pueden ser fraccionarios, y el cabezal devuelve su suma ponderada. El artículo utiliza 8 cabezales y K = 4 puntos, por lo que una consulta lee 32 posiciones por mapa de características, mientras que la autoatención estándar puntuaría cada píxel. Una versión de una sola escala y un solo cabezal se puede implementar en unas pocas líneas de PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
dim, n_points, n_queries = 32, 4, 100
value = torch.randn(1, dim, 64, 64)  # a 64x64 feature map: 4,096 positions
query = torch.randn(n_queries, dim)
ref = torch.rand(n_queries, 2) * 2 - 1  # reference points in grid_sample's [-1, 1] coordinates

# Each query predicts where to look (offsets) and how much each sample counts (weights)
to_offsets, to_weights = torch.nn.Linear(dim, n_points * 2), torch.nn.Linear(dim, n_points)
offsets = 0.1 * to_offsets(query).view(n_queries, n_points, 2)
weights = to_weights(query).softmax(dim=-1)

# Bilinearly sample the value map at reference + offset, then take the weighted sum
grid = (ref[:, None] + offsets)[None]  # (1, n_queries, n_points, 2)
sampled = F.grid_sample(value, grid, align_corners=False)[0]  # (dim, n_queries, n_points)
out = (sampled * weights).sum(dim=-1).T  # (n_queries, dim)
print(f"Output {tuple(out.shape)} from {n_queries * n_points} samples, not {n_queries * 64 * 64:,} dense scores")

Imprime Output (100, 32) from 400 samples, not 409,600 dense scores: cada una de las 100 consultas lee 4 de las 4.096 posiciones. El entrenamiento ajusta las capas de desplazamiento y peso, de modo que los puntos se acercan a las partes de la imagen que ayudan a la tarea.

Atención deformable multiescala#

Los objetos aparecen en tamaños muy distintos, así que Deformable DETR muestrea puntos de varios mapas de características a la vez. Con L niveles, cada cabezal toma K puntos de cada nivel y normaliza conjuntamente los pesos de los L × K puntos, lo que combina características finas y gruesas en un solo paso y sustituyó la red piramidal de características (FPN) de su codificador. El codificador la aplica como autoatención, con cada píxel de cada nivel como consulta, y el decodificador la aplica como atención cruzada entre las consultas de objetos y las características de la imagen.

RT-DETR, el Transformer de detección en tiempo real de Baidu disponible en el paquete Ultralytics, utiliza atención deformable multiescala como atención cruzada en sus seis capas del decodificador: cada una de sus 300 consultas de objetos muestrea 4 puntos por cabezal en 8 cabezales sobre los mapas de características con stride 8, stride 16 y stride 32, 96 posiciones en total, mientras que la atención cruzada completa sobre esos mapas de una imagen de 640 píxeles puntuaría 8.400 posiciones.

Aplicaciones#

• Detección de objetos: los sucesores de Deformable DETR, incluidos DINO y RT-DETR, mantienen la atención cruzada deformable en sus decodificadores, lo que les permite leer características de alta resolución para detectar objetos pequeños.
• Vehículos autónomos: BEVFormer proyecta imágenes de varias cámaras de un coche sobre una cuadrícula con vista cenital, donde cada celda de la cuadrícula muestrea puntos deformables de las vistas de cámara que la muestran.
• Segmentación panóptica: Mask2Former construye sus características de píxeles multiescala con atención deformable antes de predecir las máscaras.

Atención deformable frente a conceptos relacionados#

• Atención de ventana deslizante: atiende a un vecindario fijo definido por la posición, que es el mismo para todas las entradas.
• Atención dispersa: puntúa un subconjunto fijo de pares de tokens seleccionados de antemano.
• Atención Flash: calcula la atención densa habitual más rápido al dividirla en bloques en la memoria de la GPU integrada en el chip, con el mismo resultado que la atención estándar.
• Atención de Transformer de visión (ViT): compara cada parche de imagen con todos los demás, por lo que su coste aumenta con el cuadrado del número de parches.