◧ Base de Conocimiento ANE
EN ES

Capítulos · 04

Despliegue de Vision Transformers basados en atención en la ANE

Parte de la ANE Knowledge Base. Fuente: Apple ML Research, "Deploying Attention-Based Vision Transformers to Apple Neural Engine" (2024), código en references/ml-vision-transformers-ane.

Esta es la continuación del trabajo sobre transformers de 2022, que extiende los cuatro principios a los vision transformers (ViT) y sus retos adicionales: mapas de características 2D, entradas de alta resolución, partición de ventanas y position embeddings. Apple cita la búsqueda de Photos, RoomPlan y las características semánticas de ARKit como usuarios en producción de este tipo de modelos.

Se estudian dos arquitecturas:

  • DeiT/16-tiny — un ViT convencional, como baseline.
  • MOAT (tiny-MOAT-0/1/2) — una red híbrida de conv Mobile + atención: primero etapas MBConv (residual invertido), después etapas de atención. Apta para móvil; mayor precisión en ImageNet que DeiT con un número de parámetros similar.

El repositorio incluye WindowAttention, window_partition/window_reverse, la implementación de MOAT y export.py.

1. Entradas de alta resolución: atención local (de ventana)#

Problema. La atención es cuadrática en la longitud de tokens. Con una entrada de 512×512, la atención global sobre todo el mapa de características es prohibitivamente lenta.

Solución. Particionar el mapa de características en ventanas rectangulares y aplicar atención dentro de cada ventana (como en Swin/MOAT). El flujo de información entre ventanas se recupera por otros medios: Swin usa ventanas desplazadas (shifted windows), MOAT usa sus convoluciones depthwise (la mitad MBConv de cada bloque), lo cual es conveniente en la ANE porque las ventanas desplazadas no tienen soporte nativo en la ANE (WindowAttention explícitamente "supports only non-shifting window attention").

En el repositorio (vision_transformers/model.py):

  • attention_mode="global": ventana = todo el mapa de características (adecuado a baja resolución / etapas tardías).
  • attention_mode="local": ventana fija (p. ej., 8×8; los candidatos por defecto 6/8/10 se eligen para dividir el tamaño de la característica; tamaños restringidos a 6–16).
  • Etapas MOAT: ("mbconv", "mbconv", "moat", "moat") — la atención solo entra en los strides 16/32, donde el número de tokens es manejable.

2. El problema del tensor 6D: partición de ventanas "por relevos" en ≤5D#

Problema. La partición/reversión estándar de ventanas pasa por un tensor 6D (N, C, Nh, Nw, Hw, Ww) (ventanas-por-eje × tamaño-de-ventana-por-eje). La ANE soporta como máximo tensores 5D.

Solución. Realizar la partición "por relevos" (relay), factorizando una dimensión a la vez — primero la altura, luego el ancho — de modo que ningún intermedio supere 5D. De vision_transformers/attention_utils.py:

def window_partition(x, window_size):            # x: (B, H, W, C)  — NHWC!
    B, H, W, C = x.shape
    x = x.reshape((B, H // wh, wh, W, C))         # 5D: split H
    x = x.reshape((B * H // wh, wh, W, C))        # 4D: fold into batch
    x = x.reshape((B * H // wh, wh, W // ww, ww, -1))  # 5D: split W
    x = x.permute((0, 2, 1, 3, 4))
    windows = x.reshape((-1, wh, ww, C))          # (B·num_windows, wh, ww, C)
    return windows

window_reverse refleja los mismos reshapes escalonados a la inversa. El comentario del código muestra el relevo de shapes: 1,12,16,160 → 1,2,6,16,160 → 2,6,16,160 → 2,6,2,8,160 → ...

Truco generalizable: cualquier op que "necesite" rango > 5 normalmente puede escalonarse como reshapes de plegado en el batch (fold-into-batch) más un pequeño permute.

3. Layout de tensores para la partición: NHWC, no NCHW#

Problema. La ANE procesa la última dimensión en lotes de 64 bytes y la rellena hasta 64 bytes si es menor. Los tamaños de ventana son pequeños: con NCHW y una ventana de 7×7, la última dimensión contiene 7 valores FP16 = 14 bytes → 50 bytes de relleno (padding). En el peor caso (última dimensión unitaria) se produce una ralentización efectiva de 32× en FP16.

Solución. Para la ruta de partición/reversión, usar NHWC: la dimensión de canal (normalmente un múltiplo de 32) queda en último lugar y se alinea de forma natural. El bloque MOAT hace un transpose NCHW→NHWC una vez antes de LayerNorm + partición y vuelve a NCHW una vez después de la reversión de ventanas — "instead of looping on each partitioned window":

# vision_transformers/model.py — MOATBlock.forward
output = self._mbconv(inputs)                 # NCHW
shortcut = output
output = output.permute((0, 2, 3, 1))         # NHWC, once
assert output.shape[-1] % 32 == 0, "ANE buffer not aligned"
output = self._attn_norm(output)              # nn.LayerNorm on last dim (=C)
...window_partition → WindowAttention → window_reverse...
output = output.reshape((N, H, W, C)).permute((0, 3, 1, 2))  # back to NCHW, once
output = shortcut + output

Nótese la interacción de las restricciones: NHWC también permite que un nn.LayerNorm normal funcione (normaliza la última dimensión), y _build_model(channel_buffer_align=True) redondea todos los tamaños ocultos hacia arriba a múltiplos de 32 canales para mantener la última dimensión alineada a 64 bytes.

Conclusión: el layout "correcto" es dependiente de la op. Los transformers de secuencia ponen S al final (BC1S); la partición de ventanas pone C al final (NHWC). Lo que es invariante es la regla: el último eje debe ser grande y compatible con 64 bytes.

4. Interior de la atención: la misma receta que en NLP, más split_softmax#

Dentro de WindowAttention (vision_transformers/attention_utils.py) reaparece la receta de 2022:

  • Proyecciones Conv2d 1×1 separadas para Q, K, V — el código comenta "Use separate conv1x1 projection to avoid L2 cache hit" (frente a una proyección QKV fusionada).
  • torch.split por cabeza + einsum bchq,bkhc->bkhq / bkhq,bchk->bchq, un transpose sobre K ("transposeThenSplit is more efficient than the other way around").
  • split_softmax / split_head=True (por defecto): softmax aplicado por cabeza en dim=1. "Splitting on the softmax… increases the chance of L2 residency and parallelizes the computation for the softmax layer" — señalado como la clave de la ganancia de latencia, ya que el softmax es cuadrático en la longitud de tokens. (Existen alternativas de atención lineal como CosFormer, pero sacrifican precisión.)

5. Position embeddings: tamaño de archivo frente a longitud de tokens#

Para los ViT, la elección del position embedding afecta materialmente al tamaño de archivo del modelo y a la latencia en el dispositivo (enum PEType en el repositorio):

Tipo de PE Parámetros extra Crecimiento Notas
RPE (PE relativa estilo Swin) tabla num_heads × (2Wh−1)(2Ww−1), recopilada (gathered) a token_len² por cabeza cuadrático en la longitud de tokens Sobrecarga significativa de tamaño de archivo y latencia con ventanas grandes
SINGLE_HEAD_RPE una tabla compartida entre cabezas cuadrático ÷ num_heads "Reduces the file size of the positional embedding to 1/num_heads of the original RPE"
LePE_ADD (PE localmente mejorada) conv depthwise 3×3 sobre V (3·3·dim por bloque) + PE absoluta 1 × token_len × dim lineal en la longitud de tokens La conv depthwise codifica localidad en el tensor de valores; "significantly smaller than RPE"

Notas de implementación de WindowAttention:

  • El sesgo RPE se busca en relative_position_bias_table mediante un relative_position_index precomputado y se añade por cabeza antes del split softmax — con RPE de una sola cabeza, todas las cabezas suman la misma tabla (rpe_idx = 0).
  • LePE ejecuta LePE_for_Value (depthwise 3×3, groups=dim, padding="same") sobre el mapa de valores 2D y lo añade por cabeza después de attn @ v; la PE absoluta se añade a los tokens de entrada.

6. Resultados#

  • El tiny-MOAT-1 optimizado es "multiple times faster" en la ANE que la implementación de MOAT open-source de terceros.
  • A alta resolución (512×512), el MOAT optimizado también es mucho más rápido que el DeiT/16-tiny optimizado — la localidad gana cuando el número de tokens se dispara.
  • El tiny-MOAT-1 tiene además mayor precisión en ImageNet que DeiT-tiny con un número de parámetros similar.
  • La Figura 3 del artículo / los assets del repositorio muestran mediciones en dispositivo con Xcode en distintos modelos de iPhone.

7. Checklist para llevar un ViT a la ANE#

  1. Preferir un backbone híbrido conv+atención (tipo MOAT); mantener la atención en las etapas de baja resolución.
  2. Entradas de alta resolución → atención local de ventana, con la ventana dimensionada para dividir el mapa de características (6–16 por lado); sin ventanas desplazadas.
  3. Implementar la partición/reversión con el patrón de relevo 5D; NHWC durante la partición; transpose una vez a la entrada, una vez a la salida.
  4. Mantener los canales como múltiplo de 32 (channel_buffer_align).
  5. Usar conv 1×1 separadas para Q/K/V, split por cabeza, matmuls con einsum, split softmax.
  6. Preferir LePE (o RPE de una sola cabeza) frente a RPE completo para longitudes de tokens grandes.
  7. Exportar con torch.jit.trace + ct.convert(convert_to="mlprogram") y revisar el informe de rendimiento de Xcode (doc 05).

Siguiente: 05 — Deployment Workflow.

Generado desde el markdown de la base de conocimiento — cada afirmación traza a una fuente citada.