#矩阵与图竟然存在等价关系

矩阵很难理解,但换个视角或许会不一样。

在学习数学时,我们常因所学知识的难度和抽象而受挫;但有些时候,只需换个角度,我们就能为问题的解答找到一个简单又直观的解法。举个例子,小时候在学习和的平方 (a+b)² 公式时,我们可能并不理解为什么它等于 a²+2ab+b²,只知道书上这么写,老师让这么记;直到某天我们看见了这张动图:

51c大模型~合集37_大模型

登时恍然大悟,原来我们可以从几何角度来理解它!

现在,这种恍然大悟之感又出现了:非负矩阵可以等价地转换成对应的有向图!

如下图所示,左侧的 3×3 矩阵其实可以等价地表示成右侧的包含三个节点的有向图,并且这种表示方式对矩阵和图论都大有帮助。

51c大模型~合集37_大模型_02

这个例子来自致力于让每个人都能看懂数学(make math accessible for everyone)的数学家 Tivadar Danka。这位自称「混乱善良(Chaotic good)」的数学家通过一系列推文和博客文章生动地介绍了矩阵和图的这种等价性及其用途。截至目前,这些推文已被阅读了超过 200 万次,收获了超过 3200 次转发和 9100 次收藏。

51c大模型~合集37_大模型_03

 矩阵与有向图的对价性

51c大模型~合集37_大模型_04

如上图的例子所示,如果我们将其中每一行都视为一个节点,则每一个元素都可表示成一条有向且加权的边。当然,0 元素可以忽略不计。如果该元素位于第 i 行第 j 列,则对应于从节点 i 到节点 j 边。

乍一看,这似乎很复杂,但我们可以先看其中一个节点。

51c大模型~合集37_大模型_05

如图所示,对于这个 3×3 的矩阵,第 1 行对应于最顶部的节点(我们这里称之为 1 号节点),其包含 3 个元素但其中一个为 0,因此该节点延伸出了两条边。其中黄色边表示的是 (1,1) 处的元素 0.5,因此它是指向自身且权重为 0.5 的有向边。同理,蓝色边是指向 2 号节点且权重为 1 的边。

这样一来,我们便能分析出,矩阵的第 i 列便对应于指向 i 号节点的所有边。


51c大模型~合集37_大模型_06

这种等价表示有什么用?

非负矩阵与有向图之间的这种等价性既能帮助我们更好地理解矩阵及其运算,也能帮助简化一些计算过程;反过来,这也能帮助我们从新的视角理解图。

举个例子,矩阵的幂就对应于图中的游走。

51c大模型~合集37_大模型_07

如上图所示,对于 n×n 的方形矩阵 A 的 k 次幂,其中每个元素的求和过程都会纳入所有可能的 k 步游走。

举个例子,假设我们要计算上述 3×3 矩阵的平方。

51c大模型~合集37_大模型_08

如果使用矩阵乘法,则我们需要这样计算:

51c大模型~合集37_大模型_09

对于运算结果的第一个元素,我们可以得到结果 = 0.5×0.5+1×0.2+0×1.8 = 0.45。最终,我们可以得到完整的结果为:

51c大模型~合集37_大模型_10

但如果借助上述的图游走方法,则可以通过游走路径来得到结果。同样,对于结果矩阵的第一个元素,就需要对符合 a_{1,l}→a_{l,1} 的所有 2 步游走路径求和。

但是,如果这个有向图表示的是马尔科夫链的状态,其转移概率矩阵的平方本质上就表示该链 2 步之后达到某个状态的概率。

不仅如此,用图表示矩阵还能让我们深入了解非负矩阵的结构。为此,Danka 表示我们需要先了解「强连通分量(strongly connected components)」这一概念。

强连通分量

什么是强连通分量?对于一个有向图,如果能从该图中的每个节点到达其它每个节点时,我们就说该图是强连通的。如下图所示。

51c大模型~合集37_大模型_11

而强连通分量就是指有向图中能够实现强连通的部分 / 子图。如下图所示,左右各有一个强连通分量,而中间的白色边不属于任何强连通分量。

51c大模型~合集37_大模型_12

下图则展示了另一个例子,其中黄色部分是强连通分量:

51c大模型~合集37_大模型_13

对应于强连通图的矩阵是不可约矩阵,而非负矩阵中的所有其它矩阵都是可约矩阵。

51c大模型~合集37_大模型_14

Danka 通过一个例子给出了解释。(为了说明简单,例子中的所有权重均为单元权重,但实践中这些权重值可以是任意非负值。)

下面将这个包含强连通分量但本身并不强连通的图转写成对应的矩阵形式:

51c大模型~合集37_大模型_15

而这个矩阵是可约矩阵。

51c大模型~合集37_大模型_16

可以看到,在主对角线上的两个子矩阵分别表示两个强连通分量,而右上方的子矩阵表示从第 1 个强连通分量指向第 2 个强连通分量的边,左下方的则表示从第 2 个强连通分量指向第 1 个强连通分量的边(因为没有这样的边,所以全为 0)。

这种书写分块矩阵的形式被称为弗罗贝尼乌斯标准形(Frobenius normal form)。

51c大模型~合集37_大模型_17

那么,我们很自然就会问:我们能将任意非负矩阵都转换成弗罗贝尼乌斯标准形矩阵吗?

通过使用有向图来表示非负矩阵,我们可以轻松地看出答案是肯定的,因为任何表示非负矩阵的有向图都可以表示成互相连接的强连通分量。这个过程非常简单:

  1. 为非负矩阵构建对应的有向图;
  2. 找到其中的强连通分量;
  3. 换更好的方式标注各个节点。

如此便大功告成了!

用图来得到弗罗贝尼乌斯标准形

那么,这个更好的方式是什么呢?

以上述的例子为基础,我们来看看这个过程。

首先,将各个强连通分量融合成单个对象,如下图所示。这时候我们可以将每个强连通分量视为一个黑箱 —— 我们不关心其内部结构,只看其外部连接。

51c大模型~合集37_大模型_18


然后,在这个新图中,我们需要找到只有出边而没有入边的分量。这个具体示例中只有一个,我们将其标记为 0 号:

51c大模型~合集37_大模型_19

接下来一步较为麻烦:对每个分量进行编号,使得每个分量的编号都是离 0 号最远的距离。如下示例能更清晰地说明这一点:

51c大模型~合集37_大模型_20

可以看到,0 号到中间的分量有两条路径,那么选择离 0 最远的那条路径对其进行编号。最终得到:

51c大模型~合集37_大模型_21

实际上,这定义的是分量的顺序。接下来标记各个分量的内部节点:

51c大模型~合集37_大模型_22

如果该图本身来自一个矩阵,则这样的重新标注过程就能得到一个弗罗贝尼乌斯标准形矩阵!

51c大模型~合集37_大模型_23

实际上,这个重新标注的过程就是使用一个置换矩阵 P 对原矩阵执行变换,而该置换矩阵由多个转置矩阵的积构成。

51c大模型~合集37_大模型_24

以下为该定理的完整形式:

51c大模型~合集37_大模型_25

当然,用图表示矩阵的用途远不止于此,比如我们还可以使用矩阵的特征值来定义图的特征值。事实上,这一思路催生了谱图理论(spectral graph theory)这一研究领域。

结语

很显然,矩阵和图之间的这种等价关系既有助于图论研究,也能为线性代数的计算和分析提供一个新视角。其也有一些重要的实际用途,比如 DNA 数据就常被表示成矩阵或图的形式。

另外,我们都知道矩阵运算对于当前的大模型 AI 的重要性,而以知识图谱为代表的图也正通过检索增强式搜索等技术成为当前 AI 的重要助力。将这两者关联起来,或许能在 AI 可解释性以及图人工智能方面带来一些新的突破。至少,这能帮助我们更好地学习线性代数。

实际上,上述内容正是提炼自 Tivadar Danka 正在编写的《Mathematics of Machine Learning》一书。这本书将由浅入深地介绍与机器学习相关的数学知识,让读者真正知其然也知其所以然,并且 Danka 自信地宣称这会是「学习机器学习的最佳资源」。目前他已经在网上发布了两章预览,感兴趣的读者可访问:https://tivadardanka.com/mathematics-of-machine-learning-preview/




#Mamba,Transformer最强竞争者

Mamba 虽好,但发展尚早。

深度学习架构有很多,但近些年最成功的莫过于 Transformer,其已经在多个应用领域确立了自己的主导地位。

如此成功的一大关键推动力是注意力机制,这能让基于 Transformer 的模型关注与输入序列相关的部分,实现更好的上下文理解。但是,注意力机制的缺点是计算开销大,会随输入规模而二次增长,也因此就难以处理非常长的文本。

好在前段时间诞生了一种颇具潜力的新架构:结构化的状态空间序列模型(SSM)。该架构能高效地捕获序列数据中的复杂依赖关系,并由此成为 Transformer 的一大强劲对手。

这类模型的设计灵感来自经典的状态空间模型 —— 我们可以将其看作是循环神经网络和卷积神经网络的融合模型。它们可使用循环或卷积运算进行高效地计算,从而让计算开销随序列长度而线性或近线性地变化,由此大幅降低计算成本。

更具体而言,SSM 最成功的变体之一 Mamba 的建模能力已经可以比肩 Transformer,同时还能维持随序列长度的线性可扩展性。

Mamba 首先引入了一个简单却有效选择机制,其可根据输入对 SSM 进行重新参数化,从而可让模型在滤除不相关信息的同时无限期地保留必要和相关的数据。然后,Mamba 还包含一种硬件感知型算法,可使用扫描(scan)而非卷积来循环地计算模型,这在 A100 GPU 上能让计算速度提升 3 倍。

如图 1 所示,凭借强大的建模复杂长序列数据的能力和近乎线性的可扩展性,Mamba 已经崛起成为一种基础模型,并有望变革计算机视觉、自然语言处理和医疗等多个研究和应用领域。

51c大模型~合集37_大模型_26

因此,研究和应用 Mamba 的文献迅速增长,让人目不暇接,一篇全面的综述报告必定大有裨益。近日,香港理工大学的一个研究团队在 arXiv 上发布了他们的贡献。

  • 论文标题:A Survey of Mamba
  • 论文地址:https://arxiv.org/pdf/2408.01129

这份综述报告从多个角度对 Mamba 进行了总结,既能帮助初学者学习 Mamba 的基础工作机制,也能助力经验丰富的实践者了解最新进展。

Mamba 是一个热门研究方向,也因此有多个团队都在尝试编写综述报告,除了本文介绍的这一篇,还有另一些关注状态空间模型或视觉 Mamba 的综述,详情请参阅相应论文:

  • Mamba-360: Survey of state space models as transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv:2404.16112
  • State space model for new-generation network alternative to transformers: A survey. arXiv:2404.09516
  • Vision Mamba: A Comprehensive Survey and Taxonomy. arXiv:2405.04404
  • A survey on vision mamba: Models, applications and challenges. arXiv:2404.18861
  • A survey on visual mamba. arXiv:2404.15956

预备知识

Mamba 集中了循环神经网络(RNN)的循环框架、Transformer 的并行计算和注意力机制、状态空间模型(SSM)的线性特性。因此,为了透彻地理解 Mamba,就必需先理解这三种架构。

循环神经网络

循环神经网络(RNN)具有保留内部记忆的能力,因此很擅长处理序列数据。

具体来说,在每个离散时间步骤 k,标准 RNN 在处理一个向量时会连同前一时间步骤的隐藏状态一起处理,之后输出另一个向量并更新隐藏状态。这个隐藏状态就可作为 RNN 的记忆,其能保留过去已见过的输入的信息。这种动态记忆让 RNN 可处理不同长度的序列。

也就是说,RNN 是一种非线性的循环模型,可通过使用存储在隐藏状态中历史知识来有效地捕获时间模式。

Transformer

Transformer 的自注意力机制有助于捕获输入之中的全局依赖。其实现方式是基于每个位置相对于其它位置的重要程度为它们分配权重。更具体而言,首先对原始输入进行线性变换,将输入向量的序列 x 转换成三类向量:查询 Q、键 K 和值 V。

然后计算归一化的注意力分数 S 并计算注意力权重。

除了可以执行单个注意力函数,我们还可以执行多头注意力。这让模型可以捕获不同类型的关系,并从多个视角理解输入序列。多头注意力会使用多组自注意力模块并行地处理输入序列。其中每个头都独立运作,执行的计算与标准自注意力机制一样。

之后,将每个头的注意力权重汇聚组合,得到值向量的加权和。这个聚合步骤可让模型使用来自多个头的信息并捕获输入序列中的多种不同模式和关系。

状态空间

状态空间模型(SSM)是一种传统的数学框架,可用于描述系统随时间变化的动态行为。近些年来,人们已将 SSM 广泛应用于控制论、机器人学和经济学等多个不同领域。

究其核心,SSM 是通过一组名为「状态」的隐藏变量来体现系统的行为,使其能有效捕获时间数据的依赖关系。不同于 RNN,SSM 是一种具有关联(associative)属性的线性模型。具体来说,经典的状态空间模型会构建两个关键方程(状态方程和观察方程),以通过一个 N 维的隐藏状态 h (t) 建模当前时间 t 时输入 x 与输出 y 之间的关系。

  • 离散化

为了满足机器学习的需求,SSM 必需经历一个离散化过程 —— 将连续参数转变成离散参数。通常来说,离散化方法的目标是将连续时间划分为具有尽可能相等积分面积的 K 个离散区间。为了实现这一目标,SSM 采用的最具代表性的解决方案之一是 Zero-Order Hold(ZOH),其假设区间 Δ = [𝑡_{𝑘−1}, 𝑡_𝑘 ] 上的函数值保持不变。离散 SSM 与循环神经网络结构相似,因此离散 SSM 能比基于 Transformer 的模型更高效地执行推理过程。

  • 卷积计算

离散 SSM 是一个具有结合属性的线性系统,因此可以与卷积计算无缝整合。

RNN、Transformer 和 SSM 之间的关系

图 2 展示了 RNN、Transformer 和 SSM 的计算算法。

51c大模型~合集37_大模型_27

一方面,常规 RNN 的运作基于一种非线性的循环框架,其中每个计算都仅依赖于之前的隐藏状态和当前输入。

尽管这种形式可让 RNN 在自回归推理时快速生成输出,但它也让 RNN 难以充分利用 GPU 的并行计算能力,导致模型训练速度变慢。

另一方面,Transformer 架构是在多个「查询 - 键」对上并行执行矩阵乘法,而矩阵乘法可以高效地分配给硬件资源,从而更快地训练基于注意力的模型。但是,如果要让基于 Transformer 的模型生成响应或预测,则推理过程会非常耗时。

不同于仅支持一类计算的 RNN 和 Transformer,离散 SSM 灵活性很高;得益于其线性性质,它既能支持循环计算,也可支持卷积计算。这种特性让 SSM 不仅能实现高效推理,也能实现并行训练。但是,需要指出,最常规的 SSM 是时不变的,也就是说其 A、B、C 和 Δ 与模型输入 x 无关。这会限制其上下文感知型建模的能力,导致 SSM 在选择性复制等一些特定任务上表现不佳。

Mamba

为了解决上述传统 SSM 的缺点,实现上下文感知型建模,Albert Gu 和 Tri Dao 提出了可用作通用序列基础模型主干网络的 Mamba。

之后,他们俩又进一步提出了 Mamba-2,其中的结构化空间状态对偶(SSD/Structured Space-State Duality)构建了一个将结构化 SSM 与多种形式的注意力连接起来的稳健的理论框架,让我们可将原本为 Transformer 开发的算法和系统优化技术迁移用于 SSM。

Mamba-1:使用硬件感知型算法的选择式状态空间模型

Mamba-1 基于结构化状态空间模型引入了三大创新技术,即基于高阶多项式投影算子(HiPPO)的内存初始化、选择机制和硬件感知型计算。如图 3 所示。这些技术的目标是提升 SSM 的长程线性时间序列建模能力。

51c大模型~合集37_大模型_28

具体来说,其中的初始化策略可构建一个连贯的隐藏状态矩阵,以有效地促进长程记忆。

然后,选择机制可让 SSM 有能力获取可感知内容的表征。

最后,为了提升训练效率,Mamba 还包含两种硬件感知型计算算法:Parallel Associative Scan(并行关联扫描)和 Memory Recomputation(内存重新计算)。

Mamba-2:状态空间对偶

Transformer 启发了多种不同技术的发展,比如参数高效型微调、灾难性遗忘缓解、模型量化。为了让状态空间模型也能受益于这些原本为 Transformer 开发的技术,Mamba-2 引入了一个新框架:结构化状态空间对偶(SSD)。该框架在理论上将 SSM 和不同形式的注意力连接到了一起。

本质上讲,SSD 表明,Transformer 使用的注意力机制和 SSM 中使用的线性时不变系统都可被视为半可分离的矩阵变换。

此外,Albert Gu 和 Tri Dao 还证明选择式 SSM 等价于使用一种半可分离掩码矩阵实现的结构化线性注意力机制。

Mamba-2 基于 SSD 设计了一种能更高效使用硬件的计算方法,这要用到一种块分解矩阵乘法算法。

具体来说,通过这种矩阵变换将状态空间模型视为半可分离矩阵,Mamba-2 能将该计算分解为矩阵块,其中对角块表示块内计算。而非对角块则表示通过 SSM 的隐藏状态分解的块间计算。该方法可让 Mamba-2 的训练速度超过 Mamba-1 的并行关联扫描的 2-8 倍,同时性能还能媲美 Transformer。

Mamba 块

下面来看看 Mamba-1 和 Mamba-2 的块设计。图 4 比较了这两种架构。

51c大模型~合集37_大模型_29

Mamba-1 的设计是以 SSM 为中心,其中选择式 SSM 层的任务是执行从输入序列 X 到 Y 的映射。在这种设计中,经过了初始的创建 X 的线性投射之后,会使用 (A, B, C) 的线性投射。然后,输入 token 和状态矩阵会通过选择式 SSM 单元,利用并行关联扫描,从而得到输出 Y。之后,Mamba-1 采用了一个 skip 连接,以鼓励特征复用和缓解常在模型训练过程中发生的性能下降问题。最后,通过交错地堆叠该模块与标准归一化和残差连接,便可构建出 Mamba 模型。

至于 Mamba-2,则是引入了 SSD 层来创建从 [X, A, B, C] 到 Y 的映射。其实现方式是在块的起点处使用单个投射来同时处理 [X, A, B, C],这类似于标准注意力架构以并行方式生成 Q、K、V 投射的方式。

也就是说,通过移除序列线性投射,Mamba-2 块是在 Mamba-1 块的基础上进行了简化。这能让 SSD 结构的计算速度超过 Mamba-1 的并行选择式扫描。此外,为了提升训练稳定性,Mamba-2 还在 skip 连接之后添加了一个归一化层。

Mamba 模型正在发展进步

状态空间模型和 Mamba 近来发展迅猛,已经成为了一大极具潜力的基础模型骨干网络选择。尽管 Mamba 在自然语言处理任务上表现不俗,但也仍具有一些难题,比如记忆丢失、难以泛化到不同任务、在复杂模式方面的表现不及基于 Transformer 的语言模型。为了解决这些难题,研究社区为 Mamba 架构提出了诸多改进方案。现有的研究主要集中于修改块设计、扫描模式和记忆管理。表 1 分类总结了相关研究。

51c大模型~合集37_大模型_30

块设计

Mamba 块的设计和结构对 Mamba 模型的总体性能有很大的影响,也因此这成为了一大研究热点。

51c大模型~合集37_大模型_31

如图 5 所示,基于构建新 Mamba 模块的不同方法,现有研究可以分为三类:

  • 集成方法:将 Mamba 块与其它模型集成到一起,实现效果与效率的平衡;
  • 替换方法:用 Mamba 块替换其它模型框架中的主要层;
  • 修改方法:修改经典 Mamba 块内的组件。

扫描模式

并行关联扫描是 Mamba 模型内的一大关键组件,其目标是解决由选择机制导致的计算问题、提升训练过程速度以及降低内存需求。其实现方式是利用时变的 SSM 的线性性质来在硬件层级上设计核融合和重新计算。但是,Mamba 的单向序列建模范式不利于全面学习多样化的数据,比如图像和视频。

51c大模型~合集37_大模型_32

为缓解这一问题,一些研究者探索了新的高效扫描方法,以提升 Mamba 模型的性能以及促进其训练过程。如图 6 所示,在开发扫描模式方面,现有的研究成果可以分为两类:

  • 展平式扫描方法:以展平的视角看待 token 序列,并基于此处理模型输入;
  • 立体式扫描方法:跨维度、通道或尺度扫描模型输入,这又可进一步分为三类:分层扫描、时空扫描、混合扫描。

记忆管理

类似于 RNN,在状态空间模型内,隐藏状态的记忆有效地存储了之前步骤的信息,因此对 SSM 的整体性能有着至关重要的影响。尽管 Mamba 引入了基于 HiPPO 的方法来进行记忆初始化,但管理 SSM 单元中的记忆依然难度很大,其中包括在层之前转移隐藏信息以及实现无损记忆压缩。

为此,一些开创性研究提出了一些不同的解决方案,包括记忆的初始化、压缩和连接。

让 Mamba 适应多样化的数据

Mamba 架构是选择式状态空间模型的一种扩展,其具备循环模型的基本特性,因而非常适合作为处理文本、时间序列、语音等序列数据的通用基础模型。

不仅如此,近期一些开创性研究更是扩展了 Mamba 架构的应用场景,使其不仅能处理序列数据,还能用于图像和图谱等领域,如图 7 所示。

51c大模型~合集37_大模型_33

这些研究的目标是既充分利用 Mamba 能获取长程依赖关系的出色能力,也让其发挥学习和推理过程中的效率优势。表 2 简单总结了这些研究成果。

51c大模型~合集37_大模型_34

序列数据

序列数据是指以特定顺序收集和整理的数据,其中数据点的顺序具有重要意义。这份综述报告全面总结了 Mamba 在多种序列数据上的应用,包括自然语言、视频、时间序列、语音和人体运动数据。详见原论文。

非序列数据

不同于序列数据,非序列数据并不遵循特定的顺序。其数据点可以任意顺序进行组织而不会对数据的含义造成显著影响。对于专门设计用于捕获数据中时间依赖关系的循环模型(RNN 和 SSM 等)来说,这种缺乏固有顺序的数据会很难处理。

令人惊讶的是,近期的一些研究成功让 Mamba(代表性的 SSM)实现了对非序列数据的高效处理,包括图像、图谱和点云数据。

多模态数据

为了提升 AI 的感知和场景理解能力,可以整合多个模态的数据,比如语言(序列数据)和图像(非序列数据)。这样的整合能提供非常有价值和补充性的信息。

近段时间来,多模态大型语言模型(MLLM)是最受关注的研究热点;这类模型继承了大型语言模型(LLM)的强大能力,包括强大的语言表达和逻辑推理能力。尽管 Transformer 已经成为该领域的主导方法,但 Mamba 也正在崛起成为一大强劲竞争者,其在对齐混合源数据和实现序列长度的线性复杂度扩展方面表现出色,这使 Mamba 有望在多模态学习方面替代 Transformer。

应用

下面介绍基于 Mamba 的模型的一些值得注意的应用。该团队将这些应用分为了以下类别:自然语言处理、计算机视觉、语音分析、药物发现、推荐系统以及机器人和自主系统。

这里我们不再过多介绍,详见原论文。

挑战与机遇

Mamba 虽然已经在一些领域取得了出色表现,但总体而言,Mamba 研究仍还处于起步阶段,前方仍还有一些挑战有待克服。当然,这些挑战同时也是机遇。

  • 如何开发和改进基于 Mamba 的基础模型;
  • 如何充分实现硬件感知型计算,以尽可能利用 GPU 和 TPU 等硬件,提升模型效率;
  • 如何提升 Mamba 模型的可信度,这需要安全和稳健性、公平性、可解释性以及隐私方面的进一步研究;
  • 如何将 Transformer 领域的新技术用于 Mamba,如参数高效型微调、灾难性遗忘缓解、检索增强式生成(RAG)。




#硬件直接支持混合矩阵乘法

大模型终端部署新趋势

在人工智能领域,模型参数的增多往往意味着性能的提升。但随着模型规模的扩大,其对终端设备的算力与内存需求也日益增加。低比特量化技术,由于可以大幅降低存储和计算成本并提升推理效率,已成为实现大模型在资源受限设备上高效运行的关键技术之一。然而,如果硬件设备不支持低比特量化后的数据模式,那么低比特量化的优势将无法发挥。

为了解决这一问题,微软亚洲研究院推出了全新的数据编译器 Ladder 和算法 T-MAC,使当前只支持对称精度计算的硬件能够直接运行混合精度矩阵乘法。测试结果表明,Ladder 在支持 GPU 原本不支持的自定义数据类型方面,最高提速可达 14.6 倍;T-MAC 在搭载了最新高通 Snapdragon X Elite 芯片组的 Surface AI PC 上,使 CPU 上运行的大模型吞吐率比专用加速器 NPU 快两倍。此外,研究员们还设计了 LUT Tensor Core 硬件架构,这种精简设计使硬件能够直接支持各种低比特混合精度计算,为人工智能硬件设计提供了新思路。

大模型已经越来越多地被部署在智能手机、笔记本电脑、机器人等端侧设备上,以提供先进的智能及实时响应服务。但包含上亿参数的大模型对终端设备的内存和计算能力提出了极高的要求,也因此限制了它们的广泛应用。低比特量化技术因其能显著压缩模型规模,降低对计算资源的需求,成为了大模型在端侧部署和实现高效推理的有效手段。

随着低比特量化技术的发展,数据类型日益多样化,如 int4、int2、int1 等低比特数据,使得大模型在推理中越来越多地采用低比特权重和高比特权重计算的混合精度矩阵乘法(mixed-precision matrix multiplication,mpGEMM)。然而,现有的 CPU、GPU 等硬件计算单元通常只支持对称计算模式,并不兼容这种混合精度的矩阵乘法。

混合精度矩阵乘法与传统的矩阵乘法有何不同?

在传统的矩阵乘法中,参与运算的两端数值是对称的,例如 FP16*FP16、int8*int8。但大模型的低比特量化打破了这种对称性,使乘法的一端是高比特,另一端是低比特,例如在 1-bit 的 BitNet 模型中实现的 int8*int1 或 int8*int2,以及浮点数与整数的混合乘法 FP16*int4。

为了充分发挥低比特量化的优势,让硬件设备能够直接支持混合精度矩阵乘法,确保大模型在端侧设备上的高速有效运行,微软亚洲研究院的研究员们针对现有 CPU、GPU 计算算子和硬件架构进行创新:

  • 推出了数据类型编译器 Ladder,支持各种低精度数据类型的表达和相互转换,将硬件不支持的数据类型无损转换为硬件支持的数据类型指令,在传统计算模式下,使得硬件能够支持混合精度的 DNN(深度神经网络) 计算;
  • 研发了全新算法 T-MAC,基于查找表(Lookup Table,LUT)的方法,实现了硬件对混合精度矩阵乘法的直接支持,软件层面,在 CPU 上的计算相比传统计算模式取得了更好的加速;
  • 提出了新的硬件架构 LUT Tensor Core,为下一代人工智能硬件设计打开了新思路。

Ladder:自定义数据类型无损转换成硬件支持的数据类型

当前,前沿加速器正在将更低比特的计算单元,如 FP32、FP16,甚至 FP8 的运算集成到新一代的架构中。然而,受限于芯片面积和高昂的硬件成本,每个加速器只能为标准的数据类型提供有限类型的计算单元,比如 NVIDIA V100 TENSOR CORE GPU 仅支持 FP16,而 A100 虽然加入了对 int2、int4、int8 的支持,但并未涵盖更新的 FP8 或 OCP-MXFP 等数据格式。此外,大模型的快速迭代与硬件升级的缓慢步伐之间存在差距,导致许多新数据类型无法得到硬件支持,进而影响大模型的加速和运行。

微软亚洲研究院的研究员们发现,尽管硬件加速器缺乏针对自定义数据类型的计算指令,但其内存系统可以将它们转换为固定位宽的不透明数据块来存储任意数据类型。同时,大多数自定义数据类型可以无损地转换为现有硬件计算单元支持的更多位的标准数据类型。例如,NF4 张量可以转换成 FP16 或 FP32 以执行浮点运算。

基于这些发现,研究员们提出了一种通过分离数据存储和计算来支持所有自定义数据类型的方法,并研发了数据编译器 Ladder,以弥合不断出现的自定义数据类型与当前硬件支持的固有精度格式之间的差距。

Ladder 定义了一套数据类型系统,包括数据类型之间无损转换的抽象,它能够表示算法和硬件支持的各种数据类型,并定义了数据类型之间的转换规则。当处理低比特算法应用时,Ladder 通过一系列优化,将低比特数据转译成当前硬件上最高效的执行格式,包括对计算和存储的优化 —— 将算法映射到匹配的计算指令,并将不同格式的数据存储到不同级别的存储单元中,以实现最高效的运算。

51c大模型~合集37_大模型_35

图 1:Ladder 的系统架构

在 NVIDIA A100、NVIDIA V100、NVIDIA RTX A6000、NVIDIA RTX 4090 和 AMD Instinct MI250 GPU 上运行的 DNN 推理性能评估显示,Ladder 在原生支持数据类型上超越了现有最先进的 DNN 编译器,并且在支持 GPU 原本不支持的自定义数据类型方面表现出色,最高提速可达 14.6 倍。

Ladder 是首个在现代硬件加速器上运行 DNN 时,可以系统性地支持以自定义数据类型表示低比特精度数据的系统。这为模型研究者提供了更灵活的数据类型优化方法,同时也让硬件架构开发者在不改变硬件的情况下,支持更广泛的数据类型。

T-MAC:无需乘法的通用低比特混合精度矩阵乘计算

为了让现有硬件设备支持不同的数据模式和混合精度矩阵乘法,在端侧部署大模型时,常见的做法是对低比特模型进行反量化。然而,这种方法存在两大问题:首先,从性能角度来看,反量化过程中的转换开销可能会抵消低比特量化带来的性能提升;其次,从开发角度来看,开发者需要针对不同的混合精度重新设计数据布局和计算内核。微软亚洲研究院的研究员们认为,在设备上部署低比特量化的大模型,关键在于如何基于低比特的特点来突破传统矩阵乘法的实现。

为此,研究员们从系统和算法层面提出了一种基于查找表(LUT,Look-Up Table)的方法 T-MAC,帮助低比特量化的大模型在 CPU 上实现高效推理。T-MAC 的核心思想在于利用混合精度矩阵乘法的一端为极低比特(如 1 比特或 2 比特)的特点。它们的输出结果只有 2 的 1 次方和 2 的 2 次方种可能,这些较少的输出结果完全可以提前计算并存储在表中,在运算时,只需从表中读取结果,避免了重复计算,大幅减少了乘法和加法的运算次数。

具体而言,T-MAC 将传统的以数据类型为中心的乘法转变为基于位的查找表操作,实现了一种统一且可扩展的混合精度矩阵乘法解决方案,减小了表的大小并使其停留在最快的内存单元中,降低了随机访问表的成本。这一创新为在资源受限的边缘设备上部署低比特量化大模型铺平了道路。

51c大模型~合集37_大模型_36

图 2:T-MAC 示意图

在针对低比特量化的 Llama 和 1 比特的 BitNet 大语言模型的测试中,T-MAC 展现出了显著的性能优势。在搭载了最新高通 Snapdragon X Elite 芯片组的 Surface Laptop 7 上,T-MAC 让 3B BitNet-b1.58 模型的生成速率达到每秒 48 个 token,2bit 7B Llama 模型的生成速率达到每秒 30 个 token,4bit 7B Llama 模型的生成速率可达每秒 20 个 token,这些速率均远超人类的平均阅读速度。与原始的 Llama.cpp 框架相比,其提升了 4 至 5 倍,甚至比专用的 NPU 加速器还快两倍。

即使是在性能较低的设备上,如 Raspberry Pi(树莓派)5,T-MAC 也能使 3B BitNet-b1.58 模型达到每秒 11 个 token 的生成速率。T-MAC 还具有显著的功耗优势,在资源受限的设备上可以达到相同的生成速率,而它所需的核心数仅为原始 Llama.cpp 的 1/4 至 1/6。

这些结果表明,T-MAC 提供了一种实用的解决方案,使得在使用通用 CPU 的边缘设备上部署大语言模型更为高效,且无需依赖 GPU,让大模型在资源受限的设备上也能高效运行,从而推动大模型在更广泛的场景中的应用。

LUT Tensor Core:推动下一代硬件加速器原生支持混合精度矩阵乘法

T-MAC 和 Ladder 都是在现有 CPU 和 GPU 架构上,实现对混合精度矩阵乘法的优化支持。尽管这些软件层面的创新显著提升了计算效率,但它们在效率上仍无法与能够直接实现一个专门查找表的硬件加速器相比。研究员们认为,最理想的方法是重新设计硬件加速器,让 CPU、GPU 等能够原生支持混合精度矩阵乘法,但这一目标面临三大挑战:

  • 效率:设计和实现方式必须具有成本效益,通过优化芯片的利用面积,最大限度地提高低比特数据的计算效益。
  • 灵活性:由于不同的模型和场景需要不同的权重和激活精度,因此硬件中的混合精度矩阵乘法设计必须能够处理各种权重精度 (如 int4/2/1) 和激活精度 (如 FP16/8、int8) 及其组合。
  • 兼容性:新设计必须与现有的 GPU 架构和软件生态系统无缝集成,以加速新技术的应用。

为了应对这些挑战,微软亚洲研究院的研究员们设计了 LUT Tensor Core,这是一种利用查找表直接执行混合精度矩阵乘法的 GPU Tensor Core 微架构。一方面,基于查找表的设计将乘法运算简化为表预计算操作,可直接在表中查找结果,提高计算效率。另一方面,这种方法也简化了对硬件的需求,它只需用于表存储的寄存器和用于查找的多路选择器,无需乘法器和加法器。同时,LUT Tensor Core 通过比特串行设计实现了权重精度的灵活性,并利用表量化实现了激活精度的灵活性。

此外,为了与现有 GPU 微架构和软件堆栈集成,研究员们扩展了 GPU 中现有的 MMA 指令集,加入了一组 LMMA 指令,并设计了一个类似于 cuBLAS 的软件堆栈,用于集成到现有的 DNN 框架中。研究员们还设计了一个编译器,用于在具有 LUT Tensor Core 的 GPU 上进行端到端的执行计划。这些创新方法可以让 LUT Tensor Core 被无缝、快速地采用。

51c大模型~合集37_大模型_37

图 3:LUT Tensor Core 微架构概述

在 Llama 和 BitNet 模型上的测试显示,LUT Tensor Core 可以提供高达 6.93 倍的推理速度,且只占传统 Tensor Core 面积的 38.7%。在几乎相同的模型精度下,这相当于 20.7 倍的计算密度和 19.1 倍的能效提升。随着人工智能大模型规模和复杂性的不断增长,LUT Tensor Core 有助于进一步释放低比特大语言模型的潜力,推动人工智能在新场景中的应用。

“查找表方法引领了计算范式的转变。在过去,我们依赖于矩阵乘法和累加运算,而在大模型时代,得益于低比特量化技术,查找表方法将成为主流。相较于传统的浮点运算或矩阵乘法,查找表方法在计算上更轻便高效,而且在硬件层面上更易于扩展,能够实现更高的晶体管密度,在单位芯片面积上提供更大的吞吐量,从而推动硬件架构的革新。” 微软亚洲研究院首席研究员曹婷表示。

低比特量化的长尾效应:为具身智能带来新可能

低比特量化技术不仅优化了大模型在端侧设备上的运行效率,还通过减少单个参数的 “体积”,为模型参数的扩展(Scale up)提供了新的空间。这种参数扩展能力,使模型拥有了更强的灵活性和表达能力,正如 BitNet 模型所展示的,从低比特模型出发,逐步扩展至更大规模的训练。

微软亚洲研究院的 T-MAC、Ladder 和 LUT Tensor Core 等创新技术,为各种低比特量化大模型提供了高效能的运行方案,使得这些模型能够在各种设备上高效运行,并推动科研人员从低比特角度设计和优化大模型。其中部分技术已经在微软必应(Bing)搜索及其广告业务等搜索大模型中发挥作用。随着对内存和计算资源的降低,低比特大模型在机器人等具身智能系统上的部署也将成为可能,可以使这些设备更好地实现与环境的动态感知和实时交互。开发板商城 天皓智联 TB上有视觉设备哦 支持AI相关~ 大模型相关也可用whaosoft aiot

目前,T-MAC 和 Ladder 已经在 GitHub 上开源,欢迎相关研发人员测试应用,与微软亚洲研究院共同探索人工智能技术的更多可能。

相关链接:




#Body Transformer

机器人策略学习的Game Changer?伯克利提出Body Transformer

过去几年间,Transformer 架构已经取得了巨大的成功,同时其也衍生出了大量变体,比如擅长处理视觉任务的 Vision Transformer(ViT)。本文要介绍的 Body Transformer(BoT) 则是非常适合机器人策略学习的 Transformer 变体。

我们知道,物理智能体在执行动作的校正和稳定时,往往会根据其感受到的外部刺激的位置给出空间上的响应。比如人类对这些刺激的响应回路位于脊髓神经回路层面,它们专门负责单个执行器的响应。起校正作用的局部执行是高效运动的主要因素,这对机器人来说也尤为重要。

但之前的学习架构通常都没有建立传感器和执行器之间的空间关联。鉴于机器人策略使用的架构基本是为自然语言和计算机视觉开发的架构,它们常常无法有效地利用机器人机体的结构。

不过,Transformer 在这方面还是颇具潜力的,已有研究表明,Transformer 可以有效地处理长序列依赖关系,还能轻松地吸收大量数据。Transformer 架构原本是为非结构化自然语言处理(NLP)任务开发的。在这些任务中(比如语言翻译),输入序列通常会被映射到一个输出序列。

基于这一观察,加州大学伯克利分校 Pieter Abbeel 教授领导的团队提出了 Body Transformer(BoT),增加了对机器人机体上的传感器和执行器的空间位置的关注。

  • 论文标题:Body Transformer: Leveraging Robot Embodiment for Policy Learning
  • 论文地址:https://arxiv.org/pdf/2408.06316v1
  • 项目网站:https://sferrazza.cc/bot_site
  • 代码地址:https://github.com/carlosferrazza/BodyTransformer

具体来说,BoT 是将机器人机体建模成图(graph),其中的节点即为其传感器和执行器。然后,其在注意力层上使用高度稀疏的掩码,以防止每个节点关注其直接近邻之外的部分。将多个结构相同的 BoT 层连接起来,就能汇集整个图的信息,这样便不会损害该架构的表征能力。BoT 在模仿学习和强化学习方面都表现不俗,甚至被一些人认为是策略学习的「Game Changer」。

Body Transformer

如果机器人学习策略使用原始 Transformer 架构为骨干,则通常会忽视机器人机体结构所提供的有用信息。但实际上,这些结构信息能为 Transformer 提供更强的归纳偏置。该团队在利用这些信息的同时还保留了原始架构的表征能力。

Body Transformer(BoT)架构基于掩码式注意力。在这个架构的每一层中,一个节点都只能看到其自身和其直接近邻节点的信息。如此一来,信息就会依照图的结构而流动,其中上游层会根据局部信息执行推理,下游层则能汇集更多来自更远节点的全局信息。

51c大模型~合集37_大模型_38

如图 1 所示,BoT 架构包含以下组件:

1.tokenizer:将传感器输入投射成对应的节点嵌入;

2.Transformer 编码器:处理输入嵌入并生成同样维度的输出特征;

3.detokenizer:解除 token 化,即将特征解码成动作(或用于强化学习批评训练的价值)。

tokenizer

该团队选择将观察向量映射成局部观察构成的图。

在实践中,他们将全局量分配给机器人机体的根元素,将局部量分配给表示对应肢体的节点。这种分配方式与之前的 GNN 方法类似。

然后,使用一个线性层将局部状态向量投射成嵌入向量。每个节点的状态都会被馈送给其节点特定的可学习的线性投射,从而得到一个包含 n 个嵌入的序列,其中 n 表示节点的数量(或序列长度)。这不同于之前的研究成果,它们通常仅使用单个共享的可学习的线性投射来处理多任务强化学习中不同数量的节点。

BoT 编码器

该团队使用的骨干网络是一个标准的多层 Transformer 编码器,并且该架构有两种变体版本:

  • BoT-Hard:使用一个反映该图结构的二元掩码来掩蔽每一层。具体来说,他们构建掩码的方式是 M = I_n + A,其中 I_n 是 n 维单位矩阵,A 是对应于该图的邻接矩阵。图 2 展示了一个示例。这让每个节点仅能看到其自身和其直接近邻,并且能为该问题引入相当可观的稀疏性 —— 从计算成本角度看,这特别有吸引力。


51c大模型~合集37_大模型_39

  • BoT-Mix:将带有掩码式注意力的层(如 BoT-Hard 一样)与带有无掩码式注意力的层交织在一起。

detokenizer

Transformer 编码器输出的特征会被馈送给线性层,然后被投射成与该节点的肢体关联的动作;这些动作是根据相应执行器与肢体的接近程度来分配的。同样,每个节点的这些可学习的线性投射层是分开的。如果将 BoT 用作强化学习设置中的批评架构,则 detokenizer 输出的就不再是动作,而是价值,然后在机体部位上取平均值。

实验

团队在模仿学习和强化学习设置中评估了 BoT 的性能。他们维持了与图 1 相同的结构,只用各种基线架构替换 BoT 编码器,以确定编码器的效果。

这些实验的目标是解答以下问题:

  • 掩码式注意力是否能提升模仿学习的性能和泛化能力?
  • 相比于原始的 Transformer 架构,BoT 是否能表现出正面的规模扩展趋势?
  • BoT 是否与强化学习框架兼容,有哪些合理设计选择可以尽可能地提升性能?
  • BoT 策略是否可以应用于真实世界机器人任务?
  • 掩码式注意力在计算方面有哪些优势?

模仿学习实验

团队在机体跟踪任务上评估了 BoT 架构的模仿学习性能,该任务是通过 MoCapAct 数据集定义的。

结果如图 3a 所示,可以看到 BoT 的表现总是优于 MLP 和 Transformer 基线。值得注意的是,在未曾见过的验证视频片段上,BoT 相对于这些架构的优势还会进一步增大,这证明机体感知型归纳偏置能带来泛化能力的提升。

51c大模型~合集37_大模型_40

而图 3b 则表明 BoT-Hard 的规模扩展性很不错,相较于 Transformer 基线,其在训练和验证视频片段上的性能都会随着可训练参数量的增长而增长这进一步表明 BoT-Hard 倾向于不过拟合训练数据,而这种过拟合是由具身偏置引起的。下面展示了更多实验示例,详见原论文。

51c大模型~合集37_大模型_41

51c大模型~合集37_大模型_42

强化学习实验

该团队在 Isaac Gym 中的 4 个机器人控制任务上评估了 BoT 与使用 PPO 的基线的强化学习性能。这 4 个任务分别是:Humanoid-Mod、Humanoid-Board、Humanoid-Hill 和 A1-Walk。

图 5 展示了 MLP、Transformer 和 BoT(Hard 和 Mix)在训练期间的评估 rollout 的平均情节回报。其中,实线对应于平均值,阴影区域对应于五个种子的标准误差。

51c大模型~合集37_大模型_43

结果表明,BoT-Mix 的性能在样本效率和渐近性能方面始终优于 MLP 和原始 Transformer 基线。这说明将来自机器人机体的偏置整合进策略网络架构是有用的。

同时,BoT-Hard 在较简单的任务(A1-Walk 和 Humanoid-Mod)上的表现优于原始 Transformer,但在更困难的探索任务(Humanoid-Board 和 Humanoid-Hill)上表现却更差。考虑到掩码式注意力会妨碍来自远处机体部分的信息传播,BoT-Hard 在信息通信方面的强大限制可能会妨碍强化学习探索的效率。

真实世界实验

Isaac Gym 模拟的运动环境常被用于将强化学习策略从虚拟迁移到真实环境,并且还不需要在真实世界中进行调整。为了验证新提出的架构是否适用于真实世界应用,该团队将上述训练得到的一个 BoT 策略部署到了一台 Unitree A1 机器人中。从如下视频可以看出,新架构可以可靠地用于真实世界部署。

计算分析

该团队也分析了新架构的计算成本,如图 6 所示。这里给出了新提出的掩码式注意力与常规注意力在不同序列长度(节点数量)上的规模扩展结果。

51c大模型~合集37_大模型_44

可以看到,当有 128 个节点时(相当于拥有灵巧双臂的类人机器人),新注意力能将速度提升 206%。

总体而言,这表明 BoT 架构中的源自机体的偏置不仅能提高物理智能体的整体性能,而且还可受益于架构那自然稀疏的掩码。该方法可通过充分的并行化来大幅减少学习算法的训练时间。



#小模型~~

浏览器里直接能跑的SOTA小模型来了,分别在2亿、5亿和20亿级别获胜,抱抱脸出品。

秘诀只有两个:

  • 狠狠地过滤数据
  • 在高度过滤的数据集上狠狠地训练

抱抱脸首席科学家Thomas Wolf,总结团队在开发小模型时的经验,抛出新观点,引起业界关注:

合成数据目前只在特定领域有用,网络是如此之大和多样化,真实数据的潜力还没完全发挥。

目前360M模型版本已发布Demo,在线可玩(注意流量)。

在浏览器里调用本地GPU运行,连模型权重带网页前端UI,400MB搞定。

严格过滤网络数据,性能直线上升

针对微软Phi系列小模型,声称使用了一半合成数据,效果很好,但不公开数据。

开源界扛把子抱抱脸看不下去了:

造一个对标的大型合成数据集,开源它。

而且,团队隐隐暗示了,此举也有检验微软在测试集上刷榜的传闻,到底有没有这回事的考虑。

抱抱脸使用当时最好的开源模型Mixtral-8-7B构造了25B合成数据。

训练出来的模型效果还不错,但仍然在某种程度上低于Phi-1和Phi-1.5的水平。

他们尝试了让大模型在中学水平上解释各种主题,最终只有在MMLU测试上表现不好,因为MMLU是博士水平的题目。

真正的性能突破,反而来自一项支线任务:

除了用大模型从头生成合成数据,也试试用大模型筛选过滤网络数据

具体来说是使用Llama3-70B-Struct 生成的标注开发了一个分类器,仅保留FineWeb数据集中最具教育意义的网页

使用经过严格过滤的网络数据后,性能直线上升,并在大多数基准测试中超过了所有其他类似大小的模型,包括Phi-1.5。

抱抱脸团队称这项实验结果是“苦乐参半”的:虽然模型性能前所未有的高,但也显示出了合成数据还是比不过真实数据。

后来他们用同样的思路从自然语言扩展到代码,过滤的代码数据集也被证明是非常强大的。

将HumanEval基准测试成绩从13%左右直接提高到20%以上。

最终他们构造的混合数据集中,去重的过滤数据集占绝大部分,纯合成数据Cosmopedia v2只占15%。

所以总得来说,合成数据还有用吗?

团队认为,可能只对确实缺少真实数据的领域更有意义了,比如推理和数学。

即使小模型也要训练数万亿tokens

就在他们对这些新发现和结果感到兴奋时,一位新实习生Elie Bakouch加入了。

虽然他当时只是实习生,但确是一位精通各类训练技巧的专家。

在Elie的帮助下,团队将模型尺寸从1.7B开始下降到360M甚至170M,也就是对标经典模型GPT-1、GPT-2和BERT。

在这个过程中有了第二个重要发现:与过去的共识不同,即使是小模型也要在数万亿token上训练,时间越长越好。

此外数据退火(Anneal the data)也被证明是有效的,也就是在训练的最后一部分保留一组特殊的高质量数据。

最终发布的系列模型适合部署在从智能手机到笔记本电脑的各种设备上,最大的1.7B模型BF16精度只占3G内存。

作为参考,iPhone 15入门版也有6G,安卓手机就更多了。

虽然这次训练出来的基础模型足够好,但团队也还是发现一个问题。

过去的对齐和微调技术,如SFT、DPO、PPO等都是针对大模型非常有效,但对小模型效果并不理想。

团队分析,对齐数据集中包含许多对小模型来说过于复杂的概念,并且缺乏精心设计的简单任务。

下一个新坑也挖好了,有兴趣的团队可以开始搞起,没准就成了小模型大救星。

在线试玩:https://huggingface.co/spaces/HuggingFaceTB/instant-smollm

参考链接:
[1]https://huggingface.co/blog/smollm[2]https://x.com/Thom_Wolf/status/1825094850686906857