A Natureza da Modelagem Matemática e Otimização Probabilística na Engenharia em IA Generativa

The Nature of Mathematical Modeling and Probabilistic Optimization Engineering in Generative AI

October 24, 2024

Autores: Fulu Li

cs.AI

Resumo

Neste artigo, fornecemos uma análise detalhada sobre as formulações de problemas matemáticos e as explorações de otimização probabilística para alguns dos componentes-chave no modelo Transformer [33] no campo da IA generativa. Exploramos e discutimos possíveis melhorias adicionais para os métodos atuais de ponta para algumas tecnologias subjacentes-chave de modelos de IA generativa do ponto de vista da otimização algorítmica e probabilística. Em particular, apresentamos uma solução ótima para codificação de subpalavras (SWE) com base em configurações iniciais semelhantes às do algoritmo de codificação de pares de bytes (BPE) em [9] com objetivos semelhantes aos da abordagem WordPiece em [28, 31] para maximizar a probabilidade dos dados de treinamento. Também apresentamos um método de otimização de entropia cruzada para otimizar hiperparâmetros para o modelo word2vec [17]. Além disso, propomos uma combinação fatorada de codificação posicional rotativa (RoPE) [32] e atenção com viés linear (ALiBi) [23] com uma série harmônica. Também apresentamos um método probabilístico FlashAttention [6, 7] (PrFlashAttention) com uma distribuição de probabilidade sobre distâncias de bloco na matriz para decidir qual bloco é provável de participar em uma rodada de cálculo de atenção, mantendo a forma de triângulo inferior do tensor para modelos de linguagem autoregressivos por meio do remodelamento dos tensores. Por fim, apresentamos a quantização adaptativa em escada (SAQ) do cache chave-valor (KV) para atenção multi-query (MQA) com base no framework apresentado em [16] para ter uma degradação gradual na quantização enquanto se alcança qualidade de modelo razoável e economia de custos.

English

In this paper, we give an in-depth analysis on the mathematical problem formulations and the probabilistic optimization explorations for some of the key components in Transformer model [33] in the field of generative AI. We explore and discuss some potential further enhancement for current state of the art methods for some key underlying technologies of generative AI models from algorithmic and probabilistic optimization perspective. In particular, we present an optimal solution for sub-word encoding (SWE) based on similar initial settings as that of byte-pair encoding (BPE) algorithm in [9] with similar objectives as that of WordPiece approach in [28, 31] to maximize the likelihood of the training data. We also present cross entropy optimization method to optimize hyperparameters for word2vec model [17]. In addition, we propose a factored combination of rotary positional encoding (RoPE) [32] and attention with linear biases (ALiBi) [23] with a harmonic series. We also present a probabilistic FlashAttention [6, 7] (PrFlashAttention) method with a probability distribution over block distances in the matrix to decide which block is likely to participate in a given round of attention computation while maintaining the lower triangle shape of the tensor for autoregressive language models by re-shaping the tensors. Finally, we present staircase adaptive quantization (SAQ) of key-value (KV) cache for multi-query attention (MQA) based on the framework presented in [16] to have gradual quantization degradation while achieving reasonable model quality and cost savings.