2026-08-13
用范畴论为深度学习架构建一套形式化代数,把广播这类隐式约定变成可追踪的结构,完整刻画 DeepSeek-V3 全部组件,并配套可互操作的 Python/TypeScript 代码。
深度学习模型本质上是定义明确的数学函数,可我们描述模型架构的方式却惊人地随意:张量记号、框架专属代码、手绘示意图混着用。这些写法在局部管用,却给不出一个能系统推理的形式化对象。后果是,本该从架构定义里推导出来的东西,比如等价改写、高效实现、性能模型、跨框架复刻,基本靠人手推导。
作者把症结压到一个长期被忽视的概念上:广播(broadcasting)。SoftMax 写成 R^x 到 R^x,到底是逐行做还是对整张量做,取决于哪个轴被「摊开」、哪个是目标轴。广播信息一缺,数学含义就含糊,也就没法做严谨分析。而广播恰恰是把架构映射到并行硬件的关键:GPU kernel 按广播轴定义,FlashAttention 这类省带宽的 trick 也全靠把广播讲清楚。
作者的工具是范畴论(category theory),研究组合与抽象的数学。他们给深度学习架构搭了一个分层范畴,核心动作是把广播从隐式约定提升为可追踪的一等结构。具体技术是两套新范畴:一套管索引上的仿射步长变换(axis-stride),一套是数组广播范畴 Br,论文记作 [B;A],B 是基础数据类型,A 是数组轴。
关键装置叫 weave(编织)。一个 weave 给每根索引轴打上标签,标明它是「平铺」(被广播摊开的轴)还是「目标」(操作真正作用的轴),并把这套信息直接焊进操作本身。在这个底座上,作者把深度学习的积木一块块拼回去:转置、求和、外积、内积这类 Einstein 操作,用 reindexing 加虚线表示乘法、弯杯表示缩并,复刻 Penrose 图形记号;SoftMax 是多态算子,靠 weave 决定摊开到哪一轴;Dropout 被写成概率 Markov 核,底层范畴可换成概率范畴 Stoch;可学习层和归一化(RMSNorm)用 Para 构造把权重藏成一条侧入;嵌入是整数数据类型上的选择操作;卷积被拆成一步加法 reindexing 加一层线性层,连平移等变性都能顺着索引滑过去推导出来。
一个体现形式化价值的细节:Dropout 是非确定的随机操作,把一条随机值轴复制一遍,跟让随机操作在那些轴上并行各跑一次,结果并不相同。论文的代数如实反映了这件事(非确定操作上 Yoneda 滑动不成立),而随意的记号会把这种差别糊过去。
这是一篇基础设施论文,结果是能力演示,跑分不是它的目标。作者用这套框架把 DeepSeek-V3 的完整架构刻画了一遍:标准 transformer 主干、带复数 rotary 与潜在实部分量的多头潜在注意力、GeGLU 前馈、含门控的混合专家(MoE)。围绕这个表达式,他们配套实现了一组成果:
| 能力 | 说明 |
| 双语言代码包 | pyncd(Python)与 tsncd(TypeScript)镜像实现,JSON 互通 |
| 自动对齐 | 算子组合时自动匹配轴尺寸,如 qkmatmul @ softmax @ mask @ svmatmul,论文把这一步表述为一个保结构的函子 |
| 配置生成 | 对齐后剩余的自由轴即模型的自由度,可扫描出待赋值的配置项 |
| PyTorch 编译 | 代数表达式可直接生成可运行的 PyTorch 模块,作者强调 PyTorch 只是附带目标,同样能编到 TensorFlow 或 Triton |
| 超图改写 | 转成超图后可施加结合律、双函子性、对称性等代数规则 |
| 示意图生成 | TypeScript 端渲染出文中那一整套架构图 |
论文没有精度或加速数字,也不声称模型更快或更准。交付物是「能完整表达一个生产级 SOTA 模型,且这套表达机器可操纵」。
这篇的真正受众不是炼丹的人,而是做编译器、kernel 和框架的人,以及搞架构与硬件协同设计的人。线索在前作:同一作者 2025 年的 FlashAttention on a Napkin 已经证明,一旦广播被讲清楚,FlashAttention 及其性能模型是可以程序化推导出来的。本文把这件事的数学地基夯实,目标是让这类推导自动化、并推广到任意架构,而不必为每个注意力变体或每种新硬件重新手工推导一遍。
作者自己把最具杀伤力的能力划进了未来工作:自动推导底层 kernel(分块矩阵乘、注意力)和硬件感知性能模型,本文都没做,只铺了地基。Para 式的代数自动微分、量化误差的复合分析,也都只是设想。
更现实的限制:MoE 那套「稀疏轴」写法在数学上自洽,但作者明说它不足以直接计算,需要靠代数改写回到可执行的形态才能跑。PyTorch 编译也需要为 PyTorch 复杂的广播语义搭额外基础设施。最大的存疑在于,这篇的全部工程回报目前都依赖后续工具在这块地基上盖起来,地基本身还没有跑分能证明它比 einops、named tensor 这类现有写法带来可量化的收益。