Implementação da esparsificação usando a função Top-K
A operação mais fundamental realizada pelo Q-Sparse é aplicar a função de esparsificação Top-K ao tensor de entrada.
Especificamente, a arquitetura Transformer usa camadas lineares nn.Linear (multiplicação de matrizes) nas camadas de atenção e feed-forward para projeção, que pode ser representada como Y=X·W^T. (Onde X é o tensor de entrada, W representa seus pesos e Y é o tensor de saída)
No Q-Sparse, para um tensor de ativação de entrada X, primeiro é calculado seu valor absoluto |X| e ordenado, para encontrar os K maiores elementos em valor absoluto.
Aqui, K é um hiperparâmetro predefinido que determina o grau de esparsificação.
Em seguida, o Q-Sparse cria um tensor de máscara binária M com a mesma forma de X, definindo as posições correspondentes aos K maiores elementos em |X| como 1 e as demais posições como 0.
Depois, o tensor de entrada X é multiplicado elemento a elemento (produto Hadamard) com o tensor de máscara M para obter o tensor esparso X_sparse.
Durante a propagação direta, o tensor esparso X_sparse substituirá o tensor de entrada original X nos cálculos subsequentes (como multiplicação de matrizes).
Como a maioria dos elementos em X_sparse foi definida como zero, isso pode reduzir significativamente a quantidade de cálculos e a demanda por largura de banda de memória.
Durante a retropropagação, o Q-Sparse usa o Estimador Direto (Straight-Through Estimator, STE) para calcular o gradiente da função Top-K.
No treinamento tradicional, geralmente é necessário calcular o gradiente da função de perda em relação aos parâmetros da rede e usar o método de descida do gradiente para atualizar os parâmetros a fim de minimizar a perda.
No entanto, quando existem operações não diferenciáveis como quantização e Top-K na rede, o cálculo do gradiente encontra problemas, pois o gradiente da saída em relação à entrada é zero na maioria dos pontos para essas operações, impedindo a propagação efetiva do gradiente.
O STE evita o problema do desaparecimento do gradiente passando o gradiente diretamente para o tensor antes da esparsificação.
Na retropropagação normal, o gradiente da função de perda L em relação a x é ∂L/∂x=∂L/∂y⋅∂y/∂x, mas não pode ser calculado diretamente devido à não diferenciabilidade.
A solução do STE é calcular apenas o gradiente da função de perda em relação ao tensor esparso y e então copiá-lo diretamente para o tensor original x, ou seja, usar ∂L/∂y diretamente como uma estimativa de ∂L/∂x.
Para as camadas feed-forward, o Q-Sparse usa a função ReLU ao quadrado em vez da função de ativação ReLU regular, onde a operação de quadrado pode aumentar ainda mais a esparsidade da ativação (⊙ denota o produto Hadamard).
Além disso, para se adaptar a modelos quantizados, o Q-Sparse quantiza o tensor de entrada antes de aplicar a esparsificação Top-K, para garantir que a operação de esparsificação seja compatível com a representação quantizada, cuja função é representada como:
Onde ε é uma pequena constante usada para evitar divisão por zero.
Em particular, para pesos quantizados em 1 bit, o Q-Sparse usa a seguinte função de quantização, onde α é o valor absoluto médio do tensor de peso W.
60% dos parâmetros de ativação alcançam o mesmo efeito
Experimentos comparativos mostram que o Q-Sparse supera significativamente os métodos ReLU anteriores tanto em taxa de esparsidade quanto em desempenho do modelo.
Para os efeitos específicos do Q-Sparse, os autores avaliaram seu desempenho em três tarefas: treinamento do zero, continuação do treinamento e fine-tuning.
O experimento de treinamento do zero usou o modelo Llama, e os resultados mostram que nos modelos de 700M e 7B, o Q-Sparse com 70% top-K (ou seja, 40% de esparsidade geral) pode alcançar uma perda de treinamento comparável ao baseline denso.
O objetivo da continuação do treinamento é esparsificar modelos densos, e o objeto do experimento aqui foi o Mistral-7B.
Como resultado, com parâmetros de ativação de 2,9B e 3,8B, não houve diminuição significativa nas pontuações do modelo em conjuntos de dados como ARC e MMLU.
No experimento de fine-tuning, para os modelos Qwen-7B e Mistral-7B, o Q-Sparse mostrou resultados semelhantes à continuação do treinamento, alcançando desempenho muito próximo ao modelo denso com cerca de 60% dos parâmetros de ativação.
Esses resultados significam que, com o mesmo desempenho, os modelos de ativação esparsa podem reduzir significativamente os parâmetros de ativação durante a inferência em comparação com modelos densos, reduzindo assim o número de FLOPS consumidos.
Para modelos quantizados, a equipe aplicou Q-Sparse ao seu modelo BitNet b1.58 e realizou treinamento e avaliação em vários conjuntos de dados.
Pode-se ver que, em escalas de 700M e 7B, a velocidade de convergência e o valor final da função de perda dos modelos quantizados usando Q-Sparse são comparáveis aos modelos quantizados sem Q-Sparse (BitNet b1.58).
Isso indica que o Q-Sparse pode ser integrado perfeitamente em modelos quantizados sem afetar significativamente o treinamento e a convergência do modelo.
Com base nisso, os autores acreditam que combinar Q-Sparse com técnicas de quantização pode melhorar ainda mais a eficiência dos grandes modelos de linguagem na fase de inferência.
Descoberta de uma nova "Lei de Escala" para otimização de inferência
Além de avaliar o desempenho desses modelos ao adotar ativação esparsa, os autores também exploraram a relação entre desempenho do modelo, escala e taxa de esparsidade, e fizeram algumas novas descobertas.
Lei de escala de desempenho para modelos de ativação esparsa: Os autores descobriram que, semelhante aos modelos densos, o desempenho dos modelos de ativação esparsa também segue uma relação de escala de lei de potência.
Especificamente, dada uma taxa de esparsidade S, o valor da função de perda L(N,S) na convergência do modelo pode ser aproximado pela seguinte fórmula: