arXiv:2410.18441cs.LGcs.AI2024-10

从数学建模与概率优化角度,提升生成式AI核心组件的性能与效率。

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

  • 基于相似目标优化子词编码,提升训练数据似然度。
  • 提出概率化FlashAttention与阶梯自适应量化,兼顾速度与精度。
  • 适合关注模型优化、算法设计的科研与工程人员阅读。

本文深入分析了生成式AI中Transformer模型关键组件的数学问题建模与概率优化方法。针对子词编码(SWE),在与字节对编码(BPE)类似设定下,提出最大化训练数据似然的优化解;针对word2vec模型,提出交叉熵优化方法以调整超参数。进一步提出将旋转位置编码(RoPE)与线性偏置注意力(ALiBi)结合,并引入调和级数实现因子化组合。还提出概率化FlashAttention(PrFlashAttention),通过块距离的概率分布选择参与注意力计算的块,同时保持自回归模型张量的下三角结构。最后,基于[16]框架,提出阶梯自适应量化(SAQ)方法对多查询注意力(MQA)的键值缓存进行渐进式量化,在保证合理模型质量的同时实现成本节约。

原文摘要 · Abstract (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.

生成式AI数学建模概率优化注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。