声明：本文源自微信公众号机器之心（ID:almosthuman2014），经作者授权由微新创想转载发布。若能以三维方式呈现矩阵乘法的执行过程，当年学习矩阵乘法时或许就不会如此吃力。如今，矩阵乘法已成为机器学习模型的核心构建模块，是各类强大AI技术的基石。深入理解其执行方式，将帮助我们更透彻地把握AI本质，洞察这个日益智能化的世界。本文将介绍PyTorch博客中一款名为mm的可视化工具，它通过三维空间直观展示矩阵乘法及组合过程。相较于传统二维图表，mm能激发更多灵感，降低认知负担，尤其适合擅长视觉与空间思维的人群。结合三维维度与预训练权重加载能力，mm可可视化大型复合表达式（如注意力头），并观察其真实行为模式。该工具完全交互式，运行于浏览器或笔记本iframe中，完整状态保存在URL，便于生成可共享的会话（文中截图和视频均附有可打开对应可视化的链接，详见原博客）。本指南将全面介绍其功能。工具地址：https://bhosmer.github.io/mm/ref.html博客原文：https://pytorch.org/blog/inside-the-matrix
首先介绍可视化方法，通过简单矩阵乘法与表达式可视化建立直觉，再深入探讨扩展示例：介绍：为何三维可视化更优越？热身：动画——观察标准矩阵乘法分解过程热身：表达式——快速浏览基本表达式构建模块深入注意力头：深度解析NanoGPT中GPT-2一对注意力头的结构、值与计算行为并行化注意力：可视化Blockwise Parallel Transformer论文中的注意力头并行化示例注意力层大小：当整个注意力层可视化为单一结构时，MHA与FFN部分如何呈现？自回归解码过程中图像如何变化？LoRA：对注意力头架构的低秩分解可视化阐释
介绍mm的三维可视化方法基于这一前提：矩阵乘法本质上是三维运算。换言之，它可以被三维化呈现为：当我们将矩阵乘法包裹在立方体中时，参数形状、结果形状与共享维度间的关系将一目了然。此时，矩阵乘法计算获得几何意义：结果矩阵中每个位置i，j都锚定一个沿立方体深度维度k运行的向量，该向量由L的第i行水平面与R的第j列垂直面相交而成。沿此向量，来自左右参数的成对(i,k)(k,j)元素相遇相乘，再沿k维度求和，最终结果存入i，j位置。这就是矩阵乘法的直观含义：1. 将两个正交矩阵投影至立方体内部；2. 将每个交叉点的一对值相乘，形成乘积网格；3. 沿第三正交维度求和，生成结果矩阵。方向方面，工具会在立方体内部显示指向结果矩阵的箭头，蓝色箭羽源自左侧参数，红色箭羽源自右侧参数。此外，白色指示线标示每个矩阵的行轴线（尽管图中模糊）。布局约束简洁明了：左侧参数与结果沿共享高度（i）维度邻接；右侧参数与结果沿共享宽度（j）维度邻接；左侧参数与右侧参数沿共享（左宽度/右高度）维度邻接，该维度成为矩阵乘法的深度（k）维度。这种几何表示为可视化所有标准矩阵乘法分解奠定了基础，也为探索复杂矩阵乘法组合提供了直观框架。
热身：动画
在深入复杂示例前，先通过动画直观感受可视化风格，建立工具使用直觉。2a点积：以经典算法为例，通过计算左侧行与右侧列的点积计算每个结果元素。动画显示相乘值向量扫过立方体内部，每次在对应位置提交求和结果。此处L包含填充1（蓝色）或-1（红色）的行块，R包含类似填充的列块。k为24，因此结果矩阵(L@R)的蓝色值为24，红色值为-24。2b矩阵-向量积分解：分解为矩阵-向量积的矩阵乘法，呈现垂直平面，水平扫过立方体内部绘制列。随机初始化参数时，中间矩阵-向量积突出垂直模式，反映每个中间值都是左侧参数列的缩放副本。2c向量-矩阵积分解：分解为向量-矩阵积的矩阵乘法，呈现水平平面，向下穿过立方体内部绘制行。随机初始化参数时，可见类似矩阵-向量积模式，但为水平模式，对应右侧参数行缩放的副本。思考矩阵乘法如何表示参数秩与结构时，可设想这两种模式同时发生。2d对外积求和：沿k轴对向量外积逐点求和计算矩阵乘法结果，外积平面从后向前扫过，累积到结果中。随机初始化矩阵时，不仅可见数值，还能观察秩累积，因每个秩为1的外积被添加其中。直观说明低秩因式分解在近似低秩矩阵时效果最佳，这正是LoRA技术原理。
热身：表达式
探讨如何将可视化方法扩展至矩阵乘法组合表达式。之前示例可视化为L@R单次矩阵乘法，但若L和/或R本身也是矩阵乘法呢？事实证明该方法可很好地扩展至复合表达式。关键规则简单：子表达式矩阵乘法是另一个立方体，受与父表达式相同的布局约束；子矩阵乘法结果面同时是父表达式对应参数面，如同共价共享电子。在约束限制下，可按需排布子矩阵乘法各面。默认方案生成交替凸凹面立方体，最大化空间利用率同时减少遮挡（但布局完全可定制，详情见mm工具页面）。本节将可视化机器学习模型关键构建模块，帮助读者熟悉视觉表示并获取新直觉。
3a左结合表达式：介绍两个(A@B)@C形式表达式，每个具有独特形状特征。mm遵循矩阵乘法左结合约定，故(A@B)@C可简写为A@B@C。首先赋予A@B@C具有特点的FFN形状，其中”隐藏维度”比”输入”或”输出”维度宽（具体示例中B宽度大于A或C）。与单次乘法类似，浮动箭头指向结果矩阵，蓝色箭羽源自左侧参数，红色源自右侧参数。当B宽度小于A或C时，A@B@C可视化呈现类似自动编码器的瓶颈形状。交替凹凸模块模式可扩展为任意长度链，如多层瓶颈。
3b右结合表达式：可视化右结合表达式A@(B@C)。与左结合水平扩展类似，可视为从根表达式左侧参数发端，右结合表达式链垂直扩展，从根表达式右侧参数发端。有时可见以右结合形式形成的MLP，即右侧为柱状输入，权重层从右到左运行。使用上述二层FFN示例矩阵（适当转置后），C为输入，B为第一层，A为第二层。区分左右参数的视觉提示：左侧参数行与结果行共面，沿同一i轴堆叠。如(B@C)示例，这两个提示表明B是左侧参数。
3c二元表达式：可视化工具需能处理复杂表达式。真实用例中关键结构是二元表达式，左右参数均有子表达式。此处可视化为形状最简单的(A@B)@(C@D)。3d注解：分区与并行性。完整阐述超出范围，但后续注意力头部分将展示其实际效用。此处通过简单示例了解可视化如何让并行化推理更直观——通过几何分区。第一个示例将典型”数据并行”分区应用于左结合多层瓶颈示例，沿i轴分区初始左侧参数（批）、中间结果（激活）但不分区后续参数（权重），几何结构清晰显示分区参与者。第二个示例展示如何沿j轴对左侧子表达式分区、沿i轴对右侧子表达式分区、沿k轴对父表达式分区并行化二元表达式。
深入注意力头：现在观察GPT-2注意力头，具体为NanoGPT 5层第4头”gpt2″(small)配置（层数=12，头数=12，嵌入数=768），使用HuggingFace提供的OpenAI权重。输入激活取自含256个token的OpenWebText训练样本单次前向通过。选择该头因其计算常见注意力模式，且位于模型中部，激活已结构化并呈现有趣纹理。4a结构：完整注意力头可视化为单个复合表达式，始于输入终于投影输出（注：为保证自足性，按Megatron-LM描述对每个头执行输出投影）。计算包含六次矩阵乘法：Q=input@wQ//1，K_t=wK_t@input_t//2，V=input@wV//3，attn=sdpa(Q@K_t)//4，head_out=attn@V//5，out=head_out@wO//6。风车叶片为矩阵乘法1、2、3、6：前一组为输入到Q、K、V的内投影；后一组为从attn@V回到嵌入维度的外投影。中心有两个矩阵乘法：第一个计算注意力分数（凸立方体），后一个基于值向量得到输出token（凹立方体）。因果关系使注意力分数形成下三角形。建议亲自探索工具，而非仅看截图或视频，以理解结构及计算过程中的实际值。
4b计算和值：注意力头计算过程动画展示了sdpa(input@wQ@K_t)@V@wO（即矩阵乘法1、4、5、6，其中K_t和V已预计算）作为向量-矩阵积融合链的计算：序列中每项在一步内从输入穿过注意力到输出。计算值揭示许多有趣现象：低秩Q和K_t形态惊人。放大Q@K_t向量-矩阵积动画更生动：大量通道在序列中保持恒定，表明有用注意力信号可能仅由部分嵌入驱动。最熟悉的是注意力矩阵中强大但不完美的对角线，产生局部注意力：紧邻输出token位置前小邻域的值token很大程度上决定输出内容。邻域大小和token影响变化很大，可见于注意力网格非对角frost和attn[i]@V向量-矩阵积平面波动模式。但局部邻域非唯一：注意力网格最左列（对应序列首token）完全填充非零值，意味着每个输出token受首token影响。当前token邻域与初始token间注意力分数主导性存在可辨别的振荡，周期各异，初期短，沿序列向下移动变长（与每一行候选注意力token数量相关）。理解(attn@V)形成需关注V同等重要性：每个输出项是整个V向量的加权平均值。在注意力完美对角线极端情况下，attn@V只是V的精确副本。此处可见更纹理化模式：可见带状结构，特定token在注意力行的连续子序列上得分高，叠加与V相似的矩阵，因对角线较粗产生垂直遮挡。（旁注：长按或按住Control键单击显示可视化元素实际数值。）注意：位于中间层（5层），输入为中间表示而非原始token化文本，因此输入中模式发人深省——特别是强大垂直线条为特定嵌入位置，值在序列长段上统一高幅度，有时几乎充满。有趣的是，输入序列首向量独特，不仅打破高幅度列模式，且几乎每个位置携带非典型值。（旁注：此处未可视化，但该模式在多个样本输入上反复出现。）注意：可视化的是单个样本输入计算。实践中，每个头都有特征模式，能在相当多样本集合上一致表达，但查看任意激活可视化时，需记住输入完整分布可能以微妙方式影响激发的想法和直觉。建议直接探索动画！
4c注意力头有趣差异：继续通过演示展现研究模型了解其工作方式的有用性。这是GPT-2另一注意力头，行为模式与5层第4头大不相同——符合预期，毕竟位于模型非常不同部分。该头位于第一层0层的第2头：值得注意的点：注意力分布均匀。产生效果：将V的相对未加权平均值（或V的合适因果前缀）交给attn@V每一行；动画显示attn[i]@V向量-矩阵积波动很小，而非简单V缩小比例、逐渐揭示的副本。attn@V具有惊人垂直均匀性——嵌入大柱状区域中相同值模式在整个序列中持续存在。可将其看作每个token共享的属性。旁注：一方面，考虑到注意力分布非常均匀，可能期望attn@V具有一定一致性。但每一行由V的因果子序列而非整个序列构成——为何不会导致更多变化，就像沿序列向下移动时的渐进变形一样？答案必定在于V值分布的更微妙属性。最后，外投影后，该头输出在垂直方向上更加均匀。能强烈印象：该头传递的大部分信息由序列中每个token共享的属性组成。其输出投影权重构成强化这种直觉。不由得会想：该头产生的极其规则、高度结构化信息可能通过相对不那么奢华的计算手段获得。当然，这不是未经探索领域，但可视化计算信号的明确性和丰富性对于产生新想法和推理现有想法都非常有用。
5并行化注意力：上面5层第4头动画可视化了注意力头中6个矩阵乘法中的4个。它们被可视化为一条向量-矩阵积融合链，证实了从输入到输出的整个左结合链沿共享i轴呈层状，且可并行化。5a示例：沿i分区。为实践并行计算，可将输入沿i轴划分为块。工具中可通过指定轴划分为特定数量块（示例中为8，数字无特别之处）。可视化清晰显示每次并行计算都需要完整wQ（内投影）、K_t和V（注意力）及wO（外投影），因它们沿未分区维度与已分区矩阵邻接。5b示例：双重分区。展示沿多个轴分区示例。选择可视化近期创新成果Block Parallel Transformer（BPT），基于Flash Attention等研究（论文链接：https://arxiv.org/pdf/2305.19370.pdf）。首先，BPT如上所述沿i分区，并将此水平分区延伸至注意力层另一半边（FFN）。为完全解决上下文长度问题，向MHA添加第二个分区——注意力计算本身的分区（即沿Q@K_t的j轴分区）。这两个分区将注意力分成块构成的网格。可视化清楚显示：这种双重分区能有效解决上下文长度问题，因现在能以视觉划分注意力计算中每次出现的序列长度。第二次分区的”范围”：几何结构显示K和V的内投影计算可与核心双矩阵乘法一起分区。注意一个微妙细节：视觉暗示还可沿k并行化后续矩阵乘法attn@V，并以split-k风格对部分结果求和，从而并行化整个双重矩阵乘法。但sdpa（）中的逐行softmax增加了要求：在计算attn@V相应行之前，需将其所有分段归一化，这会在注意力计算与最终矩阵乘法之间添加逐行步骤。
6注意力层大小：众所周知，注意力层前一半（MHA）因二次复杂度计算需求高，但后一半(FFN)也有自身需求，因其隐藏维度宽度通常是模型嵌入维度宽度的4倍。可视化完整注意力层”生物量”有助于建立该层两半部分如何相互比较的直觉认识。6a可视化完整注意力层：下面是完整注意力层，前一半（MHA）位于后面，后一半(FFN)位于前面。同样，箭头指向计算方向。注：该可视化描绘的不是单个注意力头，而是显示未切片Q/K/V权重和围绕中心双重矩阵乘法的投影。当然，未忠实地可视化完整MHA运算，但目标是更清晰地了解该层两半的相对矩阵大小，而非每半相对计算量。（此外，此处权重使用随机值而非真实权重。）维度有所收缩以保证浏览器能轻松加载，但比例保持不变（来自NanoGPT small配置）：模型嵌入维度=192（原本768）、FFN嵌入维度=768（原本3072）、序列长度=256（原本1024），尽管序列长度对模型无根本性影响。（从视觉上看，序列长度变化将表现为输入叶片宽度变化，从而导致注意力中心大小和下游垂直平面高度变化。）6b可视化BPT分区层：简单回顾Blockwise Parallel Transformer，此处在整个注意力层语境中可视化BPT并行化方案（与上面一样省略了各个头）。特别要注意沿i（序列块）分区如何通过MHA和FFN两半边扩展。6c对FFN进行分区：可视化方法建议进行与上面描述正交的额外分区——在注意力层FFN半边，将双重矩阵乘法（attn_out@FFN_1）@FFN_2分开，首先沿j进行attn_out@FFN_1，然后沿k与FFN_2执行后续矩阵乘法。这种分区会切片两个FFN权重层，减少计算中每个参与组分的容量要求，但代价是部分结果的最终求和。下面是将此分区方法应用于未分区注意力层的样子：下面则是应用于以BPT方式分区层的情况：6d可视化一次一个token解码过程：在自回归式一次一个token解码过程中，查询向量由单个token构成。可在头脑中想象这种情况下的注意力层样子，很有启发性——单个嵌入行穿过巨大平铺的权重平面。除了强调与激活相比权重的巨大性外，这种观点还能让人想起K_t和V功能类似于6层MLP中动态生成的层，尽管MHA本身的mux/demux计算会使这种对应关系不精确。
7LoRA：近期LoRA论文《LoRA: Low-Rank Adaptation of Large Language Models》描述了一种高效微调技术，该技术基于这一思想：微调期间引入的权重δ是低秩的。根据论文，这”允许我们通过在适应过程中优化密集层变化的秩分解矩阵来间接地训练神经网络中的一些密集层…… 同时保持预训练权重处于冻结状态。”7a基本思想：简而言之，关键一步是训练权重矩阵因子而非矩阵本身：用I x K张量和K x J张量的矩阵乘法替代I x J权重张量，其中保证K为较小值。若K足够小，则尺寸方面可能有很大优势，但也有权衡：降低K会降低积可表达的秩。这里通过随机128×4左侧参数和4×128右侧参数的矩阵乘法示例说明尺寸节省与结果结构化影响——即128×128矩阵秩为4的分解。注意L@R中的垂直和水平模式：7b将LoRA应用于注意力头：LoRA将这种分解方法应用于微调过程，为每个权重张量创建要进行微调的低秩分解，并训练其因子，同时保持原始权重冻结；微调后，将每对低秩因子相乘，得到原始权重张量形状的矩阵，并将其添加到原始预训练权重张量中。下面可视化了一个注意力头，其权重张量wQ、wK_t、wV、wO被低秩分解wQ_A@wQ_B等替换。从视觉上看，因子矩阵呈现为沿风车叶片边缘的低栅栏