Indexação Avançada de Tensores (Máscaras Booleanas)

Visão geral

A indexação por máscara booleana é uma forma de indexação avançada no PyTorch que seleciona ou atualiza elementos de tensores usando um tensor booleano (torch.bool) com o mesmo formato (ou formato compatível). Ela é amplamente usada em código de ML para tarefas como:

filtrar tokens de padding
zerar contribuições de perda para entradas inválidas
mascarar logits de atenção antes do softmax na Arquitetura Transformer
atribuição condicional e limpeza de dados

No PyTorch, a indexação por máscara booleana se comporta mais como uma operação de coleta/espalhamento (gather/scatter) do que como uma visão por fatiamento (slice view). Isso tem consequências importantes para formatos, broadcasting, desempenho e autograd.

Ideia básica: “manter elementos onde a máscara é True”

Uma máscara booleana é um tensor de valores True/False, geralmente criado por comparações:

import torch

x = torch.tensor([0.2, -1.3, 4.5, 0.0])
mask = x > 0          # tensor([ True, False,  True, False])
y = x[mask]           # tensor([0.2000, 4.5000])

Isso é equivalente a torch.masked_select(x, mask):

y2 = torch.masked_select(x, mask)

Requisitos de dtype e dispositivo da máscara

A máscara deve ser torch.bool (código antigo pode usar uint8, mas bool é o dtype moderno e correto).
A máscara deve estar no mesmo dispositivo do tensor que você indexa (CPU vs CUDA).

x = torch.randn(10, device="cuda")
mask = (x > 0)        # mask is also cuda:0
x[mask]               # OK

Indexação avançada vs básica: por que isso importa

A indexação no PyTorch se divide em duas categorias amplas:

Indexação básica (fatias como x[:, 5:10], escalares inteiros como x[3]): frequentemente retorna uma visão (barata, sem cópia).
Indexação avançada (máscaras booleanas, tensores de índices inteiros, listas): retorna um novo tensor (cópia) e tipicamente envolve coleta (gather) de elementos.

Máscaras booleanas são indexação avançada. Isso leva a:

Alocação de um novo tensor de saída
Potencialmente trabalho não trivial de kernel (muitas vezes envolve converter internamente a máscara em índices)
Regras de formato diferentes das fatias

Semântica de formatos: qual formato você obtém de volta?

A indexação por máscara booleana pode:

selecionar ao longo de uma dimensão específica (máscara é 1D e usada em uma posição do índice), ou
selecionar elementos individuais (máscara corresponde ao formato completo do tensor), produzindo uma seleção achatada.

Caso A: máscara 1D usada para uma dimensão (filtragem de linhas/colunas)

Se você usa uma máscara 1D em uma posição específica do índice, ela seleciona ao longo daquela dimensão e preserva as demais dimensões.

Selecionar linhas:

x = torch.arange(12).view(4, 3)
row_mask = torch.tensor([True, False, True, False])  # shape (4,)
x_sel = x[row_mask]   # shape (2, 3)

Selecionar colunas:

col_mask = torch.tensor([False, True, True])  # shape (3,)
x_sel = x[:, col_mask]  # shape (4, 2)

Regra-chave: o comprimento da máscara deve corresponder ao tamanho da dimensão que ela indexa.

Caso B: máscara corresponde ao formato do tensor (filtragem de elementos → resultado 1D)

Se a máscara tem o mesmo formato do tensor indexado, você obtém um tensor 1D com os elementos selecionados em ordem row-major.

x = torch.arange(6).view(2, 3)
mask = x % 2 == 0      # True on even entries, shape (2, 3)
y = x[mask]            # shape (3,)  -> tensor([0, 2, 4])

Esse “achatamento para 1D” surpreende quem vem de modelos mentais de “seleção de linhas”. A distinção é:

x[row_mask] (máscara indexa uma dimensão) → mantém a estrutura
x[full_mask] (máscara indexa elementos) → achata os elementos selecionados

Caso C: máscara sobre dimensões iniciais (comum em ML)

Um padrão muito comum em ML é x com formato (B, T, H) e uma máscara com formato (B, T):

x = torch.randn(2, 4, 3)              # (B=2, T=4, H=3)
mask = torch.tensor([[1,1,0,0],
                     [1,0,1,0]]).bool()  # (2,4)

y = x[mask]  # shape (num_true, H)
print(y.shape)  # torch.Size([4, 3])

Interpretação: a máscara booleana indexa as duas primeiras dimensões, e as dimensões finais remanescentes são preservadas.

Isso é extremamente útil para “filtrar tokens válidos ao longo de um batch”, mas também destrói a estrutura (B, T). Se você precisa manter o formato, normalmente deve usar masked_fill ou where em vez disso (veja abaixo).

Broadcasting e “fazendo os formatos se alinharem”

Máscaras booleanas não fazem broadcasting automaticamente de todas as formas que iniciantes esperam. Na prática, muitas vezes você precisa de unsqueeze para alinhar dimensões.

Padrão: manter o formato e mascarar a última dimensão

Suponha que x seja (B, T, H) e mask seja (B, T) e você queira zerar estados ocultos para tokens de padding sem mudar o formato:

x = torch.randn(2, 4, 3)
mask = torch.tensor([[1,1,0,0],
                     [1,0,1,0]]).bool()  # (B,T)

x_masked = x.masked_fill(~mask.unsqueeze(-1), 0.0)  # (B,T,H)

mask.unsqueeze(-1) vira (B, T, 1) e faz broadcasting ao longo de H.
masked_fill preserva o formato, enquanto x[mask] colapsaria (B, T) em uma dimensão.

Padrão: mascarar logits no formato (B, heads, T, T)

Máscaras de atenção frequentemente têm formato (B, 1, 1, T) e fazem broadcasting sobre cabeças e posições de consulta:

B, heads, T = 2, 4, 8
logits = torch.randn(B, heads, T, T)

key_padding_mask = torch.randint(0, 2, (B, T)).bool()  # True = keep
attn_mask = ~key_padding_mask[:, None, None, :]        # (B,1,1,T) True = mask out

neg_inf = torch.finfo(logits.dtype).min
logits = logits.masked_fill(attn_mask, neg_inf)

Este é um passo padrão antes do softmax em muitas Mecânicas de Atenção.

Misturando máscaras booleanas com outros tipos de índice

O PyTorch permite combinar máscaras booleanas com fatias/inteiros/tensores de índices, mas você precisa acompanhar quais dimensões cada índice consome.

Exemplo: selecionar apenas as linhas mascaradas e, em seguida, pegar uma fatia específica de colunas:

x = torch.arange(20).view(5, 4)
row_mask = torch.tensor([True, False, True, False, True])

y = x[row_mask, 1:3]  # shape (3, 2)

Quando você mistura índices avançados (booleanos ou tensores de índice) em múltiplas dimensões, o PyTorch segue regras de “indexação avançada” semelhantes às do NumPy: índices avançados podem alterar ordem/formato de maneiras nem sempre intuitivas. Se clareza de formato for importante, muitas vezes é melhor usar nonzero + index_select/gather explicitamente.

Padrões comuns em código de ML

1) Filtragem de exemplos/tokens (dados de comprimento variável)

Você frequentemente quer remover tokens de padding para uma perda calculada por token:

logits = torch.randn(2, 4, 10)     # (B,T,V)
targets = torch.randint(0, 10, (2, 4))
valid = torch.tensor([[1,1,0,0],
                      [1,0,1,0]]).bool()

# Flatten batch and time, then filter valid positions
logits_valid = logits[valid]       # (N_valid, V)
targets_valid = targets[valid]     # (N_valid,)

loss = torch.nn.functional.cross_entropy(logits_valid, targets_valid)

Isso é conciso e correto, mas tenha em mente que muda o formato de (B,T,V) para (N_valid,V).

2) Obtendo índices a partir de uma máscara (`nonzero`)

Às vezes você quer índices explícitos em vez de filtrar:

mask = torch.tensor([[True, False, True],
                     [False, True, False]])

idx = mask.nonzero(as_tuple=False)   # shape (N_true, 2), each row is [i, j]
i, j = mask.nonzero(as_tuple=True)   # two 1D tensors

Esses índices são úteis com index_select, gather ou para depuração.

3) Atribuição mascarada (atualização in-place)

Você pode atribuir em um tensor usando uma máscara booleana:

x = torch.randn(5)
mask = x < 0
x[mask] = 0.0

As regras de formato para atribuição importam:

Se x[mask] produz uma seleção 1D (máscara de elementos), o lado direito (RHS) deve poder fazer broadcasting para (#true,) ou ser um escalar.
Se a máscara seleciona ao longo de uma dimensão (ex.: máscara de linhas), o RHS deve corresponder ao formato do bloco selecionado ou ser broadcastable.

Exemplo de atribuição de linhas:

x = torch.zeros(4, 3)
row_mask = torch.tensor([True, False, True, False])
x[row_mask] = 1.0      # fills two rows with 1.0

4) Mascarando logits com segurança antes do softmax

Ao mascarar logits, evite -1e9 hardcoded (pode underflow de maneira estranha em float16/bfloat16). Use “infinito negativo” sensível ao dtype:

logits = torch.randn(2, 5)
mask_out = torch.tensor([[False, True, False, False, True],
                         [False, False, False, True, False]])

neg_inf = torch.finfo(logits.dtype).min
masked_logits = logits.masked_fill(mask_out, neg_inf)
probs = torch.softmax(masked_logits, dim=-1)

Esta é uma boa prática em implementações da Arquitetura Transformer.

5) Mantendo o formato: `where` vs indexação booleana

Se você precisa preservar o formato e apenas “desligar” certos valores:

x = torch.randn(3, 3)
mask = x > 0

y = torch.where(mask, x, torch.zeros_like(x))  # same shape as x

Use indexação booleana (x[mask]) quando você quer mudar o formato ao filtrar. Use where / masked_fill quando você quer preservar o formato.

Características de desempenho e alternativas

A indexação booleana é conveniente, mas pode ser mais lenta e exigir mais alocações do que alternativas.

Por que a indexação booleana pode ser mais lenta

Internamente, a seleção booleana frequentemente envolve:

converter a máscara em índices (nonzero)
coletar/espalhar elementos
alocar um novo tensor de saída dimensionado pelo número de entradas True

Na GPU, nonzero e saídas de tamanho variável podem ser relativamente caros e podem inibir fusão/otimização.

Alternativas e quando usá-las

Use fatiamento para intervalos contíguos (rápido, visão)

Se você está selecionando um bloco contíguo, fatiamento é o melhor:

x = torch.randn(1024, 1024)
y = x[:, 100:200]   # view, very cheap

Máscaras booleanas não conseguem produzir visões.

Use `index_select` para selecionar ao longo de uma dimensão com índices inteiros

Se você já tem índices inteiros (ou pode calculá-los uma vez e reutilizar), index_select costuma ser mais claro e pode ser mais rápido/mais previsível:

x = torch.randn(1000, 64)
idx = torch.tensor([0, 5, 9, 10], device=x.device)
y = torch.index_select(x, dim=0, index=idx)  # (4, 64)

Você pode derivar idx a partir de uma máscara:

idx = row_mask.nonzero(as_tuple=True)[0]
y = x.index_select(0, idx)

Mas note: se você fizer isso repetidamente, calcular nonzero toda vez pode dominar o tempo de execução.

Use `gather` para seleção elementwise em batch (mantém o formato)

gather é a escolha padrão quando cada elemento do batch seleciona posições diferentes, mas você quer preservar a estrutura do batch:

# x: (B, T, H), choose one position per batch element
B, T, H = 4, 8, 16
x = torch.randn(B, T, H)
pos = torch.randint(0, T, (B, 1, 1))          # (B,1,1)
pos = pos.expand(B, 1, H)                      # (B,1,H)

chosen = x.gather(dim=1, index=pos)            # (B,1,H)
chosen = chosen.squeeze(1)                     # (B,H)

Uma máscara booleana não consegue expressar essa seleção de “um índice por linha” sem processamento extra.

Prefira `masked_fill`/`where` para mascaramento que mantém formato

Se sua operação é “mascarar valores, mas manter o formato do tensor”, esses métodos normalmente são mais rápidos e mais idiomáticos do que seleção booleana + tentativas de reshape:

x = torch.randn(2, 4, 3)
mask = torch.rand(2, 4) > 0.5

x2 = x.masked_fill(~mask.unsqueeze(-1), 0.0)  # keeps (2,4,3)

Comportamento do autograd e armadilhas

A indexação por máscara booleana interage com o autograd de formas previsíveis, mas importantes.

Gradientes fluem apenas pelos elementos selecionados

Se y = x[mask], então durante o backprop, os gradientes são espalhados de volta em x nas posições selecionadas; todas as outras posições recebem gradiente zero.

Isso geralmente é o que você quer ao filtrar tokens válidos em uma perda.

Máscaras não são diferenciáveis

Uma máscara booleana é discreta. Se sua máscara é computada a partir de x usando comparações (mask = x > 0), a criação da máscara é não diferenciável. Gradientes não se propagam “através da decisão”.

Para gating aprendível, normalmente você usa uma máscara suave (soft mask) (pesos float) e multiplica, ou usa estimadores substitutos (fora do escopo aqui, mas veja Retropropagação e Diferenciação Automática).

Atribuição mascarada in-place pode quebrar o autograd (erro comum)

Atualizações in-place em tensores que exigem gradientes podem causar erros se os valores originais forem necessários para o backward:

x = torch.randn(5, requires_grad=True)
mask = x < 0
x[mask] = 0.0  # may error later depending on computation graph

Padrões mais seguros:

Use operações out-of-place:

x = x.masked_fill(mask, 0.0)

Ou faça clone antes da modificação in-place:

x2 = x.clone()
x2[mask] = 0.0

Observe também: operações in-place em tensores folha (leaf tensors) com requires_grad=True são especialmente sensíveis.

Indexação avançada retorna uma cópia (sem semântica de visão)

Um mal-entendido frequente:

x = torch.arange(6).view(2, 3)
mask = x > 2
y = x[mask]
y[:] = 0
print(x)  # x is unchanged

y não é uma visão de x; é um novo tensor. Se você pretende modificar x, deve atribuir em x usando a máscara (x[mask] = ...) ou usar masked_fill_.

Nota sobre determinismo em atualizações por índice

A atribuição geral por indexação avançada (index_put_) pode ser não determinística na GPU quando há índices duplicados, porque atualizações podem entrar em corrida. Máscaras booleanas puras normalmente não criam posições duplicadas, mas se você misturar indexação booleana e inteira de formas complexas (ou construir índices manualmente), tenha isso em mente se você depende de execuções de treino determinísticas.

Boas práticas e dicas práticas

Prefira masked_fill / where quando você quer preservar o formato.
Use indexação booleana (x[mask]) quando você explicitamente quer filtrar e está ok com o formato mudando para (N_true, ...).

Para mascaramento de atenção/logits, use infinito negativo sensível ao dtype:

neg_inf = torch.finfo(logits.dtype).min
logits = logits.masked_fill(mask, neg_inf)

Evite recomputar mask.nonzero() dentro de loops críticos; se a máscara for estável entre passos, compute os índices uma vez e reutilize com index_select/gather.
Lembre: indexação booleana é indexação avançada → aloca e não é uma visão; fatiamento costuma ser mais barato quando aplicável.

Resumo

A indexação por máscara booleana no PyTorch é uma ferramenta poderosa e expressiva que mapeia naturalmente para muitas necessidades em ML (filtrar tokens válidos, aplicar atualizações condicionais, mascarar logits). Os pontos-chave a internalizar são:

Ela é indexação avançada: produz cópias, frequentemente muda formatos e pode ser mais lenta do que fatiamento.
O comportamento de formato depende de a máscara indexar uma dimensão (ex.: máscara de linhas) ou corresponder a elementos (máscara completa → achatamento 1D).
Broadcasting geralmente exige unsqueeze explícito para alinhar dimensões.
Para mascaramento que mantém formato, masked_fill e where geralmente são as ferramentas certas.
O autograd funciona como espalhamento de gradientes de volta para as posições selecionadas, mas escritas mascaradas in-place podem quebrar o cálculo de gradientes se você não tomar cuidado.

Usadas com critério, máscaras booleanas tornam o código PyTorch mais limpo e mais fiel às operações conceituais em modelos modernos de deep learning.