Mixture of Depths (MoD)

Mixture of Depths (MoD) é uma técnica de Transformer em que um encaminhador aprendido decide, bloco a bloco, quais os tokens processados e quais ignoram o bloco através da sua ligação residual, pelo que tokens diferentes recebem quantidades diferentes de computação. David Raposo, Adam Santoro e colegas da Google DeepMind apresentaram-na em abril de 2024 em Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models. Uma capacidade fixa k determina quantos tokens cada bloco encaminhado processa, pelo que o orçamento computacional e as formas dos tensores são conhecidos antecipadamente, embora a escolha dos tokens dependa da entrada. Os grandes modelos de linguagem (LLMs) MoD igualaram a qualidade de referência com os mesmos FLOPs de treino e tempo decorrido, usando apenas uma fração dos FLOPs por passagem direta, e deram passos a uma velocidade mais de 50% superior durante a amostragem. Com um orçamento de treino fixo, o melhor modelo MoD era maior do que o melhor modelo de referência e atingiu uma perda inferior.

Como funciona o encaminhamento do MoD#

Antes de um bloco encaminhado, um encaminhador linear atribui uma pontuação escalar a cada token. Os k tokens com as pontuações mais altas passam pela autoatenção e pelo MLP do bloco; as respetivas saídas são multiplicadas pelas pontuações do encaminhador e adicionadas novamente ao fluxo residual. A multiplicação pela pontuação coloca o encaminhador no percurso do gradiente, para que a retropropagação lhe possa ensinar quais os tokens que beneficiam. Os restantes tokens passam sem alterações e são excluídos da atenção desse bloco, pelo que os tokens posteriores não lhes podem prestar atenção nesse bloco e o custo da atenção diminui com o quadrado da capacidade: metade dos tokens custa um quarto dos produtos consulta-chave. A melhor configuração do artigo encaminhava um bloco em cada dois, com uma capacidade de 12,5% da sequência. Uma versão mínima de um bloco encaminhado, neste caso apenas com o MLP, em PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Imprime Block ran on 16 of 128 tokens; 112 skipped it unchanged: o MLP do bloco executou um oitavo do trabalho habitual, e o tensor de saída mantém a forma da sequência completa para o bloco seguinte.

Encaminhamento durante a geração#

A seleção top-k não é causal: a classificação de um token entre os k primeiros da sequência depende das pontuações dos tokens que vêm depois dele, os quais ainda não existem quando um modelo gera texto token a token. Os autores tornaram o encaminhamento causal de duas formas: através de uma função de perda auxiliar de entropia cruzada binária, que ensina a saída do encaminhador a prever a inclusão nos top-k e alterou a perda de modelação da linguagem em apenas 0,2–0,3%; ou através de um pequeno preditor MLP separado, com entradas sem gradiente, que atingiu rapidamente uma precisão de 97–99%. Em qualquer dos casos, a amostragem autorregressiva perdeu pouca qualidade. Os encaminhadores treinados também revelam um padrão: alguns tokens ativam todos os blocos encaminhados ao longo da profundidade do modelo, enquanto outros contornam blocos sempre que possível.

MoD e técnicas relacionadas#

• Mistura de especialistas (MoE): Encaminha cada token para um ou alguns MLP especialistas dentro de uma camada, acrescentando parâmetros e mantendo fixa a computação por token. O artigo também combinou as duas técnicas no MoDE, em que um token pode ser encaminhado para um especialista ou para uma operação nula, acumulando os ganhos.
• Saída antecipada: Um token termina após uma camada e ignora todas as restantes. O MoD pode ignorar um bloco e continuar a utilizar os seguintes.
• Fusão de tokens (ToMe): Combina tokens semelhantes num Transformer de visão, para que as camadas seguintes processem menos tokens, sem recorrer a um encaminhador aprendido.
• Poda de modelos: Remove pesos permanentemente, para que todas as entradas recebam a mesma rede, mais pequena.
• Encaminhamento de modelos: Envia um pedido completo para um de vários modelos; a mesma ideia aplicada ao nível de modelos completos.

A maioria das redes tem uma profundidade fixa quando é construída. O Ultralytics YOLO26 define-a por escala do modelo: o multiplicador de profundidade no seu YAML do modelo é 0.50 para as escalas n, s e m e 1.00 para l e x; por isso, uma etapa C3k2 indicada com 2 repetições executa 1 bloco em YOLO26n e 2 em YOLO26x, e todas as imagens passam por todos os blocos.