Mistura de profundidade: uma técnica para alocar recursos de computação dinamicamente em transformers

• O Mixture-of-Depths (MoD) do Google DeepMind faz com que modelos de linguagem baseados em transformers não usem a mesma quantidade de FLOPs para todos os tokens; em vez disso, em cada camada, apenas os tokens necessários participam dos cálculos de attention e MLP
• Um roteador por camada gera um peso escalar para cada token, e apenas os tokens dentro da capacidade top-k predefinida passam pelo bloco; os demais desviam pela residual connection
• Ao fixar k previamente, o grafo de computação e os tamanhos dos tensores permanecem estáticos, mas a distribuição de computação por token muda dinamicamente conforme o contexto
• Nos experimentos, o MoD conseguiu, em base isoFLOP, obter loss menor que o transformer básico no mesmo tempo de treinamento ou, com o mesmo desempenho, reduzir os FLOPs por forward pass para acelerar as etapas de treinamento e amostragem
• A melhor configuração aplicava blocos com capacidade de 12,5% em blocos alternados, e na amostragem auto-regressiva usava roteamento baseado em predictor para reduzir o problema de não causalidade do top-k

---

Desperdício de computação que o MoD pretende atacar

• Um modelo de linguagem transformer comum usa a mesma quantidade de computação para todos os tokens no forward pass
• O MoD parte da premissa de que nem todos os tokens e sequências exigem o mesmo tempo e esforço de previsão, distribuindo a computação para os tokens necessários enquanto reduz o orçamento computacional total
• Métodos anteriores de computação condicional podem criar grafos de computação dinâmicos, o que pode não se alinhar ao grafo de computação estático preferido pelo hardware atual
• Se o usuário definir o orçamento computacional total antes do treinamento, o modelo aprende a decidir onde gastar essa computação por token e por camada dentro desse limite
  • É possível prever com antecedência a quantidade de computação, o uso de memória e os FLOPs por forward pass
  • Os tokens escolhidos para roteamento mudam de acordo com o contexto de entrada

Estrutura de roteamento e forma de reduzir computação

• O MoD usa um roteador como um transformer MoE, mas em vez de escolher entre vários experts, envia o token para um de dois caminhos
  • Cálculo de bloco padrão com attention e MLP
  • Residual connection que apenas repassa o valor
• A residual connection tem baixo custo computacional, e a saída do bloco é determinada pelo valor de entrada
• Se a capacidade do bloco for configurada para um valor menor que o número total de tokens T, os FLOPs por forward pass caem em relação a um transformer básico
  • Por exemplo, se a capacidade for reduzida para T/2, a multiplicação de matrizes query-key da attention passa de T² para (T/2)², reduzindo os FLOPs para cerca de 25%
• Se a capacidade for 0, o modelo fica muito rápido, mas o desempenho piora por não usar a maior parte dos parâmetros do transformer; se a capacidade for T, ele se torna igual ao transformer básico
• O objetivo é encontrar, entre esses dois extremos, um ponto em que o modelo mantenha desempenho igual ou melhor que o transformer básico enquanto aumenta a velocidade por step

Roteamento top-k do tipo expert-choice

• Em cada camada l, o roteador gera um peso escalar para cada embedding de token
• De acordo com a capacidade C definida pelo usuário, apenas os tokens do top-k da sequência participam dos cálculos de attention e MLP
• Os tokens não selecionados pulam o cálculo do bloco e seguem para a próxima camada pela residual connection
• Esse método se baseia em expert-choice routing
  • No token-choice routing, o token escolhe o caminho, o que pode causar problemas de balanceamento de carga
  • No expert-choice routing, cada caminho escolhe seus tokens top-k, preenchendo exatamente a capacidade
• No MoD, os caminhos são na prática apenas dois — “cálculo do bloco” e “desvio” — então uma única operação top-k basta para dividir os tokens em dois grupos
• Os pesos do roteador são multiplicados pela saída do bloco, ficando no caminho do gradiente e sendo afetados pelo gradiente descendente durante o treinamento de modelagem de linguagem

O problema de não causalidade na amostragem

• A operação top-k do expert-choice routing é não causal (non-causal)
  • Decidir se um token entra no top-k exige conhecer também os pesos do roteador dos tokens seguintes
  • Na amostragem auto-regressiva, não é possível usar informações de tokens futuros
• Para evitar isso, foram testados dois métodos
  • Uma loss auxiliar de binary cross-entropy para que o sigmoid da saída do roteador previsse se o token seria escolhido pelo top-k
  • Um pequeno predictor MLP auxiliar, com a mesma entrada do roteador, mas com stop-gradient aplicado
• O primeiro método afetou o objetivo principal de modelagem de linguagem em cerca de 0,2~0,3%
• O segundo método, baseado em predictor, não afetou o objetivo de modelagem de linguagem e também não teve impacto significativo na velocidade por step
• A tarefa auxiliar de predição alcançou experimentalmente 99% de acurácia rapidamente

Resultados experimentais e observações

• Os experimentos iniciais buscaram hiperparâmetros do MoD com um orçamento total de 6e18 FLOPs
• A melhor configuração do MoD aplicava roteamento em blocos alternados, com top-k igual a 256
  • Com comprimento de sequência 2048, isso significa 256 tokens, ou 12,5%, passando por attention e MLP
  • Os outros 1792 tokens, ou 87,5%, desviam do bloco
• A variante MoD de 220M parâmetros teve desempenho ligeiramente melhor que o baseline ótimo em isoFLOP e foi mais de 60% mais rápida por step durante o treinamento, enquanto o wall-clock time total de treinamento no mesmo hardware foi quase igual
• Nas análises isoFLOP de 6e18, 2e19 e 1e20 FLOPs, os modelos MoD ótimos em FLOPs tendiam a ter mais parâmetros que o baseline
• Existem variantes MoD com o mesmo desempenho ou loss menor usando menos FLOPs por forward pass
• O ganho de velocidade por step vem de dois fatores
  • Como parte dos tokens desvia do bloco, os FLOPs por parâmetro caem para o mesmo tamanho de modelo
  • É possível escolher variantes MoD menores que atingem o mesmo desempenho do baseline
• O MoD com stochastic routing mostrou desempenho bem inferior ao transformer básico e ao MoD comum, indicando que o roteamento aprendido era importante
• Na avaliação auto-regressiva, foram usados 256.000 sequências e 500M tokens de held-out data, e a troca para roteamento baseado em predictor causou apenas pequena queda de desempenho
• O MoD também pode ser combinado com MoE para formar o Mixture-of-Depths-and-Experts (MoDE)
  • O staged MoDE decide antes da attention se o token vai para o bloco ou desvia
  • O integrated MoDE insere um expert “no-op” entre os experts MLP existentes para integrar o roteamento do MoD
• Na análise de roteamento, observou-se um padrão em que alguns tokens passam com frequência por vários blocos, enquanto outros desviam sempre que possível
• Uma análise preliminar sugeriu que os tokens que passam mais vezes pelos blocos se correlacionam com maior entropy na predição de saída, o que pode indicar que são tokens mais difíceis de prever

[GN⁺]

Opiniões no Hacker News

• Roteamento mais complexo certamente deve se tornar mais comum
  Em especial, acho que algum dia iremos para um roteamento recursivo, em que se passa novamente por um conjunto de especialistas. Daqui para a frente, o “raciocínio em cadeia” (chain-of-thought) deve acontecer de forma recursiva dentro do modelo

  • Esses objetos hipotéticos poderiam ser chamados de Recursive Neural Networks
  • O que foi descrito aqui parece um pouco com a linha de pesquisa dos Universal Transformers. A ideia é passar os embeddings de entrada várias vezes por um único bloco Transformer, enquanto um módulo separado decide se o embedding já “amadureceu” o suficiente para ser retirado
    Mais próximo da ideia de “especialistas”, há o artigo Sparse Universal Transformers do ano passado, que combina Universal Transformer com mistura esparsa de especialistas, fazendo um mecanismo de gating decidir quais blocos Transformer usar e em que ordem
    Não é minha área de especialidade, mas entendi que é difícil de treinar corretamente e que, para obter resultados parecidos com os de um Transformer comum, exige mais computação total na inferência. Ainda assim, é uma direção interessante, e vejo o fato de haver um limite superior no número de etapas de computação por token como uma das grandes desvantagens da arquitetura Transformer clássica
  • Acho que o motivo de isso ainda não ter acontecido é que, no momento do treinamento, não há como decidir quantas vezes a recursão deve ser feita
    Se você escolhe um número aleatório de vezes ou tenta várias profundidades de recursão, a saída fica “embaçada”. Ou seja, deixa de ficar claro se a saída de uma camada deve produzir informações importantes para o resultado final ou a melhor entrada para a próxima rodada de recursão
  • Atenção é basicamente roteamento, e essas outras formas de roteamento podem dar ao modelo opções mais grosseiras, potencialmente facilitando o aprendizado
  • A tendência certamente é de aumento do roteamento dinâmico, mas vejo MoE/MoD/MoDE mais como algo que permite armazenar fatos adicionais nos pesos, reduzindo sobreposição, do que como algo que possibilita raciocínio mais profundo
    Acho que o raciocínio mais profundo virá de dinamicidade no nível dos tokens, não no nível das camadas. Por exemplo, há o artigo recente Quiet-STaR, que faz o modelo emitir tokens de justificativa que depois descarta: https://arxiv.org/abs/2403.09629

• Pode ser o artigo mais importante de 2024
  A ideia de querer modelos que não gastem a mesma quantidade de computação em todos os tokens existe há tempos, mas é a primeira vez que vejo um mecanismo convincente para isso

  Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
  Isso não é meio surpreendente?

  • Sparse Universal Transformer é mais antigo e já fazia encerramento antecipado baseado em roteamento
  • Mais importante a esse ponto? O fato de que nem todo token precisa de toda a janela de contexto deveria ser uma otimização óbvia

• Uma versão introdutória simples:
  Imagine que você tem um assistente inteligente capaz de entender e processar fala. Normalmente, esse assistente presta a mesma atenção a todas as palavras, independentemente de quão importante cada uma seja para o significado geral
  Agora suponha que encontramos uma forma de ensinar o assistente a usar seus “recursos cerebrais” de maneira mais inteligente. Em vez de dar a mesma atenção a todas as palavras, ele se concentra mais nas palavras mais importantes para entender o significado e ajusta esse foco na hora, conforme o contexto
  Para evitar sobrecarregar o assistente, também colocamos um limite no total de “recursos cerebrais” que ele pode usar em determinado momento. É como dar um orçamento e dizer: “você só pode gastar recursos em um certo número de palavras por vez”. Então o assistente precisa decidir quais palavras são mais importantes
  Mesmo com esse limite, o assistente ajusta o uso de recursos de forma flexível. Ele gasta mais em algumas palavras e menos em outras, priorizando de acordo com a situação
  Ao aprender assim, o assistente funciona de modo eficiente enquanto presta atenção de forma inteligente. Ele entende tão bem quanto um assistente que dá a mesma atenção a todas as palavras, mas usa menos recursos cerebrais no total, então consegue responder e processar novas informações muito mais rápido

  • Sei que é um ELI5, mas o método descrito não é justamente o que a atenção já faz? Ela se concentra especificamente nas palavras contextualmente importantes da sequência anterior

• Com base no que entendi, fiz um pequeno resumo aqui:
  https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

  • O texto é bom. Lembra o estilo da New Scientist. Tem aquela sensação de “explicar à medida que avança, mas mantendo curto”, boa para pegar o jeito desta área

• É muito parecido com mistura de especialistas. Só que, em vez de rotear tokens para vários especialistas, é como “colocá-los em um único especialista que pode ser pulado dinamicamente”

  • Misturar os dois seria bem legal. Poderia reduzir ainda mais a computação de MoE mantendo o desempenho

• “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
  É esse o caminho que possibilita uma janela de contexto de 10 milhões de tokens? Ou isso se refere a janelas de contexto mais longas no futuro?

• Depois de tentar entender e implementar alguns algoritmos RASP, a conclusão que tirei foi que certas funções precisam de um certo número de camadas Transformer para funcionar
  Seguindo essa lógica, fica claro que as funções aprendidas por Transformers podem se espalhar por várias cabeças. Repetir essas funções pode ser muito valioso para entender e resolver problemas, mas o modo atual de inferência não consegue executar repetidamente conjuntos consecutivos de cabeças. Este artigo realmente parece apontar para uma direção promissora
  [1] https://arxiv.org/pdf/2106.06981.pdf
  [2] https://www.youtube.com/watch?v=t5LjgczaS80

• Talvez a única desvantagem de os LLMs avançarem rápido demais seja que os artigos saem mais rápido do que qualquer pessoa — pelo menos fora do Google — consegue aprender e testar as melhorias
  Quando comecei com deep learning, ReLU e dropout estavam em alta, e dava para mudar uma ou duas linhas de código em uma 1080 de consumidor e testar em poucas horas se havia melhora. Agora parece que temos de esperar algumas semanas até algum lugar como a Mistral experimentar

  • Bem-vindo à classe pobre em GPUs
    Eu estou focando em abordagens de quantização e testando com GPUs antigas de gerações anteriores

• hu-po fez uma análise em livestream aprofundada de artigos de IA
  Recomendo muito, e aqui ele aborda o artigo mixture-of-depths em discussão: https://www.youtube.com/watch?v=Teru_qIdB8Y