科学空间|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下载论文源码，解压并进入到源码，执行

点击阅读全文...

从《线性注意力简史：从模仿、创新到反哺》中我们可以看到，DeltaNet的并行形式涉及到了形如$(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-)^{-1}$的逆矩阵。近日读者 @Arch123 提出，通过实验可观察到该逆矩阵的元素总是在$[-1, 1]$内，问是否可以从数学上证实或证伪它。

在这篇文章中，我们将通过两种不同的方式证明这个结论是严格成立的。

问题描述

首先，我们准确地重述一下问题。设有矩阵$\boldsymbol{K}=[\boldsymbol{k}_1,\boldsymbol{k}_2,\cdots,\boldsymbol{k}_n]^{\top}\in\mathbb{R}^{n\times d}$，其中每个$\boldsymbol{k}_i\in\mathbb{R}^{d\times 1}$是模长不超过1的列向量，$\boldsymbol{M}\in\mathbb{R}^{n\times n}$是一个下三角的掩码矩阵，定义为
\begin{equation}M_{i,j} = \left\{\begin{aligned} &1, &i \geq j \\ &0, &i < j\end{aligned}\right.\end{equation}
$\boldsymbol{I}$是单位阵，$\boldsymbol{M}^- = \boldsymbol{M} - \boldsymbol{I}$。我们要证明的是：
\begin{equation}(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-)^{-1}\quad\in\quad [-1, 1]^{n\times n}\end{equation}

点击阅读全文...

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

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

最速下降

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

点击阅读全文...

在《让炼丹更科学一些（五）：基于梯度精调学习率》中，我们进入了基于梯度来调度学习率的新篇章。但上文末也提到，在推导动态梯度下终点损失的最优学习率时，我们遇到了证明上的困难，具体来说，我们基于变分法“猜”出来的最优学习率序列，代入结论中进行放缩验证会十分困难，因此别说最优解了，我们甚至无法判断这个序列是否是可行解。

而在本文中，我们将通过一个精妙的构造得到更精准的结论，从而解决这个问题。就证明过程来看，这一次的结论可能已经达到了无法改进的精度。这个突破依然出自论文《Optimal Linear Decay Learning Rate Schedules and Further Refinements》。

问题回顾

先重温一下之前的结论。上文末，我们得到了《让炼丹更科学一些（四）：新恒等式，新学习率》结论的一般版本：
\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{R^2}{2\eta_{1:T}} + \frac{1}{2}\sum_{t=1}^T\frac{\eta_t^2 G_t^2}{\eta_{\min(t+1, T):T}}\label{leq:last-2}\end{equation}

点击阅读全文...