科学空间|Scientific Spaces

上一篇文章《MoE环游记：6、最优分配促均衡》中，我们通过求解如下最优分配问题来实现负载均衡
\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_j x_{i,j} = k,\quad \sum_i x_{i,j} = \frac{mk}{n}\end{equation}
其中$\sum_j x_{i,j} = k$表示每个Token恰好激活$k$个Expert，而$\sum_i x_{i,j} = mk/n$表示每个Expert恰好被激活$mk/n$次。然而，仔细思考就会发现，其实前者对训练和推理都不是必要的，我们真正需要的是后者，它意味着“平均来说每个Token激活$k$个Expert”以及每个Expert的负载均衡，这足以达成MoE的目标，所以本文考虑更简化的问题
\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_i x_{i,j} = \frac{mk}{n}\label{eq:target-dyn}\end{equation}

点击阅读全文...

我们知道，负载均衡（Load Balance）是MoE架构中基本且关键的一环，直接影响模型的效率和性能。本系列已经有两篇文章介绍了两种实现负载均衡的主流思路，分别是《MoE环游记：2、不患寡而患不均》介绍的经典方案Aux Loss，以及《MoE环游记：3、换个思路来分配》中的由DeepSeek提出的Loss-Free方案。两者各有所长，亦各有局限。

本文将探讨第三种思路：最优分配，它将负载均衡视为等式约束下的线性规划问题。从最终形式上看，它仍属于Loss-Free，但基于截然不同的原理，提供了更准确且无超参的更新方式。

方法回顾

现有的两种方法中，Aux Loss的思路相对朴素，核心是“哪里不稳罚哪里”，通过正则项对负载不均施加惩罚。然而，Aux Loss有两个问题：首先，惩罚系数不好调，过大会干扰主Loss的优化，过小则均衡效果差；其次，Aux Loss的背后是STE（Straight-Through Estimator），这意味着它的梯度是次优的，它可能会带来负载均衡以外的未知影响。

点击阅读全文...

在上一篇文章《MuP之上：1. 好模型的三个特征》中，我们提出了前向稳定性、依赖稳定性、更新稳定性这三个核心指标，并给出了相应的数学定义。同时，我们提出以它们是否满足$\Theta(1)$来刻画一个模型的好坏，这将作为我们后续分析和计算的理论基石。接下来，我们将会把它们跟最速下降思想结合，给每个参数定制“稳中求快”的更新规则。

\begin{align}
&\text{前向稳定性:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c1} \\[5pt]
&\text{依赖稳定性:}\quad\max_{\boldsymbol{x}_1,\boldsymbol{x}_2} \Vert \boldsymbol{f}(\boldsymbol{x}_1;\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x}_2;\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c2} \\[5pt]
&\text{更新稳定性:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega} + \Delta\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c3}
\end{align}

我们以线性层作为第一个例子，其结果对部分读者来说应该不陌生，它就是去年逐渐兴起的Muon优化器。当然，我们的目的并不是重新发现Muon，而是展示从第一性原理出发来设计模型和优化器的过程，为我们后续处理其他参数提供统一的方法论。

点击阅读全文...

最近笔者刷到论文《Why Adam Works Better with β1=β2: The Missing Gradient Scale Invariance Principle》，顾名思义，它声称Adam在$\beta_1=\beta_2$时表现更优。经同事提醒，去年论文《In Search of Adam's Secret Sauce》也表达了相同的观点。无独有偶，昨天刚出来的《The Effect of Mini-Batch Noise on the Implicit Bias of Adam》也有类似发现。

\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\
&\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\
&\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\
&\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\
&\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}})
\end{aligned}\right.\end{equation}

众多论文都指向了$\beta_1=\beta_2$，它有什么理论上的好处呢？本文我们来学习一下相关推导。

点击阅读全文...

认识比较久的老读者或许都知道，笔者算是一个比较坚定的“古法编程爱好者”——至今仍不用IDE、不用代码补全、编辑器只需语法高亮即可。就连科学空间的博文，都是笔者一个个字敲的HTML源码（当然这有些历史原因），还有 Cool Papers 的从前端到后端的整站代码，也都是笔者手敲的。

但即便是笔者这样的“老顽固”，也不得不承认：AI Agent对某些任务来说具有无与伦比的优势——有些任务你想手写，都有种无从下手的感觉。这篇文章就介绍一个经典例子，用一行代码翻译arXiv论文。

先说方法

首先在本地部署好 kimi-cli ，配置好使用 kimi-k2.5 ，然后还要装一个LaTeX编译环境（笔者用的是MacTeX），最后从arXiv下载论文源码，解压并进入到源码，执行

点击阅读全文...

从机器学习时代的数据白化预处理，到深度学习时代的BatchNorm、InstanceNorm、LayerNorm、RMSNorm等花样百出的Normalization方法，本质上都体现了我们对“各向同性（Isotropy）”的偏爱。为什么我们会倾向于各向同性的特征呢？它有什么实际上的好处呢？这个问题能找到很多答案，比如对齐尺度、减少冗余、去相关性等等，但多是流于表面的感觉。

近日，笔者在读论文《The Affine Divergence: Aligning Activation Updates Beyond Normalisation》时，悟到了该问题在优化视角下的一个新理解，个人认为它相对来说还是比较贴近本质的，所以写出来跟大家分享和讨论一下。

最速下降

我们从最简单的线性层出发
\begin{equation}\boldsymbol{Y} = \boldsymbol{X}\boldsymbol{W}\end{equation}

点击阅读全文...

时至今日，LDM（Latent Diffusion Models）依旧是扩散模型的主流范式。借助Encoder对原始图像进行高倍压缩，LDM能显著减少训练与推理的计算成本，同时还能降低训难度，可谓一举多得。然而，高倍压缩也意味着信息损失，而且“压缩、生成、解压缩”的流水线也少了些端到端的美感。因此，始终有一部分人执着于“回到像素空间”，希望让扩散模型直接在原始数据上完成生成。

本文要介绍的《Back to Basics: Let Denoising Generative Models Denoise》正是这一思路的新工作，它基于原始数据往往处于低维子流形这一事实，提出模型应预测数据而不是噪声，由此得到“JiT（Just image Transformers）”，显著地简化了像素空间的扩散模型架构。

信噪之比

毋庸置疑，当今扩散模型的“主力军”依然是LDM，即便是前段时间颇为热闹的RAE，也只是声称LDM的Encoder已经“过时”了，要给它换一个新的更强的Encoder，但依然没改变“先压缩后生成”这一模式。

点击阅读全文...

这段时间，相信很多读者已经刷到过Muon优化器的相关消息。实际上，Muon的提出时间大致是去年的10月份，由 Keller Jordan 在推特上提出，距今也不过一年多一点。然而，就在这一年里，Muon已经经历了百亿、千亿乃至万亿参数模型的训练考验，足以表明它是一个相当有竞争力的优化器。

如今，Muon已经内置在Torch、Keras等训练框架中，就连Megatron这样的大型框架也逐渐开始支持，这意味它已经获得了业界的普遍认可。不过，对于仅熟悉Adam的读者来说，如何快速有效地切换到Muon，可能依然是一件让人困惑的事情。所以，本文试图给出一个快速上手教程。

简要介绍

Muon的正式提出者是 Keller Jordan ，目前任职于OpenAI。开头说了，Muon最早发表在推特上，而直到现在，作者也只是多写了篇博客《Muon: An optimizer for hidden layers in neural networks》而不是一篇Paper，作者的观点是“是否写成Paper，跟优化器是否有效，没有任何关系^[原文]”。

点击阅读全文...