[Tradução] MoD (Mixture-of-Depths): uma abordagem para otimizar a computação em modelos de linguagem baseados em Transformer, e MoDE (MoD+MoE)

(discuss.pytorch.kr)

7 pontos por ninebow 2024-04-13 | Ainda não há comentários. | Compartilhar no WhatsApp

PyTorchKR

Métodos de MoE (Mixture-of-Experts), como Jamba, Qwen1.5-MoE e DBRX, que processam tokens de entrada selecionando um ou mais blocos de “especialistas”, têm sido muito usados recentemente. Essa abordagem MoE mostra melhor desempenho ao alocar recursos computacionais de forma eficiente por meio do controle dinâmico de partes de LLMs muito grandes.
Se o MoE ajusta dinamicamente a largura (width) do modelo ao escolher apenas alguns especialistas, o MoD (Mixture-of-Depths), apresentado aqui, adota uma abordagem para ajustar dinamicamente a profundidade (depth) do modelo. Ou seja, para determinados tokens, em vez de executar computação em todas as camadas do modelo, ele realiza computação dinâmica pulando algumas camadas. Vamos ver isso juntos.

MoD (Mixture-of-Depths): uma abordagem para otimizar a computação em modelos de linguagem baseados em Transformer (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Introdução

Tradicionalmente, modelos Transformer alocam os mesmos recursos computacionais para todas as partes da sequência de entrada. No entanto, o artigo sobre o Transformer MoD (Mixture-of-Depths) apresentado aqui tenta alocar recursos de computação dinamicamente para posições específicas da sequência no Transformer, variando essa alocação ao longo das diferentes camadas. Para usar essa estratégia de Mixture-of-Depths (MoD), os autores empregam um mecanismo de roteamento top-k para decidir quais tokens serão processados e, assim, controlar o volume total de computação. Essa técnica exige muito menos FLOPs, mantendo desempenho comparável ou superior ao dos modelos existentes, e oferece aceleração de até 50% durante a amostragem após o treinamento.

Arquitetura Transformer apresentada no artigo Attention is All You Need

Os modelos Transformer e LLMs (Large Language Models) mais recentes continuam crescendo em escala e complexidade. Esses modelos utilizam bilhões de parâmetros para se aproximar de capacidades humanas de compreensão e geração de linguagem, e como resultado apresentam desempenho impressionante. No entanto, esse avanço vem acompanhado de custos computacionais e consumo de energia significativos, e o aumento da escala do modelo eleva drasticamente o tempo de treinamento e inferência.

Isso acontece porque o Transformer executa a mesma quantidade de computação para todos os tokens da sequência de entrada. Mas a ideia de partida é que nem todos os tokens fornecem a mesma quantidade de informação: alguns desempenham papel central na compreensão do contexto, enquanto outros são relativamente menos importantes.

O Transformer MoD permite que o modelo concentre recursos onde for necessário ao alocar computação dinamicamente para cada token. Isso reduz o custo computacional total do modelo e, ao mesmo tempo, aumenta a eficiência do processamento de informação ao dedicar mais atenção aos tokens importantes. Especialmente em modelos de grande escala como LLMs, a abordagem do MoD pode se tornar uma forma importante de economizar tempo e recursos durante treinamento e inferência.

Introdução à técnica MoD (Mixture-of-Depths)

A técnica MoD integra um mecanismo de roteamento top-k que seleciona dinamicamente quais tokens serão processados em cada camada do Transformer. Com isso, apenas uma parte dos tokens selecionados passa pelo processo completo de computação, enquanto os demais pulam uma ou mais camadas.

Figura 1 / Esquerda: visão geral do Transformer MoD (Mixture-of-Depths): de forma semelhante à arquitetura MoE (Mixture-of-Experts), o caminho de computação é escolhido por meio de um roteador. Mas, diferentemente do MoE, aqui é decidido se o bloco padrão (Self-Attention e MLP) será executado ou não. Se o token de entrada ($X_i$) não ativar a computação completa de acordo com a decisão do roteador, ele não passa pela computação daquela camada.
Figura 1 / Direita: resultado das decisões de roteamento para uma sequência curta reduzida a 64 tokens. (Eixo X: sequência, Eixo Y: camada) Roxo indica que a computação daquela camada foi executada por completo, e laranja indica que ela foi pulada.

Como descrito acima, o MoD implementa um roteador preditivo dentro da arquitetura Transformer para decidir, para cada token, se a computação da camada será executada ou não. Isso permite concentrar recursos computacionais em tokens relativamente “importantes”.

Treinamento do Transformer MoD

O processo de treinamento do Transformer MoD é basicamente semelhante ao de um Transformer padrão. A sequência de entrada é fornecida ao modelo e os cálculos necessários são executados à medida que ela passa por cada camada. Porém, o Transformer MoD inclui uma etapa adicional em que o roteador decide qual caminho computacional cada token seguirá.

Cálculo dos pesos do roteador: o roteador calcula, em cada camada do modelo, um peso para cada token. Esse peso representa quanto de computação o token deve receber e se baseia em sua importância.
Seleção dos tokens do topo: com base nos pesos calculados, o roteador seleciona, para cada camada, os tokens no topo. Esses tokens são encaminhados para o caminho que executa a computação completa.
Atribuição do caminho de computação: os tokens selecionados são atribuídos ao caminho que executa a computação padrão (Self-Attention e MLP). Os demais tokens seguem pela conexão residual e passam para a próxima camada sem computação adicional.
Cálculo da perda e retropropagação: a saída do modelo é avaliada pela função de perda final, e os gradientes são propagados de volta pelo modelo para atualizar os pesos. Nesse processo, o roteador também é treinado, podendo tomar decisões de roteamento mais eficientes com o tempo.

Esquemas de roteamento do Transformer MoD (Routing Schemes)

A implementação do roteamento é uma parte central do Transformer MoD. Os autores compararam três esquemas de roteamento e confirmaram que o roteamento de tokens no esquema Expert-Choice MoD apresenta o melhor desempenho. Vamos compará-lo com os outros esquemas apresentados:

Roteamento por escolha do token (Token-choice)

No roteamento por escolha do token, cada token escolhe por conta própria o caminho pelo qual vai passar. O modelo gera, para cada token, uma distribuição de probabilidade sobre vários caminhos computacionais possíveis (por exemplo, diferentes camadas ou blocos de processamento) e, com base nessa distribuição, envia o token para o caminho correspondente.

Esse tipo de roteamento permite grande flexibilidade, já que cada token pode escolher o caminho computacional mais adequado. Além disso, o modelo pode considerar com mais precisão a importância ou o contexto de cada token, permitindo otimização em nível individual.

Roteamento por escolha do especialista (Expert-choice)

No roteamento por escolha do especialista, o modelo escolhe diretamente quais tokens cada caminho (ou “especialista”) irá processar. Isso é feito usando um mecanismo top- que seleciona os tokens do topo com base nos pesos do roteador de cada token.

Com esse roteamento por escolha do especialista, os tokens podem ser distribuídos de maneira mais equilibrada entre os caminhos, o que espalha melhor a carga computacional e permite usar os recursos do modelo com mais eficiência. Além disso, como os tokens do topo são definidos de forma clara, a previsibilidade e a estabilidade do modelo melhoram.

Expert-choice MoD

O Expert-choice MoD aplica o conceito de roteamento por escolha do especialista à estrutura do Transformer MoD. Nesse método, o modelo seleciona os tokens do topo para cada camada ou etapa de processamento, e apenas os tokens escolhidos seguem pelo caminho que realiza computação. Os demais seguem um caminho de menor custo por meio das conexões residuais.

Como apenas os tokens selecionados executam computação, o uso total de FLOPs do modelo cai bastante. Além disso, para tokens importantes, a computação em toda a camada é garantida, enquanto computação desnecessária é minimizada, preservando ou até melhorando o desempenho do modelo. Por fim, por meio do roteamento dinâmico via roteador, o modelo pode alocar recursos computacionais de forma flexível conforme as características da sequência de entrada.

Implementação do roteamento

A implementação do roteamento pode ser feita por meio das três etapas principais a seguir:

Embeddings de tokens e definição dos pesos do roteador: para cada camada do modelo, existe um conjunto de embeddings de tokens $X^i$ da sequência dada. Para cada embedding de token $x^l_i$, o roteador calcula, por meio de uma transformação linear, um peso escalar $r^l_i$ que representa a probabilidade de esse token participar da computação.
Determinação dos pesos do topo: no conjunto de pesos do roteador dos tokens $R^l$, encontra-se a probabilidade percentil do peso de ordem $\beta$, denotada por $P_\beta(R^l)$. Aqui, $\beta$ determina a proporção de tokens que o modelo processará em cada camada. Assim, os tokens com peso do roteador maior que $P_\beta(R^l)$ participam da computação.
Seleção do caminho de computação: cada token é comparado com o peso do roteador $r^l_i$ e, se for maior ou igual a $P_\beta(R^l)$, ele executa a computação daquela camada; se for menor, é enviado diretamente para a próxima camada. Com isso, a computação é realizada apenas para os tokens selecionados, enquanto os demais seguem por um caminho que não gera custo computacional.

Na implementação desse roteamento, é preciso considerar principalmente os seguintes pontos:

Decisão dinâmica baseada em pesos: o processo de roteamento é decidido dinamicamente com base no peso de cada token, refletindo sua importância. Tokens importantes recebem mais computação para otimizar o desempenho do modelo, enquanto os menos importantes reduzem custo ao seguir pelo caminho residual.
Eficiência computacional e otimização de desempenho: por meio desse mecanismo de roteamento, o Transformer MoD pode usar os recursos computacionais com eficiência e, ainda assim, manter ou melhorar o desempenho do modelo. Em especial, ao processar apenas os tokens importantes em camadas de alto custo computacional, é possível reduzir o uso total de FLOPs.
Simplicidade e eficiência da implementação: como o roteamento é implementado com transformação linear e cálculo de percentil, ele é relativamente simples e eficiente. Isso torna possível a alocação dinâmica de recursos computacionais sem adicionar carga computacional relevante ao treinamento e à inferência.

Exemplo de funcionamento do Transformer MoD

Com base no que vimos até agora, vamos explicar o funcionamento do Transformer MoD com um exemplo.

Suponha, por exemplo, que a sequência de entrada tenha 100 tokens e que o roteador tenha definido como 10. Nesse caso, em cada camada o roteador seleciona os 10 tokens mais importantes, e apenas esses tokens são processados pelo caminho de computação padrão.

Os 90 tokens restantes pulam a computação e são enviados diretamente para a próxima camada por meio das conexões residuais. Com isso, o Transformer MoD reduz o uso total de FLOPs enquanto garante a computação necessária para os tokens importantes.

Comparação de desempenho do Transformer MoD

Desempenho no treinamento

Otimização de hiperparâmetros: por meio de experimentos que treinaram uma variação do Transformer MoD com blocos de roteamento de capacidade de 12,5% posicionados alternadamente com blocos completos de attention, foi encontrada a configuração ideal de hiperparâmetros. Essa configuração apresentou desempenho superior ao melhor baseline na análise isoFLOP.

Eficiência computacional: uma determinada variação de MoD mostrou velocidade de processamento 66% maior, ao mesmo tempo em que alcançou o mesmo desempenho que modelos existentes com menos FLOPs.

Análise de roteamento

Esparsidade das decisões de roteamento: com roteamento de capacidade de 12,5%, a maioria dos tokens contornou os blocos e apenas uma pequena parcela de tokens importantes passou por eles. Isso também pôde ser observado na distribuição dos pesos do roteador, e a rede chegou a preferir certos tokens para roteamento dependendo da profundidade.

Desempenho na inferência

Inferência autorregressiva: ao trocar, na inferência, o esquema de roteamento top- não causal (non-causal) usado no treinamento por uma abordagem causal (causal) baseada em predição, quase não houve perda de desempenho. Isso parece ocorrer porque esse problema de predição é fácil de aprender, alcançando mais de 97% de precisão desde o início do treinamento.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) significa a integração da técnica MoD (Mixture-of-Depths) com a técnica MoE (Mixture of Experts). O objetivo dessa abordagem é combinar os pontos fortes das duas estratégias para otimizar ao mesmo tempo a eficiência computacional e o desempenho de modelos Transformer. Enquanto o MoD escolhe realizar computação mais profunda ou menos profunda para cada token, o MoE processa o token por meio do especialista mais adequado entre vários “especialistas”. O MoDE integra essas duas abordagens para permitir que um token pule determinados blocos ou execute seletivamente computação mais profunda e, ao mesmo tempo, escolha a forma de processamento ideal entre vários especialistas.

Como o MoDE é implementado

O MoDE possui duas formas principais de implementação: MoDE em estágios (Staged) e MoDE integrado (Integrated).

Staged MoDE (MoDE em estágios): o MoDE em estágios aplica primeiro a técnica MoD para definir o caminho computacional de cada token e, em seguida, aplica a técnica MoE para rotear, para cada bloco especialista, os tokens que serão processados dentro do caminho selecionado. Dessa forma, o modelo ajusta primeiro a profundidade da computação e depois escolhe o especialista ideal em cada etapa.
Integrated MoDE (MoDE integrado): o MoDE integrado combina as técnicas MoD e MoE em uma única etapa de roteamento. Ou seja, ao rotear um token, uma única decisão escolhe ao mesmo tempo a profundidade do caminho computacional e o especialista que irá processá-lo. Esse método pode ser mais simples e eficiente de implementar, além de ter a vantagem de decidir simultaneamente o caminho e o especialista mais adequados para cada token.

Vantagens do MoDE

Eficiência computacional: o MoDE gerencia de forma eficaz o custo computacional total do modelo ao decidir dinamicamente, para cada token, a quantidade de computação necessária e o especialista que irá processá-lo.
Otimização de desempenho: ao atribuir cada token ao caminho computacional e ao especialista mais adequados, o MoDE maximiza o desempenho do modelo. Isso pode ser especialmente útil ao resolver problemas complexos ou processar diferentes tipos de dados.
Flexibilidade e escalabilidade: o MoDE oferece flexibilidade para ser aplicado a várias arquiteturas de modelo e tarefas, além de fornecer um framework escalável capaz de melhorar simultaneamente eficiência computacional e desempenho.

O MoDE aponta uma direção para a evolução futura dos modelos Transformer e contribui de forma importante para pesquisas que buscam maximizar o desempenho do modelo enquanto minimizam o custo computacional.