""" token_compress — Compressores e indexadores para atenção de 2M tokens. Pipeline (3 estágios): 1. ToMe (Token Merging) — Bolya et al., 2023 Funde tokens similares via matching bipartido soft. Reduz N → N' com N' = N - r·T (r tokens por camada, T camadas). Custo O(N log N). 2. KNN Indexer (IVF-like) Mantém centroides via k-means. Para cada query, busca top-k tokens no índice (k << N'). Custo O(N' log N') para construir, O(k log N') por query. 3. Per-Layer Attention Pattern Cada camada tem seu próprio padrão de sparsidade, selecionável: - "full": dense attention clássica (O(N²)) — apenas para N' pequeno - "window": atenção local em janela w - "global": tokens especiais (CLS, registradores) attendem a tudo - "block": atenção em blocos (BLOCKWISE) O padrão é configurável por camada via lista `patterns`. Complexidade total: O(N log N + T · k · N') comparado a O(T · N²) do Transformer denso. Para N=2M, T=24, k=256: ganho ~ 50×. """ from __future__ import annotations import math from typing import List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # 1. ToMeCompressor # --------------------------------------------------------------------------- def _bipartite_soft_matching( metric: torch.Tensor, r: int, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Matching bipartido soft (Bolya et al., 2023 — algoritmo original). Passos: 1. Ordena tokens por score (norma do embedding — proxy para PCA[0]). 2. Separa alternadamente em dois conjuntos A (índices pares) e B (índices ímpares) — disjuntos, |A| = |B| = N//2. 3. Para cada i ∈ [0, N//2), forma par candidato (a_i, b_i). 4. Calcula similaridade s_i = ⟨a_i, b_i⟩. 5. Seleciona os r pares com maior similaridade para fundir. Os demais pares (N//2 - r) mantêm ambos os tokens. metric: (B, N, C) r: nº de pares a fundir Retorna: a_idx: (B, r) — índices de A nos pares fundidos (âncoras) b_idx: (B, r) — índices de B nos pares fundidos ungated_idx: (B, N - 2r) — índices dos tokens não fundidos """ B, N, C = metric.shape M = N // 2 # número de pares if r <= 0 or M < r: all_idx = torch.arange(N, device=metric.device).unsqueeze(0).expand(B, -1) return all_idx[:, :0], all_idx[:, :0], all_idx with torch.no_grad(): scores = metric.norm(dim=-1) # (B, N) sorted_idx = scores.argsort(dim=-1, descending=True) # (B, N) # Alterna pares/ímpares: A = sorted_idx[:, 0::2], B = sorted_idx[:, 1::2]. a_idx_all = sorted_idx[:, 0:2 * M:2] # (B, M) b_idx_all = sorted_idx[:, 1:2 * M:2] # (B, M) # Similaridade entre cada par (a_i, b_i). a_set = metric.gather(1, a_idx_all.unsqueeze(-1).expand(-1, -1, C)) # (B, M, C) b_set = metric.gather(1, b_idx_all.unsqueeze(-1).expand(-1, -1, C)) # (B, M, C) sim = (a_set * b_set).sum(dim=-1) # (B, M) # Top-r pares mais similares para fundir. top_r = sim.topk(r, dim=-1).indices # (B, r) # Máscara de pares fundidos. merged_mask = torch.zeros(B, M, dtype=torch.bool, device=metric.device) merged_mask.scatter_(1, top_r, True) # Índices de pares não-fundidos (constantes no batch). pair_idx = torch.arange(M, device=metric.device).unsqueeze(0).expand(B, -1) keep_pair = pair_idx[~merged_mask] # (B*(M-r),) keep_pair = keep_pair.reshape(B, M - r) # Token indices. a_merge = a_idx_all.gather(1, top_r) # (B, r) b_merge = b_idx_all.gather(1, top_r) # (B, r) a_keep = a_idx_all.gather(1, keep_pair) # (B, M-r) b_keep = b_idx_all.gather(1, keep_pair) # (B, M-r) # Ungated = a_keep ∪ b_keep (total = 2*(M-r) = N - 2r se N par). ungated = torch.cat([a_keep, b_keep], dim=1) # (B, 2(M-r)) return a_merge, b_merge, ungated class ToMeCompressor(nn.Module): """Token Merging (ToMe) — funde r tokens similares a cada forward. Reduz N → N - r tokens (mantém dimensão de feature). """ def __init__(self, r: int = 64, use_merge_bias: bool = False): super().__init__() self.r = r # Opcional: pequeno gate para mistura ponderada entre par (a, b). # Por padrão usamos média simples (estilo ToMe original). self.merge_gate = nn.Parameter(torch.tensor(0.5)) if use_merge_bias else None def forward(self, X: torch.Tensor, return_aux: bool = False) -> torch.Tensor: """X: (B, N, D) -> (B, N - r, D). Quando return_aux=True, retorna também os índices (para debug). """ B, N, D = X.shape M = N // 2 # Se não há pares suficientes para fundir, retorna X inalterado. if self.r <= 0 or M < self.r: return X a_pair, b_idx, ungated = _bipartite_soft_matching(X, self.r) # Merge: média do par (a, b). a_set = X.gather(1, a_pair.unsqueeze(-1).expand(-1, -1, D)) # (B, r, D) b_set = X.gather(1, b_idx.unsqueeze(-1).expand(-1, -1, D)) # (B, r, D) if self.merge_gate is not None: w = torch.sigmoid(self.merge_gate) merged = w * a_set + (1.0 - w) * b_set else: merged = 0.5 * (a_set + b_set) # Ungated tokens (2*(M-r)). ungated_set = X.gather(1, ungated.unsqueeze(-1).expand(-1, -1, D)) # (B, 2(M-r), D) # Concatena: ungated + merged = 2(M-r) + r = 2M - r = N - r (se N par). out = torch.cat([ungated_set, merged], dim=1) # (B, N - r, D) if return_aux: return out, (a_pair, b_idx, ungated) return out # --------------------------------------------------------------------------- # 2. KNNIndexer (IVF-like, sem dependência do FAISS) # --------------------------------------------------------------------------- class KNNIndexer(nn.Module): """Indexador k-NN baseado em k-means (estilo IVF). Mantém uma lista invertida de tokens por centroide. Query busca apenas nos centroides mais próximos (n_probe) e retorna top-k tokens globais. Custo por query: O(n_probe · avg_list_size + n_centroids). Funciona em batch: para cada query, retorna índices dos top-k tokens e os pesos de similaridade (cosseno). """ def __init__( self, dim: int, n_centroids: int = 256, n_probe: int = 8, k: int = 64, n_iter: int = 5, ): super().__init__() self.dim = dim self.n_centroids = n_centroids self.n_probe = n_probe self.k = k self.n_iter = n_iter # Centroides — buffer treinável em modo "build". self.register_buffer("centroids", torch.randn(n_centroids, dim) * 0.02) # Listas invertidas (chave = centroide idx, valor = índices de tokens). # Em runtime mantemos como listas Python — não persistimos em state_dict. self._lists: List[torch.Tensor] = [] self._tokens: Optional[torch.Tensor] = None # (N, D) self._built = False @torch.no_grad() def build(self, tokens: torch.Tensor) -> None: """Constrói o índice IVF sobre o banco de tokens. tokens: (N, D) """ N, D = tokens.shape assert D == self.dim device = tokens.device # 1. K-means rápido (n_iter iterações de Lloyd). if N <= self.n_centroids: self.centroids = tokens.clone() self._lists = [torch.tensor([i], device=device) for i in range(N)] else: # Inicialização: amostragem aleatória. perm = torch.randperm(N, device=device)[:self.n_centroids] self.centroids = tokens[perm].clone() for _ in range(self.n_iter): # Atribui cada token ao centroide mais próximo. d2 = (tokens.unsqueeze(1) - self.centroids.unsqueeze(0)).pow(2).sum(-1) # (N, K) assign = d2.argmin(dim=-1) # (N,) # Atualiza centroides. for k in range(self.n_centroids): mask = assign == k if mask.any(): self.centroids[k] = tokens[mask].mean(dim=0) # Constói listas invertidas. d2 = (tokens.unsqueeze(1) - self.centroids.unsqueeze(0)).pow(2).sum(-1) assign = d2.argmin(dim=-1) self._lists = [ (assign == k).nonzero(as_tuple=True)[0].to(device) for k in range(self.n_centroids) ] self._tokens = tokens self._built = True @torch.no_grad() def search(self, queries: torch.Tensor, k: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: """Busca top-k tokens para cada query. queries: (B, D) Retorna: indices: (B, k) — índices no banco _tokens weights: (B, k) — similaridade cosseno (em [-1, 1]) """ if not self._built: raise RuntimeError("Índice não construído. Chame build() antes.") k = k or self.k B = queries.shape[0] device = queries.device # 1. Encontra n_probe centroides mais próximos por query. q_norm = F.normalize(queries, dim=-1) # (B, D) c_norm = F.normalize(self.centroids, dim=-1) # (K, D) sim_c = q_norm @ c_norm.t() # (B, K) n_probe = min(self.n_probe, self.n_centroids) top_c = sim_c.topk(n_probe, dim=-1).indices # (B, n_probe) # 2. Para cada query, agrega tokens das listas invertidas correspondentes. all_idx_list = [] all_w_list = [] for b in range(B): cand_indices = torch.cat([self._lists[c.item()] for c in top_c[b] if len(self._lists[c.item()]) > 0]) if cand_indices.numel() == 0: # Fallback: busca global. cand_indices = torch.arange(self._tokens.shape[0], device=device) cand_tokens = self._tokens[cand_indices] # (M, D) q_b = q_norm[b] # (D,) sims = (cand_tokens / cand_tokens.norm(dim=-1, keepdim=True).clamp(min=1e-8)) @ q_b # (M,) kk = min(k, sims.numel()) topk = sims.topk(kk) idx = cand_indices[topk.indices] # (kk,) w = topk.values # (kk,) # Pad com zeros se necessário. if kk < k: pad_idx = torch.zeros(k - kk, dtype=torch.long, device=device) pad_w = torch.zeros(k - kk, device=device) idx = torch.cat([idx, pad_idx]) w = torch.cat([w, pad_w]) all_idx_list.append(idx) all_w_list.append(w) indices = torch.stack(all_idx_list, dim=0) # (B, k) weights = torch.stack(all_w_list, dim=0) # (B, k) return indices, weights # --------------------------------------------------------------------------- # 3. CompressionPipeline (ToMe + KNN index) # --------------------------------------------------------------------------- class CompressionPipeline(nn.Module): """Pipeline de compressão: ToMe → KNN index → retorna subconjunto + pesos.""" def __init__( self, dim: int, tome_r: int = 64, n_centroids: int = 256, n_probe: int = 8, k: int = 64, ): super().__init__() self.tome = ToMeCompressor(r=tome_r) self.indexer = KNNIndexer(dim=dim, n_centroids=n_centroids, n_probe=n_probe, k=k) self._built = False @torch.no_grad() def build(self, tokens: torch.Tensor) -> None: """Constrói índice KNN sobre tokens pré-comprimidos por ToMe.""" compressed = self.tome(tokens) self.indexer.build(compressed) self._built = True def forward( self, queries: torch.Tensor, keys: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Para cada query, retorna top-k keys + pesos. queries: (B, N_q, D) keys: (B, N_k, D) — índice construído para o batch atual. Em modo "global" (índice precompilado via build), apenas queries são processadas e keys são ignoradas. Retorna: indices: (B, N_q, k) — índices das top-k keys weights: (B, N_q, k) — similaridades (cosseno) compressed_keys: (B, N_k - r, D) — keys pós-ToMe (se r > 0) """ # Comprime keys via ToMe. compressed_keys = self.tome(keys) # Para cada batch element, faz busca local (não usa índice global pré-construído). B, Nq, D = queries.shape Nk_c = compressed_keys.shape[1] k = self.indexer.k # Similaridade cosseno entre queries e compressed_keys. q_norm = F.normalize(queries, dim=-1) # (B, Nq, D) k_norm = F.normalize(compressed_keys, dim=-1) # (B, Nk_c, D) sim = q_norm @ k_norm.transpose(-1, -2) # (B, Nq, Nk_c) topk = sim.topk(min(k, Nk_c), dim=-1) # (B, Nq, k) indices = topk.indices # (B, Nq, k) weights = topk.values # (B, Nq, k) return indices, weights, compressed_keys