SpatialLM-Mini acelerou o treinamento com contexto de 16K em 2,47 vezes em uma RTX 3050

O contexto longo custa caro antes mesmo de o modelo começar a gerar valor. Com atenção completa, cada novo token pode prestar atenção a todos os anteriores, e o volume de cálculos cresce de forma quadrática. No experimento SpatialLM-Mini, esse acesso foi limitado às posições próximas e a algumas conexões de longo alcance. Em uma RTX 3050 com 8 GB de memória, o protótipo foi treinado 2,47 vezes mais rápido que o Transformer de referência com uma janela de 16 384 tokens.

Para quem desenvolve modelos, há uma indicação concreta: reduzir o número de conexões pode acelerar experimentos com sequências longas mesmo em uma placa de vídeo barata. Mas o resultado se refere a um modelo pequeno, um único corpus e um kernel de cálculo específico. O benefício do método propriamente “espacial” de calcular a atenção ainda é uma questão em aberto.

Como os tokens escolhem quem ouvir

Imagine um documento longo. Na atenção causal convencional, um token no final pode considerar diretamente qualquer posição anterior. No SpatialLM-Mini, a lista de posições acessíveis é definida de antemão: o próprio token, alguns predecessores próximos e posições a distâncias de 1, 2, 4, 8 e assim por diante. O acesso a posições futuras é proibido.

Em seguida, o modelo decide quanta atenção dedicar a cada vizinho permitido. Para isso, o estado oculto é transformado em coordenadas tridimensionais aprendíveis de consulta e chave. Quanto menor a distância entre elas, maior pode ser o peso da conexão; parâmetros aprendíveis adicionais ajustam esse cálculo.

As coordenadas servem aqui para calcular os pesos. O modelo não procura os tokens mais próximos em todo o documento nem utiliza os núcleos RT da placa de vídeo. A lista de vizinhos é fixa, enquanto a distribuição da atenção dentro dela depende do texto.

Esse grafo reduz os cálculos, mas limita o acesso direto a fatos distantes. Por isso, o autor também testou um híbrido: dois blocos espaciais se alternam com dois blocos de atenção completa. Isso permite que parte das camadas acesse todo o contexto anterior.

O que mostrou a comparação com 8 GB

Três modelos com aproximadamente 9,3 milhões de parâmetros cada foram treinados do zero com Orgulho e Preconceito, de Jane Austen. Cada um tinha quatro camadas, o mesmo comprimento de contexto e três execuções com sementes aleatórias diferentes. Do GPT-2, foi usado apenas o tokenizador.

Com um orçamento de 180 segundos de treinamento medido, o Transformer de referência processou 13 126 tokens por segundo; o modelo espacial, 32 364; e o híbrido, 18 725. O híbrido apresentou a menor perda média no conjunto de teste. No entanto, os modelos mais rápidos conseguiram processar mais tokens e realizar mais atualizações nesse intervalo.

Na comparação após os mesmos 100 passos, a diferença na perda diminuiu consideravelmente. O modelo espacial precisou de 50,64 segundos; o Transformer, de 125,10 segundos. O principal ganho comprovado nesse experimento é a velocidade das atualizações com uma sequência longa.

Com a memória, o resultado foi diferente: o pico de memória alocada pelo PyTorch subiu de aproximadamente 499 para 671 MiB. A implementação esparsa cria tensores intermediários ao reunir os vizinhos, enquanto a atenção completa utiliza um kernel otimizado. Esses números não representam toda a memória de vídeo ocupada no computador.

O que verificar antes de adotar no seu projeto

O Transformer de referência utilizou o PyTorch Efficient Attention. O FlashAttention não estava disponível na compilação para Windows usada no experimento. Por isso, o fator de aceleração não pode ser extrapolado automaticamente para outro sistema nem para comparações com outros kernels.

Há também uma lacuna mais relevante: o conjunto de conexões e o método de cálculo dos pesos foram alterados ao mesmo tempo. Para entender a contribuição das coordenadas, é necessário usar o mesmo grafo esparso com um produto escalar convencional ou com pesos fixos. Por enquanto, a vantagem pode ser explicada principalmente pela redução do número de interações.

Aceitar 16 mil tokens também não demonstra que o modelo consegue recuperar o fato relevante no início do documento. Isso exige um teste separado. A velocidade de geração não foi medida, o protótipo não tem cache KV e um treinamento curto com um único romance não permite avaliar sua adequação a tarefas de produção.

O próximo passo prático é repetir a comparação com seu próprio corpus e hardware e, depois, acrescentar uma referência com o mesmo grafo, sem coordenadas. Segundo a descrição do autor, o código, as configurações e as medições originais estão disponíveis no repositório SpatialLM-Mini. Para uma equipe que pesquisa treinamento de baixo custo com contextos longos, esse é um ponto de partida útil para um experimento próprio.