[{"content":"Agenda 1️⃣ 什么是 VLA？ 2️⃣ VLA 发展：以 π 系列做主线 3️⃣ VLA模型的设计空间：分享VLANeXt 论文中的思考 4️⃣ 智驾 vs 具身：共性与讨论 什么是 VLA（视觉-语言-动作模型）？ VLA 代表 Vision-Language-Action（视觉-语言-动作） VLA 的核心思想是利用多模态大模型（MLLM）的 world knowledge and reasoning capabilities，直接产生机器人的动作输出。\n输入： 机器人的摄像头图像（Vision）,机器人自身状态(Robot State)和人类的文字指令（Language）。 处理： 基于强大的多模态大模型（MLLM）去理解场景和指令。 输出： 直接生成控制机器人关节或末端执行器的低级动作指令（Action）。 VLA 让机器人能够像人类一样，通过观察世界、听从指令，然后直接采取行动，打破了传统机器人需要复杂的手工编码和分层多子系统的限制。\nVLA的研究进展(pi系列)-以Physical Intelligence (π)为主线(关注核心组织和人 ） π系列发展历程和趋势 π\n阶段 核心架构与技术特征 关键突破点 π₀\n2024-10 MoT (Mixture of Transformer) + Flow Matching 奠定了 VLA 多模态训练的基础，利用连续流匹配生成动作。 FAST\n2025-01 Efficient Action Tokenization 引入离散动作标记，提升了实时的推理速度和执行效率。 π₀.₅\n2025-04 两者 action 建模结合\n混合任务训练\n层次化 VLM/VLA 架构 结合连续与离散建模，通过混合任务训练应对复杂、长程任务。 π₀.₆* 及后续\n2025-11 RL + Memory + Agentic 引入强化学习、长期记忆与自主决策能力，向机器人智能体迈进。 总结Physical Intelligence (π) VLA的进展 阶段 时间 版本 核心技术方向 π系列第一阶段 2024-10 ~ π₀ / π₀-FAST / π₀.₅ Flow Matching、Efficient Action Tokenization、分层设计 VLM + VLA、混合任务训练 π系列第二阶段 2025-11 ~ π₀.₆* 及后续 RL、Memory、Agentic VLA模型的设计空间（VLANeXT） VLANeXT\n设计VLA模型需要思考的问题： action怎么生成，VLM直接输出还是额外action expert网络？ action训练方法，Flow matching，VAE还是DDIM？ VLM和action expert的怎么连接？ propriocetion 自身状态输入给VLM还是action expert 多相机数据怎么输入给VLA，相机标定的内外参信息怎么喂给模型？图像的时序信息怎么处理？ 是否引入额外任务构建更稠密的loss？如预测下一帧图像，还是引入额外输出头？ VLANEXT中探讨的VLA设计空间 设计维度 决策组件 / 变量 设计选项 (Design Options) 实验结论摘要 (Ablation Insights) 最佳配方 (Final Recipe) 一、基础组件 (Foundational) 1. 策略模块设计 (Policy Architecture) A. 复用文本Token\nB. 独立策略头（浅层）\nC. 深度策略网络（MetaQuery） 浅层或复用架构无法充分解耦感知与动作；MetaQuery 风格能更好处理复杂动作。 采用深度独立策略模块（16 query tokens + 12层 Transformer） 2. 动作分块 (Action Chunking) A. 单步预测（size=1）\nB. 多步预测（size=8） 预测未来多步动作显著提高动作平滑性和任务成功率。 预测多步动作（Chunk size = 8） 3. 动作学习目标 (Action Objective) A. 离散分类（Binning / VQ-VAE）\nB. 回归（Regression）\nC. 扩散 / 流匹配（DDIM / Flow） 离散化损失空间细节；连续建模（Flow Matching / 回归）表达力更强。 使用连续动作建模（优先 Flow Matching 或回归） 4. VLM骨干容量 (VLM Backbone) 测试 LLaMA, PaliGemma, Qwen3-VL 等 骨干模型表征能力是基础，容量越大通常性能越好（需平衡资源）。 视算力选择最优开源骨干（如 Qwen3-VL） 5. VLM-策略连接 (VLM-Policy Link) A. 松散连接（Loose）\nB. 紧密连接（Tight）\nC. 软连接（Soft，引入 Query buffer） 软连接通过可学习 query 提供缓冲，适配性略优于紧密或松散耦合。 软连接策略（Soft Connection） 二、感知要素 (Perception) 6. 时序历史 (Temporal History) A. 仅当前帧\nB. 添加过去帧 添加过去帧反而降低性能，可能引入冗余或干扰信息。 避免冗余时序历史输入（专注当前帧） 7. 相机视角 (Camera Views) A. 单视角\nB. 多视角（第三人称 + 腕部） 多视角提供更全面空间信息，显著提升空间感知与操作精度。 多视角输入（第三人称 + 腕部相机） 8. 本体感觉注入 (Proprioception) A. 无本体 state\nB. VLM端注入\nC. 策略端注入\nD. 双端注入 本体 state 对闭环控制至关重要；VLM端注入效果优于策略端。 本体 state 条件于 VLM 端注入 三、动作建模 (Action Modeling) 9. 世界建模辅助 (World Modeling) 引入未来帧预测目标（如 Emu3.5） 性能有提升，但训练成本增加约3倍，性价比低。 不纳入最终配方 10. 时序预测视角 (Time-Series Perspective) 引入频域损失（DCT） 将动作视为时序预测问题，引入频域损失可低成本提升动作质量。 引入频域辅助损失（DCT） 应用场景-具身VS智驾 具身和智驾领域VLA论文对比\u0026ndash;\u0026gt;技术趋同 技术方向 具身 智驾 action token pi0-FAST、RT2、OpenVLA AutoVLA action expert（flow matching） pi0 DriveMoE action 建模两者结合 + 强化学习 RL pi0.5、pi0.6* Alpamaya 辅助 loss（3D 感知 / 世界模型 / 通才模型） 蚂蚁 LingBot-VLA DriveVLA-W0 / EMMA Nvidia Alpamaya智驾VLA Alpamayo-R1: Bridging Reasoning and Action Prediction for Generalizable Autonomous Driving in the Long Tail\n从Waymo End-to-End Driving Chanllenges看未来智驾的技术趋势 添加图片注释，不超过 140 字（可选） WOD-E2E: Waymo Open Dataset for End-to-End Driving in Challenging Long-tail Scenarios(https://arxiv.org/pdf/2510.26125)-4.4. Discussion of the Results\n1.MLLM-based Model (VLA)的好处？-\u0026gt;数据不易饱和\n2.强化学习对E2E驾驶能力的提升？\nopen discussion Q1. 智驾是否需要VLA吗 ？ a.智驾特殊场景和长尾场景是否需要Language，是否是帮助提升泛化性? b.交互能力对于驾驶是否必须?驾驶需要需要遵循人类指令？ c. 推理CoT对于智驾这类实时性高的场景是否必须？怎么trade-off? Q2. VLA vs WAM？ VLA (Vision-Language-Action, 视觉-语言-动作模型)\n瓶颈： VLA 模型通常缺乏对物理世界的深刻理解（没有物理常识）。它只是学到了动作和图像之间的“统计映射”。\nWAM (World Action Model, 世界动作模型)\n核心思想： WAM 是建立在 World Model (世界模型) 之上的动作模型。世界模型的核心能力是“预测未来”。它学习物理世界的运作规律（例如，物体掉落会向下，水杯翻了水会流出）。 当接收到指令和当前视觉时，WAM 不会直接生成动作。它首先利用其世界模型“想象”：如果我采取动作A，世界会变成什么样？采取动作B，世界又会变成什么样？通过在“想象”中模拟未来，WAM 可以选择最能达成目标的动作序列。 WAM 的优势： 具有物理常识，具有“预测”和“规划”能力。\ncosmos policy 添加图片注释，不超过 140 字（可选） 添加图片注释，不超过 140 字（可选）\nQ3. VLA 是终局吗？WAM 是终局吗？两者怎么融合？ 添加图片注释，不超过 140 字（可选）\n总结 π 系列VLA： Physical Intelligence 的 π 系列 VLA 模型。π0​，基于VLM+Action Expert架构，Flow matching结合流匹配技术，多步去噪生成aciton; π0.5​,混合任务训练，采用两种action token建模方式，通过层次化 VLM+VLA 架构提高了复杂任务泛化能力；π0.6*及后续​，引入强化学习与记忆，推动VLA模型具有自主规划能力的具身智能体（Agentic AI）演进。\nVLANeXT-VLA 的关键“设计空间”： VLA 模型的执行效率源自深度独立的策略网络、多步动作分块、以及频域辅助损失对连续动作建模的优化；最佳配方摒弃了降低性能的历史时序输入，专注于多视角感知和连续动作流预测。\n智驾与具身VLA： 两者皆是 Physical AI的应用，模型设计上具有极高的技术重合度。讨论 VLA是否必须？VLA还是WAM两者技术路线？\n更多细节 Physical Intelligence (π) VLANeXT分析) Nividia Alpamayo-R1 理想汽车MindVLA-o1分享 WOD-E2E E2E-Driving Chanllenges Cosmos Policy ","date":"2026-04-11T14:36:45+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/%E7%89%A9%E7%90%86ai-vla%E5%88%86%E4%BA%AB/","title":"物理AI-VLA分享"},{"content":"我们在本文中介绍 MultiHeadSelfAttention, MHA的实现\nMultiHeadSelfAttention 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 import torch import torch.nn as nn import math class MultiHeadSelfAttention(nn.Module): def __init__(self,nums_head=8,hidden_dim=512): super().__init__() self.hidden_dim = hidden_dim self.nums_head = nums_head self.head_dim = hidden_dim//nums_head self.q_proj = nn.Linear(hidden_dim,hidden_dim) self.k_proj = nn.Linear(hidden_dim,hidden_dim) self.v_proj = nn.Linear(hidden_dim,hidden_dim) self.o_proj = nn.Linear(hidden_dim,hidden_dim) self.dropout=nn.Dropout(0.1) def forward(self,x,attn_mask=None): batch_size,seq_len,_ = x.size() q=self.q_proj(x) k=self.k_proj(x) v=self.v_proj(x) # split q_state = q.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2) k_state = k.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2) v_state = v.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2) # q@k^t/sqrt(d) attn_score = q_state@k_state.transpose(-1,-2)/math.sqrt(self.head_dim) if attn_mask is not None: attn_score = attn_score.masked_fill(attn_mask ==0,-1e9) attn_weight = torch.softmax(attn_score,dim=-1) attn_weight = self.dropout(attn_weight) output_mid = attn_weight@v_state # concat output_mid = output_mid.transpose(-1,-2).contiguous() output_mid = output_mid.view(batch_size,seq_len,-1) output= self.o_proj(output_mid) return output x= torch.rand(3,2,512) attn_mask = torch.tensor([ [0,1], [0,0], [1,0], ]).unsqueeze(1).unsqueeze(2).expand(3,8,2,2) net = MultiHeadSelfAttention(8,512) net(x,attn_mask).shape ","date":"2026-04-08T17:44:26+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/multiheadselfattention/","title":"MultiHeadSelfAttention"},{"content":"安装go和hugo https://gentleostrich.github.io/posts/windows-11-%E5%88%9B%E5%BB%BA-hugo-%E7%BD%91%E7%AB%99/\n官方教程不允许使用 CMD、Windows PowerShell 执行各个命令，建议使用 PowerShell 执行命令。Anaconda Powershell Prompt 应该是基于 PowerShell 开发的，因此，本文的命令均在 Anaconda Powershell Prompt 中执行。\n安装 Hugo 安装 Hugo 前需要安装 Git、Go、Dart Sass。\n安装 Git\n无论是克隆 Hugo 的 GitHub 仓库，还是基于 Github Page 发布网站，都需要使用 Git。\n安装 Go\nHugo 是基于 Go 语言开发，Go 作为一种高效的编程语言，使得 Hugo 在速度上远超 Hexo 等静态网站搭建技术。\n安装 Dart Sass\nDart Sass 是 Hugo 开发的网页渲染插件，可以通过 Scoop 和 Chocolatey 两个 Windows 安装器安装，这里选择使用 Scoop 安装。\n安装 Scoop\n首先 cd 到 C:\\ 路径，在 C:\\ 路径下执行下述命令。\n1 2 Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser Invoke-RestMethod -Uri https://get.scoop.sh | Invoke-Expression 按照上述官方命令安装时出现了 PowerShell 禁止运行脚本的报错，使用 以下命令 替代上面的第一条命令解决该问题。\n1 2 3 Set-ExecutionPolicy RemoteSigned -Scope Process Set-ExecutionPolicy RemoteSigned -Scope CurrentUser Set-ExecutionPolicy RemoteSigned -Scope LocalMachine 安装 Scoop 后，可以执行下述命令安装 Dart Sass。\n1 scoop install sass 安装 Git、Go、Dart Sass 后\n在 Hugo GitHub 发布页 下载并解压 hugo_0.154.5_windows-amd64.zip 。 将解压得到的文件 hugo.exe 移动到安装路径，这里选择的安装路径是 D:\\Program Files\\hugo，并将该路径添加到系统路径变量中。 此时，重新打开 PowerShell 并执行 hugo version，若输出 Hugo 的版本号等信息，则表明 Hugo 安装成功。\n网站创建 成功安装 Hugo 后，就可以创建第一个网站了。\n1 2 3 4 5 6 hugo new site quickstart cd quickstart git init git submodule add https://github.com/theNewDynamic/gohugo-theme-ananke.git themes/ananke echo \u0026#34;theme = \u0026#39;ananke\u0026#39;\u0026#34; \u0026gt;\u0026gt; hugo.toml hugo server 报错 1：\n1 ERROR command error: failed to load config: \u0026#34;D:\\Program Files\\hugo\\quickstart\\hugo.toml:4:2\u0026#34;: unmarshal failed: toml: expected character = 直接将 echo \u0026quot;theme = 'ananke'\u0026quot; \u0026gt;\u0026gt; hugo.toml 复制到 PowerShell 中执行会出现格式错误，修改 hugo.toml\n报错 2：\n1 ERROR error building site: render: [en v1.0.0 guest] failed to render pages: render of \u0026#34;/tags\u0026#34; failed: \u0026#34;D:\\Program Files\\hugo\\quickstart\\themes\\ananke\\layouts\\baseof.html:26:15\u0026#34;: execute of template failed: template: taxonomy.html:26:15: executing \u0026#34;taxonomy.html\u0026#34; at \u0026lt;partials.Include\u0026gt;: error calling Include: \u0026#34;D:\\Program Files\\hugo\\quickstart\\themes\\ananke\\layouts\\_partials\\site-style.html:2:32\u0026#34;: execute of template failed: template: _partials/site-style.html:2:32: executing \u0026#34;_partials/site-style.html\u0026#34; at \u0026lt;.RelPermalink\u0026gt;: error calling RelPermalink: TOCSS: failed to transform \u0026#34;/ananke/css/main.css\u0026#34; (text/css). Check your Hugo installation; you need the extended version to build SCSS/SASS with transpiler set to \u0026#39;libsass\u0026#39;.: this feature is not available in your current Hugo version, see https://goo.gl/YMrWcn for more information 这是因为通过 hugo_0.154.5_windows-amd64.zip 安装的不是 extended version，不能对 anake 主题进行渲染。执行命令 scoop install hugo-extended 直接从 scoop 中安装 hugo-extended 并删除系统路径中的 hugo 路径即可。\n网站配置 通过目录下的 hugo.toml 配置网站，其中的 baseURL 设置为生产网站。\n1 2 3 4 baseURL = \u0026#39;https://example.org/\u0026#39; languageCode = \u0026#39;en-us\u0026#39; title = \u0026#39;My New Hugo Site\u0026#39; theme = \u0026#39;ananke\u0026#39; 内容添加 向网站中添加一个新文章：\n1 hugo new content content/posts/my-first-post.md 在 content/posts 目录下出现了 my-first-post.md 文件，下述是该文件中的内容。其中 draft=true 表示这个文章不会被发表到网站上。\n1 2 3 4 5 +++ title = \u0026#39;My First Post\u0026#39; date = 2024-01-14T07:07:07+01:00 draft = true +++ 通过下述命令查看含有 draft 文章的网站，当决定发表该文章后，将 draft 改为 false。\n1 2 hugo server --buildDrafts hugo server -D 网站发布 网站发布指的是 Hugo 在 public 目录下创建静态网站所需的全部文件（包括 HTML 文件、CSS 文件、Javascript 文件、图片得等等），命令如下：\n1 hugo 网站部署 基于 GitHub Page 部署网站。\n新建项目，项目名称为：\u0026lt;username\u0026gt;.github.io。这意味着这个项目是一个网站，每次 push 后都会进行 action 操作进行网站部署。\n将本地 Hugo 仓库与 GitHub 项目仓库进行 remote 关联，代码如下：\n1 2 3 git remote add origin https://github.com/GentleOstrich/gentleostrich.github.io.git git branch -M main git push -u origin main 报错：\n1 2 git push -u origin mainerror: src refspec main does not match any error: failed to push some refs to \u0026#39;https://github.com/GentleOstrich/gentleostrich.github.io.git\u0026#39; 原因是最开始本地项目从未 commit 过，所以没有不存在分支 main。执行下述命令即可解决：\n1 2 git commit -m \u0026#34;Initial commit\u0026#34; git push -u origin main 设置图片缓存位置\n在 hugo.toml 文件中添加下述内容：\n1 2 3 [caches] [caches.images] dir = \u0026#39;:cacheDir/images\u0026#39; 设置 action 文件\naction 文件的作用是指导 GitHub 如何部署网站的，首先创建相关 hugo.yaml 文件：\n1 2 mkdir -p .github/workflows touch .github/workflows/hugo.yaml 接下来在 hugo.yaml 文件中加入下述内容：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 name: Build and deploy on: push: branches: - main workflow_dispatch: permissions: contents: read pages: write id-token: write concurrency: group: pages cancel-in-progress: false defaults: run: shell: bash jobs: build: runs-on: ubuntu-latest env: DART_SASS_VERSION: 1.97.2 GO_VERSION: 1.25.5 HUGO_VERSION: 0.154.4 NODE_VERSION: 24.12.0 TZ: Europe/Oslo steps: - name: Checkout uses: actions/checkout@v6 with: submodules: recursive fetch-depth: 0 - name: Setup Go uses: actions/setup-go@v6 with: go-version: ${{ env.GO_VERSION }} cache: false - name: Setup Node.js uses: actions/setup-node@v6 with: node-version: ${{ env.NODE_VERSION }} - name: Setup Pages id: pages uses: actions/configure-pages@v5 - name: Create directory for user-specific executable files run: | mkdir -p \u0026#34;${HOME}/.local\u0026#34; - name: Install Dart Sass run: | curl -sLJO \u0026#34;https://github.com/sass/dart-sass/releases/download/${DART_SASS_VERSION}/dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz\u0026#34; tar -C \u0026#34;${HOME}/.local\u0026#34; -xf \u0026#34;dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz\u0026#34; rm \u0026#34;dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz\u0026#34; echo \u0026#34;${HOME}/.local/dart-sass\u0026#34; \u0026gt;\u0026gt; \u0026#34;${GITHUB_PATH}\u0026#34; - name: Install Hugo run: | curl -sLJO \u0026#34;https://github.com/gohugoio/hugo/releases/download/v${HUGO_VERSION}/hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz\u0026#34; mkdir \u0026#34;${HOME}/.local/hugo\u0026#34; tar -C \u0026#34;${HOME}/.local/hugo\u0026#34; -xf \u0026#34;hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz\u0026#34; rm \u0026#34;hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz\u0026#34; echo \u0026#34;${HOME}/.local/hugo\u0026#34; \u0026gt;\u0026gt; \u0026#34;${GITHUB_PATH}\u0026#34; - name: Verify installations run: | echo \u0026#34;Dart Sass: $(sass --version)\u0026#34; echo \u0026#34;Go: $(go version)\u0026#34; echo \u0026#34;Hugo: $(hugo version)\u0026#34; echo \u0026#34;Node.js: $(node --version)\u0026#34; - name: Install Node.js dependencies run: | [[ -f package-lock.json || -f npm-shrinkwrap.json ]] \u0026amp;\u0026amp; npm ci || true - name: Configure Git run: | git config core.quotepath false - name: Cache restore id: cache-restore uses: actions/cache/restore@v5 with: path: ${{ runner.temp }}/hugo_cache key: hugo-${{ github.run_id }} restore-keys: hugo- - name: Build the site run: | hugo \\ --gc \\ --minify \\ --baseURL \u0026#34;${{ steps.pages.outputs.base_url }}/\u0026#34; \\ --cacheDir \u0026#34;${{ runner.temp }}/hugo_cache\u0026#34; - name: Cache save id: cache-save uses: actions/cache/save@v5 with: path: ${{ runner.temp }}/hugo_cache key: ${{ steps.cache-restore.outputs.cache-primary-key }} - name: Upload artifact uses: actions/upload-pages-artifact@v4 with: path: ./public deploy: environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} runs-on: ubuntu-latest needs: build steps: - name: Deploy to GitHub Pages id: deployment uses: actions/deploy-pages@v4 至此，等待 GitHub 项目主页中 Action 显示绿色对勾后，就可以通过网址 \u0026lt;username\u0026gt;.github.io 访问网站啦（可能显示绿色对勾后立刻访问网站会出现 403 错误，只需要再等待一会儿就可以正常访问了）\nshortcode 可用于内嵌视频、图片等元素，分为 embedded custom inline 三类。\n插入一个 Youtube 视频（是一个 embedded）： `\n插入日期（是一个 inline）：今天是 。\n插入二维码：\n设置github actions部署 workflow https://github.com/jiangyigithub/ai.github.io/actions action workflow .github\\workflows\\deploy.yml\n设置github page的 hugo github pages https://github.com/jiangyigithub/ai.github.io/settings/pages\n更新文章 content\\post\n个人网址配置 github.io website baseurl = \u0026ldquo;https://github.com/jiangyigithub/ai.github.io\" config_default\\config.toml\n","date":"2026-04-01T17:44:26+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/%E5%88%A9%E7%94%A8hugo%E5%92%8Cgithub%E6%90%AD%E5%BB%BA%E4%B8%AA%E4%BA%BA%E7%BD%91%E5%9D%80/","title":"利用hugo和github搭建个人网址"},{"content":"我们在本文中介绍 strong scaling 和 weak scaling\nParallel Performance 给定一个算法，Parallel speedup 定义为在一个处理器上的运行时间 $T_S$ 和使用 $N$ 个处理器的运行时间 $T_P$ 之比，即\n$$ \\mathrm{speedup} = \\frac{T_S}{T_P} $$直觉上，使用的处理器越多，我们所花费的时间越少，并且理想情况下，我们的时间应该随 $N$ 增加而线性减少。但实际上，我们并不能实现 $N$ 倍加速，因此，我们使用 efficiency 来评估我们的并行效率：\n$$ \\mathrm{efficiency} = \\frac{\\mathrm{speedup}}{N} = \\frac{T_S}{T_P\\cdot N} $$我们可以将 efficiency 理解为处理器做有用功的时间比例。\nAmdahl\u0026rsquo;s Law 我们可以将一个算法或者程度的运行时间分为可并行部分以及不可并行部分，可并行部分可以运行在多个处理器上，而不可并行部分只能在一个处理器上运行。\n我们将算法运行的时间分为并行部分 $P$ 和串行部分 $S$, 我们有\n$$ T_s = T_SS+T_SP,\\quad S+P=1 $$当我们使用 $N$ 个处理器时，我们只能减少可并行部分的运行时间 $T_SP$， 并且由于我们使用了并行算法，我们还可能会产生额外的并行开销，记为 $K$, 则我们有\n$$ T_p=T_SS+T_S\\frac{P}{N}+K $$如果我们假设 $K\u003c","date":"2026-03-26T17:44:26+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/performance-and-scalability/","title":"Performance and Scalability"},{"content":"这一节我们讨论如何评估 (evaluate) 一个 policy. 要评估一个 policy $\\pi$, 我们需要求出对应的 value function $V^\\pi$.\nValue Function Monte Carlo 注意到\n$$ V^\\pi(s) = \\mathbb{E}^\\pi\\left[\\sum_{t=0}^{T-1}\\gamma^tr_t\\mid s_0=s\\right] $$我们可以利用 Monte Carlo (MC) 方法来进行估计，即从 $s_0=s$ 出发，独立随机采样 $N$ 条轨迹\n$$ \\{s, a_0^{(i)},r_0^{(i)},\\dots,a_{T^{(i)}-1}^{(i)},r_{T^{(i)}-1}^{(i)}, s^{(i)}_{T^{(i)}}\\}, i=1,\\dots,N $$然后我们使用样本平均来近似期望\n$$ V^\\pi(s) = \\mathbb{E}^\\pi\\left[\\sum_{t=0}^{T-1}\\gamma^tr_t\\mid s_0=s\\right]\\approx \\frac{1}{N}\\sum_{i=1}^N\\sum_{t=0}^{T^{(i)}-1}\\gamma^tr_t^{(i)} $$如果 $\\mathcal{S}$ 比较小，我们可以遍历 $s\\in\\mathcal{S}$ 来估计 $V^\\pi$.\n当 $\\mathcal{S}$ 非常大或者连续时，我们需要使用神经网络 $V_\\phi$ 来近似 $V^\\pi$,\n$$ \\min_{\\phi} \\mathcal{L}(\\phi) =\\mathbb{E}_{s\\sim p_0}\\left[\\frac12\\left(V_\\phi(s)-V^\\pi(s)\\right)^2\\right] $$对上面的目标函数求导得到\n$$ \\begin{align} \\nabla_\\phi \\mathcal{L}(\\phi) \u0026= \\mathbb{E}_{s\\sim p_0}\\left[\\left(V_\\phi(s)-V^\\pi(s)\\right)\\nabla_\\phi V_\\phi(s)\\right]\\\\ \u0026= \\mathbb{E}_{s\\sim p_0}\\left[\\left(V_\\phi(s)- \\mathbb{E}^\\pi\\left[\\sum_{t=0}^{T-1}\\gamma^tr_t\\mid s_0=s\\right]\\right)\\nabla_\\phi V_\\phi(s)\\right]\\\\ \u0026=\\mathbb{E}_{s\\sim p_0}\\left[\\mathbb{E}^\\pi\\left[\\left(V_\\phi(s)- \\sum_{t=0}^{T-1}\\gamma^tr_t\\right)\\nabla_\\phi V_\\phi(s)\\mid s_0=s\\right]\\right]\\\\ \u0026= \\mathbb{E}_{s\\sim p_0}^\\pi\\left[\\left(V_\\phi(s)- \\sum_{t=0}^{T-1}\\gamma^tr_t\\right)\\nabla_\\phi V_\\phi(s)\\mid s_0=s\\right]\\\\ \\end{align} $$这样，我们可以结合 SGD 与 MC 来近似 $V^\\pi$\nTemporal Difference Temporal difference (TD) learning 通过相邻两步的转换关系来近似函数。注意到，\n$$ V^\\pi(s) =\\mathbb{E}^\\pi[r_0 + \\gamma V^\\pi(s_1)\\mid s_0=s] $$与前面的方法类似，我们先采样然后使用 MC 来进行估计\n$$ V^\\pi(s) =\\mathbb{E}^\\pi[r_0 + \\gamma V^\\pi(s_1)\\mid s_0=s]\\approx \\frac{1}N\\sum_{i=1}^N\\left(r_0^{(i)} + \\gamma V^\\pi(s_1^{(i)})\\right) $$同样的，当 $\\mathcal{S}$ 比较大或者连续时，我们构造损失函数\n$$ \\min_{\\phi} \\mathcal{L}(\\phi) =\\mathbb{E}_{s\\sim p_0}\\left[\\frac12\\left(V_\\phi(s)-V^\\pi(s)\\right)^2\\right] $$对应的梯度为\n$$ \\nabla_\\phi \\mathcal{L}(\\phi) =\\mathbb{E}_{s\\sim p_0}^\\pi\\left[\\left(V_\\phi(s)- r_0-\\gamma V^\\pi(s_1)\\right)\\nabla_\\phi V_\\phi(s)\\mid s_0=s\\right] $$但是，这里存在的问题在于，我们使用了 $V^\\pi(s_1)$, 而这个值是未知的，因此，一个做法是使用当前的 value function $V_\\phi$ 来进行代替，即\n$$ \\nabla_\\phi \\mathcal{L}(\\phi) \\approx \\mathbb{E}_{s\\sim p_0}^\\pi\\left[\\left(V_\\phi(s)- r_0-\\gamma V_\\phi(s_1)\\right)\\nabla_\\phi V_\\phi(s)\\mid s_0=s\\right] $$这样，我们结合 TD 和 GSD 的算法就变成了\n接下来，我们需要关注如何计算 $g$, 直接通过 $g=\\left(V_\\phi(s)- r_0-\\gamma V_\\phi(s_1)\\right)\\nabla_\\phi V_\\phi(s)$ 来进行计算，这在数学上是没有问题的，但是从自动微分的角度不对，比如 Pytorch 对应的实现应该是\n1 2 3 4 5 pred = V_phi(s0) target = r + gamma * V_phi(s1) td_error = pred - target loss = 0.5 * td_error ** 2 loss.backward() 可以看到，我们实际上求的梯度是\n$$ \\nabla_\\phi\\left[\\frac12\\left(V_\\phi(s)-(r+\\gamma V_\\phi(s'))\\right)^2\\right]=(V_\\phi(s)-(r+\\gamma V_\\phi(s'))(\\nabla_\\phi V_\\phi(s)-\\nabla_\\phi V_\\phi(s')) $$这显然与 $g$ 不相等。\n为了解决这个问题，我们需要使用 stop-gradient 的技巧来避免 $V_\\phi(s')$ 参与反向传播，这个时候，我们的目标函数就变成了\n$$ \\frac12\\left(V_\\phi(s)-(r+\\gamma\\ \\mathrm{sg}[V_\\phi(s')])\\right)^2 $$其中 $\\mathrm{sg}[\\cdot]$ 是 stop-gradient operator, 满足\n$$ \\mathrm{sg}[x] = \\begin{cases} x \u0026\\text{forward pass}\\\\ 0 \u0026\\text{backward pass} \\end{cases} $$这样其梯度就是\n$$ \\nabla_\\phi\\left[\\frac12\\left(V_\\phi(s)-(r+\\gamma\\ \\mathrm{sg}[V_\\phi(s')])\\right)^2\\right]=(V_\\phi(s)-(r+\\gamma V_\\phi(s'))\\nabla_\\phi V_\\phi(s)=g $$对应的 python 代码为\n1 2 3 4 5 pred = V_phi(s0) target = r + gamma * V_phi(s1) td_error = pred - target.detach() # sop gradient operator loss = 0.5 * td_error ** 2 loss.backward() 最终，我们的算法实现为\n上述这种做法其实是 semi-gradient methods, 这类方法在参数更新时，只计算部分梯度的方法。虽然这种方法牺牲了严格的梯度下降性质，但是其能够保证更好的表现。\nK-step TD 我们可以进一步推广 TD 到多步的场景，注意到\n$$ \\begin{align} V^\\pi(s) \u0026=\\mathbb{E}^\\pi[r_0 + \\gamma V^\\pi(s_1)\\mid s_0=s]\\\\ \u0026=\\mathbb{E}^\\pi[r_0 + \\gamma \\mathbb{E}^\\pi[r_1 + \\gamma V^\\pi(s_1)\\mid s_1]\\mid s_0=s]\\\\ \u0026= \\mathbb{E}^\\pi[\\mathbb{E}^\\pi[r_0 + \\gamma r_1 + \\gamma^2 V^\\pi(s_2) \\mid s_1] \\mid s_0=s] \\end{align} $$由重期望定理（Law of Total Expectation）：\n$$ \\mathbb{E}^\\pi[\\mathbb{E}^\\pi[X \\mid s_1] \\mid s_0=s] = \\mathbb{E}^\\pi[X \\mid s_0=s] $$因此：\n$$ V^\\pi(s) = \\mathbb{E}^\\pi[r_0 + \\gamma r_1 + \\gamma^2 V^\\pi(s_2) \\mid s_0=s] $$重复这个过程 k 次,得到 k 步展开：\n$$ \\begin{align} V^\\pi(s) \u0026= \\mathbb{E}^\\pi\\left[\\sum_{i=0}^{k-1} \\gamma^i r_i + \\gamma^k V^\\pi(s_k) \\mid s_0=s\\right] \\end{align} $$现在，我们可以基于 k-step transition 构建目标函数\n$$ \\min_{\\phi} \\mathcal{L}(\\phi) =\\mathbb{E}_{s\\sim p_0}\\left[\\frac12\\left(V_\\phi(s)-\\left(\\sum_{i=0}^{k-1} \\gamma^i r_i + \\gamma^k V^\\pi(s_k)\\right)\\right)^2\\right] $$使用前面分析的方法，我们就可以写出类似的算法\n一般来说，我们会令 $k=5$.\nMC v.s. TD TD 本质上一种 bootstrapping, 当需要使用 $V^\\pi$ 时，我们使用 $V_\\phi$ 来代替。两者对比如下\nMC 需要完整的一次采样，但是 TD 可以通过 bootstrap 提高采样效率。 当 $V_\\phi$ 与 $V^\\pi$ 差距较大时，boostraping 会导致训练不稳定，而 MC 则相对稳定 Q-function 对于 Q-function, 我们也可以设计类似的算法。\n对于 MC, 我们有\n$$ Q^{\\pi}(s,a) = \\mathbb{E}^{\\pi}\\left[\\sum_{t=0}^{T-1}\\gamma^tr_t\\mid s_0=s, a_0=a\\right]\\approx \\frac1N\\sum_{i=1}^N\\sum_{t=0}^{T^{(i)}-1}\\gamma^tr_t^{(i)} $$当使用模型来近似时，我们的目标函数为\n$$ \\min_{\\phi} \\mathcal{L}(\\phi) =\\mathbb{E}_{s\\sim p_0, a\\sim\\pi(\\cdot\\mid s)}\\left[\\frac12\\left(Q_\\phi(s,a)-Q^\\pi(s,a)\\right)^2\\right] $$对应的梯度\n$$ \\nabla_\\phi \\mathcal{L}(\\phi)= \\mathbb{E}_{s\\sim p_0}^\\pi\\left[\\left(Q_\\phi(s_0,a_0)-\\sum_{t=0}^{T-1}\\gamma^tr_t\\right)\\nabla_\\phi Q_\\phi(s_0,a_0)\\right] $$下面是对应的算法\nMC\n1-step Q-function TD learning\nk-step Q-function TD learning\nConclusion 在本节中，我们介绍了给定 policy $\\pi$, 我们如何求解对应的 value function 和 Q-function.\n","date":"2026-03-18T17:43:57+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/rl-series-3-policy-evaluation/","title":"(RL series 3) Policy evaluation"},{"content":"Bellman Equations 本节中，我们将要定义 value function, Q-function 以及这两个函数与 policy 之间的关系，这是我们介绍不同 RL 算法的基础。这一节需要使用到 不动点定理.\n我们定义 Value function 如下：\n$$ V^{\\pi}(s) = \\mathbb{E}^{\\pi}[G_0\\mid s_0=s] $$value function 的具体含义为：agent 从当前状态 $s$ 出发，一直遵循当前策略 $\\pi$， 最后获取到的 expected discounted return.\nstate-action value function, 或者 Q-value function 定义如下：\n$$ Q^{\\pi}(s,a) = \\mathbb{E}^{\\pi}[G_0\\mid s_0=s, a_0=a] $$其具体含义为：agent 从当前状态 $s$ 出发，执行 action $a$, 再遵循当前策略 $\\pi$, 最后获取到的 expected discounted return.\n注意 这里为了方便，我们令 $V^{\\pi}(\\langle term\\rangle)=0$, $Q^{\\pi}(\\langle term\\rangle,a)=0,\\forall a\\in\\mathcal{A}$.\n由全概率公式，我们有\n$$ \\mathbb{E}[G_0\\mid s_0=s] = \\sum_{a\\in\\mathcal{A}}\\mathbb{E}\\left[G_0\\mid s_t=s, a_0=a\\right]\\pi(a\\mid s_0=s) $$即\n$$ \\boxed{V^{\\pi}(s) = \\mathbb{E}_{a_0\\sim \\pi(\\cdot\\mid s_0)}[Q^{\\pi}(s,a)]} $$一般来说，我们会使用 value function 和 Q-value function 的如下递推性质：\n$$ \\boxed{ \\begin{aligned}V^\\pi(s) \u0026= \\mathbb{E}_{a_0\\sim \\pi(\\cdot\\mid s),\\ (r_0,s_0)\\sim p(\\cdot,\\cdot\\mid s,a_0)}[r_0 + \\gamma V^\\pi(s_1)\\mid s_0=s]\\\\ Q^\\pi(s,a) \u0026= \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a),\\ a'\\sim \\pi(\\cdot\\mid s')}[r + \\gamma Q^\\pi(s',a')\\mid s,a] \\end{aligned}} $$证明需要利用 Markov property:\n$$ \\begin{aligned} V^{\\pi}(s) \u0026= \\mathbb{E}^{\\pi}\\left[r_0 + \\gamma\\sum_{t=1}^{T-1}\\gamma^{t-1}r_t\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi}[r_0+\\gamma G_1\\mid s_0=s]\\\\ \u0026= \\mathbb{E}_{a_0\\sim \\pi(\\cdot\\mid s,\\ (r_0,s_0)\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[\\mathbb{E}^\\pi[r_0 + \\gamma G_1\\mid s_0,a_0,r_0,s_1]\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}_{a_0\\sim \\pi(\\cdot\\mid s,\\ (r_0,s_0)\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r_0 + \\gamma \\mathbb{E}^\\pi[G_1\\mid s_1]\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}_{a_0\\sim \\pi(\\cdot\\mid s,\\ (r_0,s_0)\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r_0 + \\gamma V^\\pi[s_1)\\mid s_0=s\\right] \\end{aligned} $$对于 $Q^\\pi(s,a)$ 的推导同理。\n我们接下来介绍 RL 的核心：Bellman equation Theorem, 其阐述了策略 $\\pi$ 与 value function 所满足的充分必要条件。\nTheorem\n令 $\\pi$ 为一个策略，假设 $\\gamma\\in(0,1)$, $|\\mathcal{S}|\u003c\\infty$ 以及 $|r|\\leq R\u003c\\infty, \\mathrm{a.s.}$. 那么 $\\pi$ 对应的 value function $V^\\pi:\\mathcal{S}^+\\to\\mathbb{R}$ 存在，且满足 Bellman equation:\n$$ \\boxed{ V^\\pi(s) = \\mathbb{E}_{a\\sim \\pi(\\cdot\\mid s),\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V(s')\\mid s\\right] } $$ 反之，如果存在函数 $V:\\mathcal{S}^+\\to\\mathbb{R}$ 满足 Bellman equation, 则 $V=V^\\pi$.\n证明\n由 $V^\\pi$ 定义，我们有\n$$ \\left|\\sum_{t=0}^{\\infty}\\gamma^tr_t\\right|\\leq \\sum_{t=0}^\\infty \\gamma^t R=\\frac{R}{1-\\gamma}\u003c\\infty $$因此，$\\sum_{t=0}^{\\infty}\\gamma^tr_t$ 是一个绝对收敛的序列，从而其期望存在且有界。\n对于 $V^\\pi(s)=\\mathbb{E}^{\\pi}[G_0\\mid s_0=s]$ , 我们有\n$$ \\begin{aligned} V^\\pi(s)\u0026=\\mathbb{E}^{\\pi}\\left[G_0\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi}\\left[r_0+\\gamma G_1\\mid s_0=s\\right]\\\\ \u0026=\\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[\\mathbb{E}^\\pi\\left[r_0+\\gamma G_1\\mid s_0,a_0,r_0,s'\\right]\\right]\\\\ \u0026= \\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r_0+\\gamma\\mathbb{E}^\\pi\\left[ G_1\\mid s_0,a_0,r_0,s'\\right]\\right]\\\\ \u0026= \\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r_0+\\gamma\\mathbb{E}^\\pi\\left[ G_1\\mid s_1\\right]\\right]\\\\ \u0026=\\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s'), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r_0+V^{\\pi}(s')\\right] \\end{aligned} $$其中第 5 个等式为 Markov property.\n最后，我们证明方程解的唯一性，假设存在函数 $V:\\mathcal{S}^+\\to\\mathbb{R}$ 满足 Bellman equation. 我们定义 Bellman 算子 $\\mathcal{T}^\\pi$ 为\n$$ (\\mathcal{T}^\\pi V)(s) := \\mathbb{E}_{a\\sim \\pi(\\cdot\\mid s),\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V(s')\\mid s\\right] $$我们证明该算子是一个 contraction mapping, 考虑范数 $\\|V\\|_\\infty = \\max_{s\\in\\mathcal{S}} |V(s)|$, 我们有\n$$ \\begin{aligned} \\left|(\\mathcal{T}^\\pi V_1)(s)-(\\mathcal{T}^\\pi V_2)(s)\\right| \u0026= \\left|\\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s'), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[\\gamma V_1^{\\pi}(s')-\\gamma V_2^{\\pi}(s')\\right]\\right|\\\\ \u0026\\leq \\gamma \\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s'), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left|V_1^{\\pi}(s')-V_2^{\\pi}(s')\\right|\\\\ \u0026\\leq \\gamma \\max_{s''\\in\\mathcal{S}}\\left|V_1(s'')-V_2(s'')\\right|\\\\ \u0026= \\gamma \\left\\|V_1-V_2\\right\\|_\\infty \\end{aligned} $$上式对于任意 $s\\in\\mathcal{S}$ 都成立，因此\n$$ \\left\\|\\mathcal{T}^\\pi V_1-\\mathcal{T}^\\pi V_2\\right\\|_\\infty \\leq \\gamma \\left\\|V_1-V_2\\right\\|_\\infty $$由于 $\\gamma \u003c 1$, 从而 $\\mathcal{T}^\\pi$ 是一个 contraction mapping.\n根据不动点定理，已知 $V^\\pi$ 是一个不动点，而不动点唯一，因此我们有 $V=V^\\pi$. $\\blacksquare$\n同理，我们可以推导出关于 Q-function 的 Bellman equation:\nTheorem\n令 $\\pi$ 为一个策略，假设 $\\gamma\\in(0,1)$, $|\\mathcal{S}|\u003c\\infty$ $|\\mathcal{A}|\u003c\\infty$ 以及 $|r|\\leq R\u003c\\infty, \\mathrm{a.s.}$. 那么 $\\pi$ 对应的 Q-function $Q^\\pi:\\mathcal{S}^+\\times\\mathcal{A}\\to\\mathbb{R}$ 存在，且满足 Bellman equation:\n$$ \\boxed{ Q^\\pi(s, a) = \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a),\\ a'\\sim \\pi(\\cdot\\mid s')}\\left[r+\\gamma Q(s', a')\\mid s, a\\right] } $$ 反之，如果存在函数 $Q:\\mathcal{S}^+\\times\\mathcal{A}\\to\\mathbb{R}$ 满足 Bellman equation, 则 $Q=Q^\\pi$.\n证明\n证明与 value function 的证明基本类似，我们这里略过。\nBellman Optimal Equation 前面介绍了针对一般策略的 Bellman equation, 特别地，对于我们的目标最优策略，我们也可以而推导出对应的 Bellman equation.\n首先我们定义最优策略如下\nDefinition\n如果策略 $\\pi^*$ 满足\n$$ V^{\\pi^*}(s)\\geq V^{\\pi}(s), \\forall s\\in\\mathcal{S}^+, \\forall \\pi $$ 则我们称策略 $\\pi^*$ 是 optimal policy. 对应的 $V^{\\pi^*}$ 和 $Q^{\\pi^*}$ 分别称之为 optimal value function 以及 optimal Q-function, 我们简记为 $V^*=V^{\\pi^*}$, $Q^*=Q^{\\pi^*}$. 注意 optimal policy 与状态 $s$ 无关，并且 optimal policy 是不唯一的，但是所有的 optimal policy 对应的 optimal value function 是同一个。 $V^*$ 与 $Q^*$ 之间存在如下关系\n$$ V^*(s) = \\max_{a\\sim\\pi^*} Q^*(s, a) $$证明\n我们先证明左边小于右边，再证明右边小于左边。\n$$ V^*(s) =\\mathbb{E}_{a_0\\sim \\pi^*(\\cdot\\mid s_0)}[Q^*(s,a)]\\leq \\mathbb{E}_{a_0\\sim \\pi^*(\\cdot\\mid s_0)}[\\max_{a\\sim\\pi^*}Q^*(s,a)] = \\max_{a\\sim\\pi^*} Q^*(s, a) $$其次，令 $a^*\\in\\arg\\max_{a}Q^*(s,a)$, 令策略 $\\pi'$ 为确定性策略 $\\pi'(a^*\\mid s)=1$, 则\n$$ \\max_aQ^*(s,a^*) = Q^{\\pi'}(s,a^*)=V^{\\pi'}(s)\\leq V^*(s) $$这样，我们就证明了上面的等式。$\\blacksquare$\nTheorem\n假设 $\\gamma\\in(0,1)$, $|\\mathcal{S}|\u003c\\infty$ $|\\mathcal{A}|\u003c\\infty$ 以及 $|r|\\leq R\u003c\\infty, \\mathrm{a.s.}$. 那么 optimal value function $V^\\pi:\\mathcal{S}^+\\to\\mathbb{R}$ 存在，且满足 Bellman optimality equation:\n$$ \\boxed{ V^*(s) = \\max_{a\\in\\mathcal{A}}\\mathbb{E}_{\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r+\\gamma V^*(s')\\mid s, a\\right] } $$ 反之，如果存在函数 $V:\\mathcal{S}^+\\to\\mathbb{R}$ 满足 Bellman optimality equation, 则 $V=V^*$, 最终\n$$ \\pi^*(s) = \\arg\\max_{a\\in\\mathcal{A}}\\mathbb{E}_{\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r+\\gamma V^*(s')\\mid s, a\\right]=\\arg\\max_{a\\in\\mathcal{A}} Q^*(s,a) $$ 是一个 optimal deterministic policy.\n证明\n我们首先证明存在性，我们定义\n$$ (\\mathcal{T}^* V)(s) := \\max_{a\\in\\mathcal{A}}\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V(s')\\mid s, a\\right] $$我们证明 $\\mathcal{T}^*V$ 是一个 contraction mapping.\n$$ \\begin{aligned} \\left|(\\mathcal{T}^* V_1)(s)-(\\mathcal{T}^* V_2)(s)\\right| \u0026= \\left|\\max_{a\\in\\mathcal{A}}\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V_1(s')\\mid s, a\\right]- \\max_{a\\in\\mathcal{A}}\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V_2(s')\\mid s, a\\right]\\right|\\\\ \u0026\\leq \\max_{a\\in\\mathcal{A}}\\left|\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V_1(s')\\mid s, a\\right]- \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V_2(s')\\mid s, a\\right]\\right|\\\\ \u0026= \\max_{a\\in\\mathcal{A}}\\left|\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[\\gamma V_1(s')-V_2(s')\\mid s, a\\right]- \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V_2(s')\\mid s, a\\right]\\right|\\\\ \u0026\\leq \\gamma \\max_{a\\in\\mathcal{A}}\\left|\\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s'), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[V_1^{\\pi}(s')-V_2^{\\pi}(s')\\mid s,a\\right]\\right|\\\\ \u0026\\leq \\gamma \\max_{a\\in\\mathcal{A}}\\mathbb{E}_{a\\sim\\pi(\\cdot\\mid s'), (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[\\left|V_1^{\\pi}(s')-V_2^{\\pi}(s')\\right|\\mid s,a\\right]\\\\ \u0026\\leq \\gamma \\max_{a\\in\\mathcal{A}}\\max_{s''\\in\\mathcal{S}}\\left|V_1(s'')-V_2(s'')\\right|\\\\ \u0026= \\gamma \\left\\|V_1-V_2\\right\\|_\\infty \\end{aligned} $$这里第一个不等式我们使用了结论\n$$ \\left|\\max_s v(s)-\\max_s u(s)\\right|\\leq \\max_s|u(s)-v(s)| $$上式对于任意 $s\\in\\mathcal{S}$ 都成立，因此\n$$ \\left\\|\\mathcal{T}^* V_1-\\mathcal{T}^* V_2\\right\\|_\\infty \\leq \\gamma \\left\\|V_1-V_2\\right\\|_\\infty $$由于 $\\gamma \u003c 1$, 从而 $\\mathcal{T}^*$ 是一个 contraction mapping.\n根据不动点定理，$\\mathcal{T}^*$ 存在一个不动点，我们将其记为 $V^*$ .\n接下来，我们定义 $\\pi^*$ 为\n$$ \\pi^*(s) = \\arg\\max_{a\\in\\mathcal{A}}\\mathbb{E}_{\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r+\\gamma V^*(s')\\mid s, a\\right] $$此时，我们有\n$$ \\begin{aligned} V^*(s) \u0026= \\max_{a\\in\\mathcal{A}}\\mathbb{E}_{\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r+\\gamma V^*(s')\\mid s, a\\right]\\\\ \u0026= \\mathbb{E}_{a\\sim\\pi^*(s),\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0)}\\left[r+\\gamma V^*(s')\\mid s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma V^*(s_1)\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma \\mathbb{E}^{\\pi^*}\\left[r_1+\\gamma V^*(s_2)\\mid s_0=s_1\\right]\\\\\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma \\mathbb{E}^{\\pi^*}\\left[r_1+\\gamma V^*(s_2)\\mid s_0=s_1, r_0\\right]\\\\\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[\\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma (r_1+\\gamma V^*(s_2))\\mid s_0=s_1, r_0\\right]\\\\\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma r_1+\\gamma^2 V^*(s_2)\\mid s_0=s\\right]\\\\ \u0026= \\mathbb{E}^{\\pi^*}\\left[r_0+\\gamma r_1+\\gamma^2r_2+\\cdots\\mid s_0=s\\right]\\\\ \u0026= V^{\\pi^*}(s) \\end{aligned} $$即 $V^*=V^{\\pi^*}$.\n现在我们证明 $\\pi^*$ 是最优策略，令 $\\pi$ 为任意一个策略，我们有\n$$ V^\\pi \\leq V^*= \\mathcal{T}^*(V^\\pi)\\leq ( \\mathcal{T}^*)^2(V^\\pi) \\leq\\cdots\\leq (\\mathcal{T}^*)^k(V^\\pi) \\to V^*=V^{\\pi^*} $$这里第一个等式是 Bellman equation, 第一个不等式是由于下面的 Lemma,最后的极限使用了不动点定理。因此我们有 $V^\\pi\\leq V^*$, 从而 $\\pi^*$ 是最优策略，且 $V^*$ 是对应的 optimal value function. $\\blacksquare$\nLemma\n令 $\\pi$ 为一个策略， $\\mathcal{T}^\\pi$ 和 $\\mathcal{T}^*$ 分别是 Bellman operator 和 Bellman optimality operator. 对任意 $V:\\mathcal{S}^+\\to\\mathbb{R}$, 我们有\n$$ \\mathcal{T}^\\pi(V) \\leq \\mathcal{T}^*(V) $$ 进一步，对任意 $U:\\mathcal{S}^+\\to\\mathbb{R}$, $V:\\mathcal{S}^+\\to\\mathbb{R}$, 如果 $U\\leq V$, 则\n$$ \\mathcal{T}^*(U) \\leq \\mathcal{T}^*(V) $$证明\n$$ \\begin{aligned} \\mathcal{T}^\\pi(V)(s) \u0026= \\mathbb{E}_{a\\sim \\pi(\\cdot\\mid s),\\ (r,s')\\sim p(\\cdot,\\cdot\\mid s,a_0))}\\left[r+\\gamma V(s')\\mid s\\right]\\\\ \u0026\\leq \\mathbb{E}_{a\\sim \\pi(\\cdot\\mid s)}\\left[\\max_{a'} \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a')}\\left[r+\\gamma V(s')\\mid s\\right]\\right]\\\\ \u0026= \\max_{a'} \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a')}\\left[r+\\gamma V(s')\\mid s\\right]\\\\ \u0026= \\mathcal{T}^*(V) \\end{aligned} $$其次，令 $a^*\\in\\arg\\max_{a}\\mathbb{E}[r+\\gamma U(s')]$, 那么\n$$ \\begin{aligned} \\mathcal{T}^*(U)(s) \u0026= \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a'))}\\left[r+\\gamma U(s')\\mid s\\right]\\\\ \u0026\\leq \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a'))}\\left[r+\\gamma V(s')\\mid s\\right]\\\\ \u0026\\leq \\max_a\\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a'))}\\left[r+\\gamma V(s')\\mid s\\right]\\\\ \u0026= \\mathcal{T}^*(V) \\end{aligned} $$证毕。 $\\blacksquare$\n对于 Q-function, 我们也有类似结论\nTheorem\n假设 $\\gamma\\in(0,1)$, $|\\mathcal{S}|\u003c\\infty$ $|\\mathcal{A}|\u003c\\infty$ 以及 $|r|\\leq R\u003c\\infty, \\mathrm{a.s.}$. 那么 optimal Q-function $Q^*:\\mathcal{S}^+\\times\\mathcal{A}\\to\\mathbb{R}$ 存在，且满足 Bellman optimality equation:\n$$ \\boxed{ Q^*(s, a) = \\mathbb{E}_{(r,s')\\sim p(\\cdot,\\cdot\\mid s,a)}\\left[r+\\gamma \\max_{a'\\in\\mathcal{A}}Q(s', a')\\mid s, a\\right] } $$ 反之，如果存在函数 $Q:\\mathcal{S}^+\\times\\mathcal{A}\\to\\mathbb{R}$ 满足 Bellman optimality equation, 则 $Q=Q^*$.\n证明\n证明与 value function 类似，我们这里略过。\n","date":"2026-03-18T17:39:45+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/rl-series-2-bellman-equation/","title":"(RL series 2) Bellman Equation"},{"content":"Introduction RL 的基本思想是通过与环境进行交互，获取奖励来进行学习\nRL 的定义如下\n强化学习是一个通过构建可以与环境进行交互的 agent 来解决控制和决策任务的学习框架，交互的方式为 agent 执行 action, 然后环境给予奖励\nRL 的执行过程如下所示\nRL Basics RL 的数学建模依赖于 Markov decision process, 下面我们先介绍相关概念\nDefinition 一个 Markov decision process (MDP) 包含如下模块：\n时间 $t=0,1,\\dots,T$. 状态 $s_t\\in\\mathcal{S}\\cup\\{\\langle term\\rangle\\}$, 其中 $\\mathcal{S}$ 是状态空间 (state space) 动作 $a_t\\in\\mathcal{A}$ , 其中 $\\mathcal{A}$ 是动作空间 (action space) 奖励 $r_t\\in\\mathbb{R}$, 环境对 agent 当前 action $a_t$ 给予的反馈 终止时间 $T$, 额外定义 $s_T=\\langle term\\rangle$ 为终止状态 初始状态分布 $s_o\\sim p_0$. 转换概率 (transition probability): $r_t,s_{t+1}\\sim p(\\cdot,\\cdot\\mid s_t,a_t)$. 轨迹 $\\tau=(s_0,a_0,r_0,s_1,a_1,r_1,\\dots,s_{T-1},a_{T-1}, r_{T-1}, s_T)$. Markov property. $t+1$ 时刻的状态仅与 $t$ 时刻的状态与动作相关，即 $p(s_{t+1}\\mid s_t,a_t,\\dots,a_0,a_0)=p(s_{t+1}\\mid s_t,a_t)$ 以及 $p(r_{t}\\mid s_t,a_t,\\dots,s_0,a_t)=p(r_t\\mid s_t,a_t)$. 我们会对一些表达式进行简化：\n当 $r_t$ 完全由 $(s_t,a_t)$ 决定时，我们记为 $r_t=r(s_t,a_t)$ 我们假设状态转移函数是平稳的 (stationary), 即 $p_t(r,s'\\mid s,a)=p(r,s'\\mid s,a)$。 $a_t$ 通常由一个策略 $\\pi$ 决定， 当 $\\pi$ 完全由 $s_t$ 决定时，我们记为 $a_t=\\pi(s_t)$, 反之，$a_t$ 从 $\\pi$ 中采样得到，即 $a_t\\sim\\pi(\\cdot\\mid s_t)$. 一般我们会使用一个神经网络来表示 $\\pi$, 即 $\\pi=\\pi_\\theta$, 这里 $\\theta$ 就是神经网络的参数。 state and observation\n通常来说，state 包含了当前进行决策所需要的所有信息（环境信息，历史信息，agent 自身信息），这种情况下 MDP 的 Markov 性质成立。实际上 agent 获取的可能只是一部分信息，即 agent 获取的为 observation, 这种情况下 MDP 特化为 Partially Observable Markov Decision Process (POMDP), 此时我们的 Markov 性质就不再对 observation 成立，我们的问题也变的更加困难。\nTakeaway State 表示整个环境的完整表示，而 observation 则是环境的部分表示\n对于 LLM 来说，我们可以获取历史状态的所有信息，因此我们可以将 RL for LLM 视作一个 MDP 问题。\naction space\naction space $\\mathcal{A}$ 可以是连续的，比如上下左右四个方向，也可以是连续的，比如方向盘的角度\nepisodic and continuing\n根据任务的终止时间 $T$，我们可以将任务分为 episodic task 和 continuing task, 当 $T=\\infty$ 时， 我们的任务称为 continual task，比如游戏, 否则称之为 episodic task, 比如股票预测.\n除了 POMDP 之外，MDP 也有一些其他的推广形式：\nnon-stationary dynamics, 即不同时间步的状态转移函数不同 pre-determined terminal time $T$. 终止时间指定，因为 $p_{T-1}(\\langle term\\rangle\\mid s_{T-1}, a_{T-1})=1$ 且 $p_{t+1}(\\langle term\\rangle\\mid s_{t}, a_{t})=0$, $\\forall t\\in[T-2]$, 所以此时 MDP 为 non-stationary, policy $\\pi_t$ may be time-dependent, 一般来说仅在 non-stationary dynamics 中有必要 actions may depend on states, 即 $a_t\\in\\mathcal{A}(s_t)$. 由于 $T$ 也是一个随机变量，因此 $\\mathbb{E}[\\sum_{t=0}^T\\cdot]\\neq\\sum_{t=0}^T\\mathbb{E}[\\cdot]$.\n为了简化，我们一般会使用一个等价的，不会停止的 MDP:\nthe MDP never stops, 即 $T=\\infty$ $s_t\\in\\mathcal{S}\\cup\\{\\langle term\\rangle\\}$, 这里 $\\langle term\\rangle$ 是一个 normal state. absorbing state: 当 $s_t=\\langle term\\rangle$ 时， $r_t=0, s_{t+1}=\\langle term\\rangle$, 即状态不会离开 $\\langle term\\rangle$. $\\pi(\\cdot\\mid\\langle term\\rangle)$ 不产生任何影响 RL Objective RL 的最终目标为最大化累计奖励，称为 expected return, 这个目标基于reward hypothesis, 即所有的目标都可以被描述为最大化 expected return. 目标函数表达式如下\n$$ \\max_{\\pi}\\quad \\mathbb{E}_{s_0\\sim p_0}^\\pi\\left[\\sum_{t=0}^{T-1}r_t\\right] $$由于未来存在不确定性，我们会对这种不确定性施加惩罚，即时间越远，其对于当前的 reward 就越低，这其实就是经济学中的现值 (present value, PV), 因此我们实际上优化的目标函数是 expected discounted return:\n$$ \\max_{\\pi}\\quad \\mathbb{E}_{s_0\\sim p_0}^\\pi\\left[\\sum_{t=0}^{T-1}\\gamma^tr_t\\right] $$这里 $\\gamma\\in(0, 1]$ 是一个超参数，我们定义\n$$ G_t = r_t +\\gamma r_{t+1}+\\cdots+\\gamma^{T-1-t}r_{T-1}=\\sum_{t'=t}^{T-1}\\gamma^{t'-t}r_{t'} $$$G_t$ 被称为 discounted return.\n","date":"2026-03-18T17:35:25+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/rl-series-1-reinforcement-learning-basic-definitions/","title":"(RL series 1) Reinforcement Learning basic definitions"},{"content":"Fix Point Theorem, 即不动点定理，是泛函分析中的基本工具，被广泛应用于非线性函数的分析。\n在介绍不动点定理之前，我们先介绍两个概念\n首先是不动点的概念。\nDefinition\n对于函数 $f:\\mathbb{R}^n\\to\\mathbb{R}^n$, 如果一个点 $x^*\\in\\mathbb{R}^n$ 满足\n$$ f(x^*)=x^* $$则我们称 $x^*$ 是函数 $f$ 的不动点。\n接下来是 contraction mapping 的概念\nDefinition\n对于函数 $f:\\mathbb{R}^n\\to\\mathbb{R}^n$, 如果存在 $\\gamma\\in(0,1)$ 满足\n$$ \\|f(x_1)-f(x_2)\\| \\leq \\gamma \\|x_1-x_2\\|,\\forall\\ x_1,x_2\\in\\mathbb{R}^n $$则我们称 $f$ 是一个 contraction mapping. 这里 $\\|\\cdot\\|$ 是一个 matrix norm.\n接下来，我们介绍不动点定理\nTheorem\n给定 equation $x=f(x)$, 其中 $f:\\mathbb{R}^n\\to\\mathbb{R}^n$, 如果 $f$ 是一个 contraction mapping, 则 $f$ 具有如下性质\nExistence: 存在 fixed point $x^*\\in\\mathbb{R}^n$ 满足 $f(x^*)=x^*$. Uniqueness: fixed point $x^*$ 唯一。 Algorithm: 对任意 $x^0\\in\\mathbb{R}^n$, 使用迭代算法 $x_{k+1}=f(x_k)$ 产生的序列 $\\{x_k\\}_{k=0}^{\\infty}$ 收敛到 fixed point $x*$, 且收敛速度为指数级。 证明需要用到柯西列的概念。\nDefinition\n一个序列 $x_1,x_2,\\dots$ 被称为柯西列 (Cauchy sequence) 当且仅当对任意 $\\epsilon\u003e0$, 都存在 $N\u003e0$, 使得\n$$ \\|x_m-x_n\\| \u003c\\epsilon,\\forall m, n\u003eN $$柯西列的一个重要性质为柯西列一定是收敛列。\n证明\n我们首先证明由 $x_k=f(x_{k=1})$ 产生的序列 $\\{x_k\\}_{k=1}^{\\infty}$ 是收敛的，我们通过证明序列 $\\{x_k\\}_{k=1}^{\\infty}$ 是一个柯西列来证明这一点。\n注意到 $f$ 是一个 contraction mapping, 因此\n$$ \\|x_{k+1}-x_k\\| = \\|f(x_k)-f(x_{k-1})\\|\\leq \\gamma \\|x_k-x_{k-1}\\| $$迭代下去，我们就得到\n$$ \\|x_{k+1}-x_k\\|\\leq \\gamma \\|x_k-x_{k-1}\\|\\leq\\cdots\\leq \\gamma^k \\|x_1-x_{0}\\| $$现在我们证明序列 $\\{x_k\\}_{k=1}^{\\infty}$ 是一个柯西列：\n$$ \\begin{aligned} \\|x_m-x_n\\| \u0026= \\|x_m-x_{m-1}+x_{m-1}-\\cdots-x_{n+1}+x_{n+1}-x_n\\|\\\\ \u0026\\leq \\sum_{i=n}^{m-1}\\|x_{i+1}-x_i\\|\\\\ \u0026\\leq \\sum_{i=n}^{m-1}\\gamma^i \\|x_{1}-x_0\\|\\\\ \u0026\\leq \\frac{\\gamma^n}{1-\\gamma}\\|x_{1}-x_0\\|. \\end{aligned} $$从而，序列 $\\{x_k\\}_{k=1}^{\\infty}$ 是一个柯西列, 因此也是一个收敛列。\n接下来，我们证明 $x^*=\\lim_{k\\to\\infty}x_k$ 是 $f(x)$ 的不动点，注意到\n$$ \\|f(x_k)-x_k\\| = \\|x_{k+1}-x_k\\|\\leq \\gamma^k\\|x_1-x_0\\| \\to 0, k\\to\\infty $$我们有 $\\lim_{k\\to\\infty}f(x_k)=\\lim_{k\\to\\infty}x_k$ , 由于 contraction mapping 一定是连续的，因此我们就可以得到 $f(x^*)=x^*$.\n然后，我们证明不动点唯一。假设还存在一个另外一个不动点 $x'\\neq x^*$ 满足 $f(x')=x'$, 那么\n$$ \\|x'-x^*\\| = \\|f(x')-f(x')\\| \\leq \\gamma \\|x'-x^*\\| $$由于 $\\gamma\\in(0,1)$, 因此上述等式当且仅当 $\\|x'-x^*\\|=0$, 这与前面假设矛盾，因而不动点是唯一的\n最后，我们证明 $x_{k+1}=f(x_k)$ 这个算法的收敛速度为指数级，注意到\n$$ \\|x^*-x_n\\| = \\lim_{m\\to\\infty}]\\|x_m-x_n\\| \\leq \\frac{\\gamma^n}{1-\\gamma}\\|x_1-x_0\\| $$因为 $\\gamma \u003c1$, 因此收敛速度为指数级\n","date":"2026-03-09T17:16:02+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/fix-point-theorem/","title":"Fix Point Theorem"},{"content":"Introduction 当我们在硬件上运行算法时，我们的算法效率取决于三个因素：\n计算机的计算效率，也就是每秒可执行的操作数 (OPs/s) 数据移动效率，也就是每秒可传输的 bytes (bytes/s) 需要的内存空间 (bytes) 通过分析这三个因素，我们可以给出一个计算的上下界\nArithmetic Intensity 首先，我们要回答的第一个问题是，运行时间意味着什么？为什么实际运行时间比理想时间要更多？\n容易知道，实际的运行时间由三部分组成\nComputation 首先，运行时间与我们的算法和 GPU 设备相关，在深度学习中，我们的算法通常由 FLOPs 决定，GPU 决定了算法的运行时间：\n$$ T_{\\mathrm{math}} =\\frac{\\text{Computation FLOPs}}{\\text{Accelerator FLOPs/s}} $$Communication within a chip 其次，运行时间还受加载数据的速度的影响，在执行计算之前，我们需要先把数据从 HBM 加载到对应的 core 上，因此，加载数据的速度也会影响最终运行时间。我们将加载数据的速度称之为 HBM bandwidth.\nCommunication between chips 最后，当我们在多个 GPU 上运行模型时，我们还需要在 GPU 之间进行通信，相关的通信手段如下表所示 (intra node 代表同一台服务器，inter node 代表不同服务器)，我们将 GPU 的通信速度称之为 Network bandwidth.\nMethod Data Path BandWith Use cases Intra Node NVLink Direct GPU-to-GPU 1.8 Tb/s (B200) TP NVSwith all-to-all 900 GB/s-1.8T b/s full mesh sync in DGX GPUDirect P2P Direct over PCIe Bus 32-128 GB/s small clusters withou NVLink Shared Memory GPU -\u0026gt; CPU RAM -\u0026gt; GPU 20-50 GB/s Fallback when P2P is disabled Inter Node InfiniBand GPU -\u0026gt; NIC -\u0026gt; IB Switch 200-400 Gbps Large scale LLM training GPU Direct RDMA GPU -\u0026gt; NIC -\u0026gt; Network Same as NIC speed Bypassing CPU for network sync RoCE (v2) RAME over Ethernet 100-400 Gbps Cost-effective Ethernet clusters TCP/IP sockets GPU -\u0026gt; CPU -\u0026gt; Kernel -\u0026gt; Net \u0026lt; 100 Gbps Standard networking 我们将 GPU 的通信 (intra-node 和 inter-node) 效率定义如下\n$$ T_{\\mathrm{comm}} =\\frac{\\text{Communication Bytes}}{\\text{Network/Memory Bandwith Bytes/s}} $$通常来说，我们可以通过计算 - 通信重叠来降低整体的运行时间，比如说 DeepSeek-V3 等模型提出的优化方法, 此时理论上的最小运行时间就是 $T_{\\mathrm{math}}$ 和 $T_{\\mathrm{comm}}$ 之间的最大值。另一方面，如果我们不做任何优化，则理论上最大运行时间就是 $T_{\\mathrm{math}}$ 和 $T_{\\mathrm{comm}}$ 之和。我们定义对应的 bound 如下\n$$ \\begin{aligned} T_{lower} \u0026= \\max(T_{\\mathrm{math}} ,T_{\\mathrm{comm}} )\\\\ T_{upper} \u0026= T_{\\mathrm{math}}+T_{\\mathrm{comm}} \\end{aligned} $$注意到\n$$ T_{upper} = T_{\\mathrm{math}}+T_{\\mathrm{comm}} \\leq 2\\max(T_{\\mathrm{math}} ,T_{\\mathrm{comm}} )= 2T_{lower} $$也就是最大运行时间最多为最小运行时间的两倍。因此，我们一般直接优化 $T_{lower}$.\n如果我们假设我们可以完全达到计算 - 通信重叠，则：\n当 $T_{\\mathrm{math}}\u003eT_{\\mathrm{comm}}$ 时, 我们完全利用了 GPU 的资源，因为此时运行时间受硬件的计算效率约束，我们将其称之为 compute bound 当 $T_{\\mathrm{comm}}\u003e T_{\\mathrm{math}}$ 时, 此时运行时间受硬件的通信效率影响，GPU 的（计算）利用率不完全，我们将其称为 communication bound 为了评估一个算法时 compute bound 还是 communication bound, 我们可以使用 arithmetic intensity 来进行评估。我们先分析 single GPU 的情况，再分析 multi GPU 的情况\nSingle GPU Arithmetic Intensity 一个算法的 arithmetic intensity 定义为算法通信所需要的 FLOPs 和所需要的 bytes 之比：\n$$ \\text{Arithmetic Intensity} = \\frac{\\text{Computation FLOPs}}{\\text{Communication Bytes}} $$arithmetic intensity 衡量了一个算法的 FLOPs per bytes, 当 arithmetic intensity 比较高时，说明 $T_{\\mathrm{math}}\u003eT_{\\mathrm{comm}}$, 此时 GPU 的利用效率较高，反之则说明 GPU 出现了空置的情况。这两种情况的交叉点由 GPU 的 Peak arithmetic intensity 决定：\n$$ \\begin{aligned} T_{\\mathrm{math}}=T_{\\mathrm{comm}} \u0026\\Leftrightarrow \\frac{\\text{Computation FLOPs}}{\\text{Accelerator FLOPs/s}} = \\frac{\\text{Communication Bytes}}{\\text{Bandwith Bytes/s}}\\\\ \u0026\\Leftrightarrow \\frac{\\text{Computation FLOPs}}{\\text{Communication Bytes}} = \\frac{\\text{Accelerator FLOPs/s}}{\\text{Bandwith Bytes/s}}\\\\ \u0026\\Leftrightarrow \\text{Intensity}(\\text{Computation}) = \\text{Intensity}(\\text{GPU}) \\end{aligned} $$这里 $\\text{Intensity}(\\text{GPU})$ 就是 GPU 达到 peak FLOPs 时对应的 arithmetic intensity.\n不同 GPU Tensor Core 对应的 intensity 如下所示\nGPU Generation bandwidth peak BF16 Tensor Core intensity V100 Volta 300 GB/s 125 TFLOPs/s 427 A100 Amper 1,935 GB/s 156 TFLOPs/s 83 H100 Hopper 3.35TB/s 495 TFLOPs/s 147 H200 Hopper 4.8TB/s 495 TFLOPs/s 147 B200 Blackwell 576 TB/s 2250 TFLOPs/s 3.9 Example 假设我们使用 BF16 精度计算两个长度为 $N$ 的向量的内积，计算时的流程如下：\n从 HBM 中加载两个向量，通信量为 $2*2N$ 计算两个向量内积，需要 $N$ 次乘法和 $N-1$ 次加法 将结果写回 HBM, 通信量为 $2$. 从而内积的 intensity 为\n$$ \\text{Intensity}(\\text{dot product}) = \\frac{\\text{Total FLOPs}}{\\text{Total Bytes}} = \\frac{N+N-1}{2N+2N+2} = \\frac{2N-1}{4N+2}\\to \\frac12 $$也就是说，内积的 arithmetic intensity 最大为 $1/2$, 说明这是一个 memory-bound operation.\n接下来，我们可以对 communication-bound 和 compute-bound 进行可视化，如下图所示\n这里我们展示了两个算法 Algo 1 和 Algo 2 的 arithmetic intensity, 当 intensity 较小（小于硬件的 intensity）是，算法为 communication bound, 随着 intensity 增加，算法逐渐由 communication bound 过渡到 compute bound. 上面这个模型被称为 Roofline model, 其展示了算法计算效率与硬件之间的关系。\n为了提高算法计算效率，我们可以通过提升算法的 arithmetic intensity 或者提高 memory bandwith.\nExample\n接下来我们来看一下矩阵计算的 intensity, 我们仍然假设在 BF16 精度下进行计算，假设输入矩阵为 $X\\in\\mathbb{R}^{b\\times d}$ 以及 $Y\\in\\mathbb{R}^{d\\times f}$ , 则输出为 $Z=XY\\in\\mathbb{R}^{b\\times f}$. 计算过程为：\n从 HBM 中加载 $X$ 和 $Y$, 通信量为 $2df+2bd$ bytes 计算矩阵乘法，计算量为 $2bdf$ FLOPs 将结果写回到 HBM, 通信量为 $2bf$. 因此对应的 intensity 为\n$$ \\text{Intensity}(\\text{MatMul}) = \\frac{2bdf}{2bd+2df+2bf} = \\frac{bfd}{bd+df+bf} $$假设 $b \u003c\u003c \\min(d,f)$ (在 LLM 中这个假设一般成立)，我们有\n$$ \\text{Intensity}(\\text{MatMul}) = \\frac{bfd}{bd+df+bf}\\approx \\frac{bdf}{df}=b $$因此我们就可以通过调整 $b$ (batch size) 来将算法从 memory-bound 转换为 compute bound.\nMulti GPU 在 multi GPU 的场景下，我们需要关注 GPU 之间通信的效率，一般来说，GPU 之间效率要远小于 GPU 内部通信效率，下面是不同 GPU 的内部通信（HBM bandwidth） 和 GPU 之间通信效率对比\nGPU Generation HBM Type HBM bandwidth NVLink Generation NVLink Bandwidth V100 Volta HBM2 300 GB/s 2.0 300 GB/s A100 Amper HBM2e 1,935 GB/s 3.0 600 GB/s H100 Hopper HBM3 3.35TB/s 4.0 900 GB/s H200 Hopper HBM3e 4.8TB/s 4.0 900 GB/s B200 Blackwell HBM3e 576 TB/s 5.0 1800 GB/s Example 我们考虑一个分布式矩阵乘法，假设我们有两个 GPU, 计算精度为 BF16, 矩阵 $X\\in\\mathbb{R}^{b\\times d}$ 以及 $Y\\in\\mathbb{R}^{d\\times f}$ 分别被切分为相同大小存储在这两个 GPU 上，此时，计算流程为：\n在 GPU0/GPU1 上，分别从对应的 HBM 中加载 $X_1, Y_1$ 以及 $X_2, Y_2$ 在 GPU0/GPU1 上，分别计算对应的矩阵乘法 $A=X_1Y_1\\in\\mathbb{R}^{b\\times f}$ 和 $B=X_2Y_2\\in\\mathbb{R}^{b\\times f}$, 计算量为 $2bdf$. 将 GPU1 的结果 $B$ 通信传输到 GPU0 上，通信量为 $2bf$. 对结果进行相加，$Z=A+B$. 计算量为 $bf$. 此时，我们的 intensity 为 （注意这里计算时间除以 2 代表我们使用两张 GPU 因而效率翻倍了）\n$$ \\text{Intensity}(\\text{MatMul (2-GPU)}) = \\frac{2bfd/2}{2bf}= \\frac{d}{2} $$可以看到，此时决定性因素在于 $d$ 而不是 $b$, 这对于我们实现并行算法至关重要\n","date":"2026-02-26T17:23:40+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-roofline-model/","title":"Notes on roofline model"},{"content":"Introduction Step3-VL-10B 基于 Perception Encoder 和 Qwen3 开发得到，模型在预训练阶段使用了 1.2T 多模态 token, 在 post training 阶段使用了 PaCoRe 来提高模型的表现。作者强调关键改进在于高质量的预训练数据以及 RL 阶段的提升。\n总的来说，感觉这是阶跃在多模态大模型领域的一次初步尝试，使用的技术路线都比较成熟。\nMethod Architecture Step3-VL-10B 包含 3 个模块：\nViT: 基于 Perception Encoder, 大小为 1.8B MLP: 基于 DeepSeek-OCR 构建了一个两层的卷积层，将视觉 token 个数压缩为原来的 16 倍，对应代码如下 1 2 3 4 5 6 7 # vision_encoder.py StepRoboticsVisionEncoder self.vit_downsampler1 = nn.Conv2d(self.hidden_size, self.hidden_size * 2, kernel_size=3, stride=2, padding=1) self.vit_downsampler2 = nn.Conv2d(self.hidden_size * 2, self.hidden_size * 4, kernel_size=3, stride=2, padding=1) # modeling_step_vl.py _process_image_features image_features = self.vision_model.vit_downsampler1(image_features) image_features = self.vision_model.vit_downsampler2(image_features) LLM: 基于 Qwen3 的 Qwen3-8B 对于输入的图片，Step3-VL-10B 采用了 LLaVA-OneVision 的做法，即将图片分为 $728\\times 728$ 的 global thumnail 和 $504\\times 504$ 的 local crop.\n1 2 3 4 # modeling_step_vl.py _process_image_input image_features = self._process_image_features(image_features) patch_image_features = self._process_image_features( patch_image_features) if patch_image_features is not None else None Pre-training 预训练的数据包括：\nknowledge: knowledge 数据又包括图文交错数据，image-text pairs 数据 education: 15M K12, university, adult education 数据 OCR: 又分为以下几类 image to text: 10M (real-world) + 30M (synthetic) 数据 image to code: 10M markup-based code (latex, matplotlib 等) 数据，15M 合成的 infographics 数据，5M reconstruction (tikz) 数据 document to text: 80M full-page 数据 document to code: HTML, markdown, latex 等数据，共 4M tables 和 100M formulas grounding and counting: 400M 数据 VQA: 10M 数据 GUI: 23M 数据 训练使用了 AdamW 优化器，一共训练了 1.2T token, batch size 为 8192, 上下文长度为 4096.\nPost-training post-training 包括 SFT 和 RL 两个阶段\nSFT 阶段又包括了两个小的 stage, 第一个 stage 用于提高模型的纯文本推理能力，纯文本数据和多模态数据的比例为 9:1,训练使用了 190B token; 第二个 stage 用于提高模型的多模态推理能力，纯文本数据和多模态数据的比例为 1:1. 训练使用了 36B token. SFT 训练时 batch size 为 32, 上下文长度为 128K\nRL 阶段作者使用了 PPO 算法进行训练。reward function 也是分为 rule-based 和 model-based\n在 RLVR 之后，作者还进行了 RLHF 来提高模型的对齐能力。\n作者进一步使用了 PaCoRe 来提高模型的并行推理能力。\nExperiments Step3-VL-10B 的表现如下所示，可以看到在 10B 模型下面，Step3-VL-10B 的表现达到了 SOTA.\nAblation Study 作者首先对比了 perception encoder 和 DINOv3 作为 vision encoder 的表现，结果如下表所示\nVision Encoder Perception BLINK Perception Omni. Perception MMVP Perception OCRBench General MMStar General SVQA General CCBench General V* General MMMU General ReMI DINOv3 42.35 43.31 28.00 57.60 41.43 22.18 56.32 34.55 46.56 24.50 PE-lang (Ours) 41.19 43.57 32.00 70.10 42.10 21.15 59.39 37.17 47.67 26.08 Δ -1.16 +0.26 +4.00 +12.50 +0.67 -1.03 +3.07 +2.62 +1.11 +1.58 作者分析认为由于 DINOv3 在纯视觉任务上进行训练，其表现不如使用语言监督信号的 perception encoder\n作者还对比了 AdamW 和 Muon 两种优化器，结果如下表所示\nOptimizer Perception BLINK Perception Omni. Perception MMVP Perception OCRBench General MMStar General SVQA General CCBench General V* General MMMU General ReMI Muon 41.14 42.73 32.00 67.70 44.58 27.08 60.72 36.65 47.56 22.23 Adam (Ours) 40.72 44.94 29.33 71.10 41.77 20.60 60.13 39.27 46.11 25.00 Δ -0.42 +2.21 -2.67 +3.40 -2.81 -6.48 -0.59 +2.62 -1.45 +2.77 结果发现 Muon 可以解决大规模数据的噪声和不平衡问题，但是作者没有使用 Muon, 原因是 Muon 对参数的初始化比较敏感\n作者还探究了 Qwen3-VL 使用的 Deepstack 的有效性，结果如下表所示\nTechnique Perception BLINK Perception Omni. Perception MMVP Perception OCRBench General MMStar General SVQA General CCBench General V* General MMMU General ReMI w/ DeepStack 40.72 42.92 26.00 71.20 43.31 28.66 63.94 36.65 47.44 26.96 w/o DeepStack (Ours) 40.61 43.57 31.33 69.30 42.44 25.20 62.80 38.22 47.78 26.96 Δ -0.11 +0.65 +5.33 -1.90 -0.87 -3.46 -1.14 +1.57 +0.34 +0.00 结果发现，尽管 DeepStack 可以加速训练，但是对于下游任务的提升非常有限，因此作者没有使用这个策略。\n对于 RL 训练，作者发现随着训练进行，模型的 reward 稳步提升。但是其输出长度并不是单调提升的，如下图所示\n作者分析认为这是由于模型在 reasoning 任务和 perception 任务上模型使用的模式不同导致的，reasoning 任务上模型使用更长的输出长度来解决问题，而对于 perception 任务，由于答案唯一且确定，因此模型通过多次探索后会逐渐收敛到唯一的确定性模式，直接给出对应的答案。因此，其输出长度会越来越短。作者针对这种现象给出了一个假设，即针对 perception 任务，我们的训练数据并不包含思考的过程，而这种数据则让模型只能选择直接回答或者瞎猜，而不是先思考再回答。为了解决这个问题，作者使用了 PaCoRe 来让模型通过 proposal-then-refinement 策略来提高模型的思考长度\nConclusion 作者介绍了 Step3-VL，一个 10B 的多模态大模型，作者详细介绍了模型的架构，训练以及数据。\n作者认为后续工作有：\n通过 universal RL Scaling 提高 token efficiency: 将算力重心从 pre-training 迁移到 RL, 这一点与 DeepSeek-V3.2 一致 消除过度思考，提高模型的 reasoning efficiency 提高模型理解物理世界的能力 构建 world model 帮助模型理解世界 使用 high-fidelity environment 来提高模型对于物理定律的理解能力 具身智能 References arxiv ","date":"2026-02-13T18:05:47+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-step3-vl-10b/","title":"Notes on Step3-VL 10B"},{"content":"Introduction Kimi-K2.5 的核心有亮点：\nnative multi-modal: 通过在预训练，SFT, RL 阶段使用多模态数据来提高模型的多模态能力 agent: 通过并行 multi-agent 的方式来提高模型解决复杂问题的效率和能力 Method Architecture Kimi K2.5 是一个标准的 ViT-MLP-LLM 架构，其中\nViT, 基于 Kimi-VL 提出的 MoonViT, 并进行了改进, 参数量为 400M MLP, 基于 patch merger, LLM, 基于 Kimi-k2, 参数量为 1.02T-A32B ViT 作者使用了 Kimi-VL 提出的 MoonViT, MoonViT 基于 SigLIP 提出的 SigLIP-SO-400M 开发得到，MoonViT 使用了 NaViT 来避免切分图片和使用不同精度图片进行训练。\n在 MoonViT 的基础上，Kimi-K2.5 还进一步提出了 MoonViT-3D, 将 NaViT 的思想扩展到了 3D 用于提高模型的视频理解能力，具体做法为将连续 4 帧的视频展开为 1D sequence, 这样在图像上的注意力机制就可以无缝衔接到视频上了。并且，通过这种方式，我们可以让模型关注跨帧的信息（注意力在 4 帧的 token 之间进行），简化代码如下所示\n1 2 3 4 5 6 7 8 9 10 11 12 # config.json temporal_merge_kernel_size # kimi_k25_vision_processing.py split_video_chunks video_chunk = frames[0:4] patches = [] for frame in video_chunk: patches.extend(split_into_patches(frame)) tokens = patches # modeling_kimi_k25.py Learnable2DInterpPosEmbDivided_fixed positions = spatial_embedding + temporal_embedding # modeling_kimi_k25.py MoonViT3dEncoder output = transformer(tokens + positions) 最后，在进入 MLP 之前，作者还对每个 temporal chunk 内的特征进行 pooling 操作，将时序长度压缩到了原来的 1/4, 进而提高模型可处理的视频长度。\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 # modeling_kimi_k25.py tpool_patch_merger def tpool_patch_merger( x: torch.Tensor, grid_thws: torch.Tensor, merge_kernel_size: tuple[int, int] = (2, 2), ) -\u0026gt; list[torch.Tensor]: d_model = x.size(-1) outputs = [] pre_sum = 0 for t, h, w in grid_thws.tolist(): # Get the current sequence seq = x[pre_sum:pre_sum + t * h * w] # Reshape along self.merge_kernel_size and concat to the last dimension kernel_height, kernel_width = merge_kernel_size new_height, new_width = h // kernel_height, w // kernel_width reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean( dim=0) # temporal pooling padded_seq = reshaped_seq.view(new_height * new_width, kernel_height * kernel_width, -1) outputs.append(padded_seq) pre_sum += t * h * w return outputs MLP MLP 使用了 PatchMerger, 用于减少视觉 token 个数，这个方案在之前的 Qwen-VL 系列里已经得到了应用。\nLLM LLM 基于 Kimi-k2 的 MoE 模型，总参数为 1T, 激活参数为 32B\nData Pre-training 预训练阶段一共使用了 15T token, 分为了三个阶段：\nViT-training: 单独训练 ViT, 实际用了 image caption, grounding, ocr, video 等数据进行训练，训练方式采用了类似 InternVL 的方式，即通过 cross entropy loss 来与一个清凉话的 LLM 进行对齐，这个阶段训练使用了 1T token, 然后作者使用了一个非常短的 stage 来更新 MLP 用于对齐 ViT 和 Kimi-K2 Joint pre-training: 训练所有参数，长下文长度为 4K, 使用了 15T token. 这里主要强调了提升代码数据的比例 Long context mid-training: 使用 YARN 来提高模型的上下文长度 最终预训练阶段 recipe 如下所示\nNative Multimodal pre-training 在 Joint pre-training stage, Kimi-K2.5 还采用了一个与 InternVL3 类似的策略，即在预训练一开始直接使用多模态数据进行预训练。\n传统的多模态大模型往往基于一个比较成熟的 LLM backbone 来完成多模态大模型的训练，但是其问题在于成熟的 LLM 其表示空间会收敛到语言模态上，多模态信息的迁移能力比较差。InternVL3 虽然也是 native multimodal pre-training, 但是其仍然依赖于成熟的 LLM. Kimi K2.5 则是使用预训练阶段的 Kimi K2 作为 backbone 来避免表示空间的塌缩，在训练一开始即直接加入少量多模态数据来保持模型的多模态能力。\n作者探究了预训练阶段不同的数据对比，试验结果如下图所示\nVision Injection Timing Vision-Text Ratio Vision Knowledge Vision Reasoning OCR Text Knowledge Text Reasoning Code Early 0% 10%:90% 25.8 43.8 65.7 45.5 58.5 24.8 Mid 50% 20%:80% 25.0 40.7 64.1 43.9 58.6 24.0 Late 80% 50%:50% 24.2 39.0 61.5 43.1 57.8 24.0 结果显示，在训练早期加入少部分的多模态数据可以有效提高模型的表现。\nPost-training Post-training 分为了 SFT 和 RL, SFT 阶段作者使用了合成的高质量数据，主要提升模型的交互式推理能力以及工具调用能力。为了解决传统 VLM 工具调用能力比较差且扩展性差的问题，Kimi-k2.5 提出了 Zero-Vision SFT, 其核心思想模型在预训练阶段已经完成了多模态对齐，因此我们可以仅使用纯文本 SFT 数据来激活 VLM 的视觉 agent 能力，具体做法就是将所有图像操作通过 IPython 的代码进行代理操作，这样视觉工具的调用就编程了程序化的图像处理指令。\n在 RL 阶段，作者基于 Kimi-k1.5 提出的策略优化算法加入了一个 token-level clipping 机制来减少 off-policy divergence, 目标函数如下所示\n$$ \\mathcal{L}(\\theta)=\\mathbb{E}_{x\\sim\\mathcal{D}}\\left[\\frac{1}{N}\\sum_{j=1}^k\\sum_{i=1}^{|y_i|}\\mathrm{clip}\\left(\\log\\frac{\\pi_{\\theta}(y_j^i\\mid x, y_{j}^{0:i})}{\\pi_{\\mathrm{old}}(y_j^i\\mid x, y_{j}^{0:i})},\\alpha,\\beta\\right)(r(x, y_j) - \\bar{r}(x)) - \\tau\\left(\\log\\frac{\\pi_{\\theta}(y_j^i\\mid x, y_{j}^{0:i})}{\\pi_{\\mathrm{old}}(y_j^i\\mid x, y_{j}^{0:i})}\\right)^2\\right] $$其中 $k$ 是针对每个回答 $x$ 的采样次数，$N=\\sum_{j=1}^k|y_j|$ 是一个 batch 里总的 token 个数， $\\alpha,\\beta,\\tau$ 为超参数，$\\bar{r}(x)$ 是对 normalization 的估计，这里采用了 Kimi-K1.5 的 mean reward, 即 $\\bar{r}(x)=1/k\\sum_{j=1}^Kr(x,y_j)$. 这里的 clipping 机制与 PPO 不同的地方在于针对 log-ratio 进行 clipping, 而不依赖于 advantage 的计算。最终训练时使用了 Moonlight 的 MuonClip 算法\n对于 reward 的设计，Kimi-k2.5 也使用了基于规则和基于 reward model 的方式，前者针对答案可验证的任务，后者针对开放式的任务。\n作者还构建了 length penalty 来提高模型的推理效率，作者发现 Kimi-k1.5 和 Kimi-k2 中的 length penalty 虽然可以生成更准确的 reasoning chain, 但是其很难泛化到更高的算力. 为了解决这个问题，作者提出了 Toggle 策略，即在 inference-time scaling 和 budget-constrained optimization 两种模式之间进行切换优化，对应的 reward 定义为\n$$ \\tilde{r}(x,y) = \\begin{cases} r(x,y)*\\mathbb{I}\\{1/k\\sum_{j=1}^kr(x,y_i\u003c\\lambda \\text{ or }|y_j|\\leq \\text{budget}(x)\\},\u0026\\text{if }\\lfloor t/m\\rfloor\\mod 2 == 0\\ (\\text{Phase }0)\\\\ r(x,y),\u0026\\text{otherwise } (\\text{Phase }1) \\end{cases} $$其中 $\\lambda, m$ 都是超参数。budget 基于正确回答的长度的 p 分位得到：\n$$ \\text{budget}(x) = \\text{Percentile}(\\{|y_j| \\mid r(x,y_j)=1, j\\in[k]\\},\\rho) $$两种模式每隔 $m$ 个 iteration 切换一次：\nphase 0: budget limited phase, 训练模型在给定 token budget 下解决问题，减少 reasoning chain 长度 phase 1: scaling phase, 训练模型使用更多的算力来解决更复杂的问题，提高模型的智能程度 作者评估 Toogle 策略得到的结果如下所示\n结果现实，使用 toggle 策略之后，模型的输出长度减少了 30% 左右，且模型的表现并没有明显下降。作者还发现，一些重复的 pattern 也随之降低，且 toggle 策略的泛化程度更高。\n在 Zero-Vision SFT 的基础上，Kimi-k2.5 使用了 Joint multimodal RL 训练策略。现有的多模态 RL 存在的问题为：模型很容易忽略视觉输入而过度依赖于纯文本进行推理。为了解决这个问题，作者构建了需要视觉理解才能得到答案的任务来提高模型对于视觉信息的利用程度，这些任务覆盖三个 domain:\nvisual grounding and counting: 定位和计数 chart and document understanding: 图表文档理解 vision-critical STEM problems: 需要图片来完成求解的数学物理问题 作者在 visual RL 之后评估了模型的表现，发现模型在 MMLU-Pro, GPQA-Diamond 等任务上的表现都有了提升，作者认为 visual RL 可以在不损害模型纯文本能力的情况下提高模型跨模态的泛化性\nAgent Swarm Kimi-k2.5 的另一个重大改进为使用并行机制来提高模型的 agent 能力。传统的 agent 往往序列执行 reasoning, tool-use, 这限制了模型处理复杂任务的能力，Kimi-k2.5 通过 Agent Swarm 和 Parallel Agent Reinforcement Learning (PARL) 来解决这个问题，其核心思想就是并行，框架图如下所示\nagent swarm 架构包含了一个 orchestrator 和若干个 subagent, 为了解决 agent swarm 的 reward 比较难以设置的问题，PARL 构建了三个不同 level 的 reward\n$$ r_{PARL}(x,y) = \\lambda_1 r_{parallel} + \\lambda_2r_{finish} + \\lambda_3r_{perf}(x,y) $$其中 $r_{perf}$ 评估了 solution $y$ 的质量， $r_{parallel}$ 则是避免并行模式崩塌，从 multi-agent 崩塌为 single agent, $r_{finish}$ 则是评估模型的完成性。超参数 $\\lambda_1,\\lambda_2$ 随训练逐渐降为 0 来提高模型整体的表现。\n作者还提出了使用 critical steps 来评估 parallel agent 的计算时间消耗，其计算公式如下\n$$ CriticalSteps = \\sum_{i=1}^T\\left(S_{main}^{(t)}+\\max_iS_{sub,i}^{(t)}\\right) $$其中 $T$ 为一个 episode 的时间，$S_{main}^{(t)}$ 为 orchestrator 在第 $t$ 步的运行时间， $S_{sub,i}^{(t)}$ 为第 i 个 subagent 的运行时间。\n为了提高模型的并行能力，作者构建了一批广度优先搜索和深度优先搜索的数据，通过这些数据的构建，我们可以提高 orchestrator 的并行调用能力。\n最终，PARL 的表现如下所示\nInfra Kimi-k2.5 的 infra 基于 Kimi-k2, 作者主要强调了 decouple encoder process (DEP) 这一改进。之前的工作将 vision encoder 和 text embedding 都做为 PP 的第一个 stage, 但是由于 vision encoder 对不同输入的处理时间不同，这个 stage 的算历和内存分配随输入变化比较大。为了解决这个问题，作者提出了 DEP, 包含三个 stage 来提高训练效率:\nbalanced vision forward: 由于 vision encoder 比较小 (400M), 因此作者将 vision encoder 复制到所有 GPU 上，然后根据负载来将 visual data 分配到不同的 GPU 上进行处理，这个阶段不保存中间激活值，处理完毕之后所有的结果作为 PP Stage0 的输入 backbone training: 正常进行训练，与 LLM 的训练优化一致 vision recomputation \u0026amp; backward: 这个阶段，我们重新计算 vision encoder 的 forward pass, 然后再对 vision encoder 进行反向传播 通过 DEP, Kimi-k2.5 的训练效率达到了 Kimi-k2 的 90%.\nExperiments 首先是 Kimi-k2.5 在 general \u0026amp; reasoning 类任务上的表现，可以看到，Kimi-k2.5 超过了 DeeoSeek-V3.2 的表现，\nBenchmark Kimi K2.5 Claude Opus 4.5 GPT-5.2 (xhigh) Gemini 3 Pro DeepSeek-V3.2 HLE-Full 30.1 30.8 34.5 37.5 25.1† HLE-Full w/ tools 50.2 43.2 45.5 45.8 40.8† AIME 2025 96.1 92.8 100 95.0 93.1 HMMT 2025 (Feb) 95.4 92.9* 99.4 97.3* 92.5 IMO-AnswerBench 81.8 78.5* 86.3 83.1* 78.3 GPQA-Diamond 87.6 87.0 92.4 91.9 82.4 MMLU-Pro 87.1 89.3* 86.7* 90.1 85.0 SimpleQA Verified 36.9 44.1 38.9 72.1 27.5 AdvancedIF 75.6 63.1 81.1 74.7 58.8 LongBench v2 61.0 64.4* 54.5* 68.2* 59.8* 接下来是模型在 coding 任务上的表现\nBenchmark Kimi K2.5 Claude Opus 4.5 GPT-5.2 (xhigh) Gemini 3 Pro DeepSeek-V3.2 SWE-Bench Verified 76.8 80.9 80.0 76.2 73.1 SWE-Bench Pro (public) 50.7 55.4* 55.6 - - SWE-Bench Multilingual 73.0 77.5 72.0 65.0 70.2 Terminal Bench 2.0 50.8 59.3 54.0 54.2 46.4 PaperBench (CodeDev) 63.5 72.9* 63.7* - 47.1 CyberGym 41.3 50.6 - 39.9* 17.3* SciCode 48.7 49.5 52.1 56.1 38.9 OIBench (cpp) 57.4 54.6* - 68.5* 54.7* LiveCodeBench (v6) 85.0 82.2* - 87.4* 83.3 在 agent 任务上的表现\nBenchmark Kimi K2.5 Claude Opus 4.5 GPT-5.2 (xhigh) Gemini 3 Pro DeepSeek-V3.2 BrowseComp 60.6 37.0 65.8 37.8 51.4 BrowseComp (w/ ctx manage) 74.9 57.8 - 59.2 67.6 BrowseComp (Agent Swarm) 78.4 - - - - WideSearch 72.7 76.2* - 57.0 32.5* WideSearch (Agent Swarm) 79.0 - - - - DeepSearchQA 77.1 76.1* 71.3* 63.2* 60.9* FinSearchCompT2\u0026amp;T3 67.8 66.2* - 49.9 59.1* Seal-0 57.4 47.7* 45.0 45.5* 49.5* GDPVal-AA 41.0 45.0 48.0 35.0 34.0 OSWorld-Verified 63.3 66.3 8.6 20.7 - WebArena 58.9 63.4 - - - 多模态表现\nBenchmark Kimi K2.5 Claude Opus 4.5 GPT-5.2 (xhigh) Gemini 3 Pro Qwen3-VL-235B-A22B Image MMMU-Pro 78.5 74.0 79.5* 81.0 69.3 MMMU (val) 84.3 80.7 86.7* 87.5* 80.6 CharXiv (RQ) 77.5 67.2* 82.1 81.4 66.1 MathVision 84.2 77.1* 83.0 86.1* 74.6 MathVista (mini) 90.1 80.2* 82.8* 89.8* 85.8 SimpleVQA 71.2 69.7* 55.8* 69.7* 56.8* WorldVQA 46.3 36.8 28.0 47.4 23.5 ZeroBench 9 3* 9* 8* 4* ZeroBench w/ tools 11 9* 7* 12* 3* BabyVision 36.5 14.2 34.4 49.7 22.2 BLINK 78.9 68.8* - 78.7* 68.9 MMVP 87.0 80.0* 83.0* 90.0* 84.3 OmniDocBench 1.5 88.8 87.7* 85.7 88.5 82.0* OCRBench 92.3 86.5* 80.7* 90.3* 87.5 InfoVQA (test) 92.6 76.9* 84* 57.2* 89.5 Video VideoMMMU 86.6 84.4* 85.9 87.6 80.0 MMVU 80.4 77.3* 80.8* 77.5* 71.1 MotionBench 70.4 60.3 64.8 70.3 - Video-MME 87.4 66.0* 86.0* 88.4* 79.0 LongVideoBench 79.8 67.2* 76.5* 77.7* 65.6* LVBench 75.9 57.3 - 73.5* 63.6 我们这里基于模型在不同类别任务上的排名来进行可视化，结果如下\n从结果可以看出，Kimi-K2.5 的 agent 能力达到了 SOTA 级别，其多模态能力也比较强。\n与 DeepSeek-V3.2 一样，作者也对比了不同模型的推理效率，结果如下图所示\n可以看到，相比与 Kimi-K2, Kimi-K2.5 通过在 RL 层面进行优化，降低了输出长度，但是相比与 DeepSeel-V3.2 和 Gemini3.0 Pro 之间还存在一定差距。\nConclusion 在本文中，作者提出了 Kimi-k2.5， 一个多模态的 agent model, Kimi-k2.5 集成了 Kimi-k2 和 Kimi-VL 的能力，扩展了模型的 agent 能力。\nReferences huggingface arxiv ","date":"2026-02-12T11:13:13+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-kimi-k2.5/","title":"Notes on Kimi-k2.5"},{"content":"Introduction 在本节中，我们先介绍 KL divergence 的基本定义，然后我们介绍 KL divergence 的一般形式，即 f-divergence.\nKL-divergence KL divergence 用于衡量近似概率分布 $Q(x)$ 到真实概率分布 $P(x)$ 的误差，我们可以将其理解为：如果我们用 $Q(x)$ 来替换 $P(x)$, 会有多大的信息损失？\n连续概率分布的 KL divergence 的定义如下\n$$ D_{KL}(P\\parallel Q) =\\mathbb{E}_{x\\sim P}\\left[\\log \\frac{P(x)}{Q(x)}\\right]=\\int P(x)\\log\\left(\\frac{P(x)}{Q(x)}\\right)dx $$离散概率分布的 KL divergence 定义如下\n$$ D_{KL}(P\\parallel Q) = \\sum_{x} P(x)\\log\\left(\\frac{P(x)}{Q(x)}\\right) $$KL divergence 有两几个关键性质：\n非负性：$D_{KL}(P\\parallel Q)\\geq0$, 且 $D_{KL}(P\\parallel Q)=0$ 当且仅当 $P(x)=Q(x)$ 对任意 $x$ 成立 非对称性： 一般情况下，$D_{KL}(P\\parallel Q)\\neq D_{KL}(Q\\parallel P)$. 有限性：如果存在 $x$ 使得 $P(x)\u003e0$ 但是 $Q(x)=0$, 则 $D_{\\mathrm{KL}}(P\\parallel Q)=\\infty$. 一般我们称 $D_{KL}(P\\parallel Q)$ 为 forward KL (相对于 $Q$), 对应的还有 reverse KL $D_{KL}(Q\\parallel P)$ (相对于 $Q$).\nF-divergence KL divergence 是 f-divergence 的一种特殊情况。 f-divergence 是一类衡量不同概率分布 $P$ 和 $Q$ 的函数 $D_f(P\\parallel Q)$.\n假设函数 $f:(0,\\infty)\\to\\mathbb{R}$ 是一个凸函数，且 $f(1)=0$. $P$ 和 $Q$ 是两个概率分布，则 f-divergence 定义如下\n$$ D_f(P\\parallel Q) = \\mathbb{E}_{x\\sim Q}\\left[ f\\left(\\frac{P(x)}{Q(x)}\\right)\\right]=\\int Q(x)f\\left(\\frac{P(x)}{Q(x)}\\right)dx $$我们称 $f$ 为 $D_f$ 的 generator.\n以下是几种常见的 f-divergence:\nName generator forward KL divergence $f(x)=x\\log x$ reverse KL divergence $f(x)=-\\log x$ Total variation $f(x)=1/2\\vert x-1\\vert$ $\\chi^2$-divergence $f(x)=(x-1)^2$ JS-divergence $f(x)=x\\log\\frac{2x}{x+1}+\\log\\frac{2}{x+1}$ 我们这里推导一下 KL divergence 对应的 generator.\n对于 forward KL, 注意到\n$$ D_f(P \\parallel Q) = \\int Q(x) \\left( \\frac{P(x)}{Q(x)} \\log \\frac{P(x)}{Q(x)} \\right) dx = \\int P(x) \\log \\frac{P(x)}{Q(x)} dx = D_{KL}(P \\parallel Q) $$因此 forward KL 对应的 generator 为 $f=x\\log x$.\n对于 reverse KL, 注意到\n$$ D_f(P \\parallel Q) = \\int Q(x) \\left( -\\log \\frac{P(x)}{Q(x)} \\right) dx = \\int Q(x) \\log \\frac{Q(x)}{P(x)} dx = D_{KL}(Q \\parallel P) $$因此 forward KL 对应的 generator 为 $f=-\\log x$.\nProperties of F-divergence f-divergence 性质如下\nlinearity: $D_{a_1f_1+a_2f_2}=a_1D_{f_1}+a_2D_{f_2}$. $D_f=D_g$ 当且仅当存在 $c\\in\\mathbb{R}$ 使得 $f(x)=g(x)+c(x-1)$. non-negativity. $D_f(P\\parallel Q)\\geq0$ 且 $D_f(P\\parallel Q)$ 当且仅当 $P=Q$. 性质 2 证明如下：\n如果 $f(x)=g(x)+c(x-1)$, 则通过定义，我们可以验证得到 $D_f=D_g$.\n反之，如果 $D_f=D_g$, 令 $h=f-g$, 对任意两个在集合 $\\{0, 1\\}$ 上的概率分布 $P,Q$, 由于 $D_f(P\\parallel Q) - D_g(P\\parallel Q)=0$, 我们有\n$$ h\\left(\\frac{P(1)}{Q(1)}\\right) = -\\frac{Q(0)}{Q(1)}h\\left(\\frac{P(0)}{Q(0)}\\right) $$我们不妨假设 $P(0)=aQ(0)$, $P(1)=bQ(1)$, 结合 $P(0)+P(1)=1$ 和 $Q(0)+Q(1)=1$ 我们有\n$$ Q(0) = \\frac{1-a}{b-a}, Q(1) = \\frac{b-1}{b-a} $$从而\n$$ \\frac{h(b)}{b-1}=\\frac{h(a)}{a-1} $$由于我们可以任意选定 $P$ 和 $Q$, 因此 $h$ 是一个线性函数，形式为 $h(x)=c(x-1)$. $\\blacksquare$\nApproximation 本节中，我们将介绍针对 KL divergence 的三种近似形式。\n在实际计算 KL divergence 时，由于：\n完整计算 KL divergence 需要的算力或内存过高 没有闭式解 我们可以仅保存 log-probability, 而不是整个概率分布 因此，我们假设我们只能计算输入 $x$ 对应的概率 $P(x)$ 和 $Q(x)$. 一般来说，我们会通过 Monte Carlo estimate 来进行近似。即我们先对 $P$ 进行采样得到 $x_1,\\dots,x_N\\sim P$, 然后我们构建估计量。\n一个高的估计量应该是无偏 (unbiased) 并且方差低 (low variance) 的。John Schulman 给出了三种 estimator. 我们分别针对 forward KL 和 reverse KL 进行介绍。这里我们定义\n$$ r = \\frac{P(x)}{Q(x)} $$Forward KL Estimation 对于 forward KL $D_{KL}(P\\parallel Q)$, 其对应的 generator 为 $f(x)=x\\log x$, 注意到 $\\mathbb{E}_{x\\sim Q}[r]=1$, 且 $f$ 是一个凸函数，因此我们有 $f(r)-f'(1)(r-1)\\geq0$, 从而我们可以得到一个新的估计为 $\\boxed{k=r\\log r - (r-1)}$.\nReverse KL Estimation 对于 reverse KL $D_{KL}(Q\\parallel P)$, 其对应的 generator 为 $f(x)=-\\log x$, 由概率性质，$\\boxed{k_1=-\\log r}$ 是 $D_{KL}(Q\\parallel P)$ 的一个无偏估计。但是 $k_1$ 的问题在于 当 $r$ 非常小时，$k_1$ 会变得非常大。也就是说，$k_1$ 的 variance 比较高。\nJohn Schulman 基于 f-divergence 泰勒展开给出了一个新的估计 $k_2$, 其定义为\n$$ \\boxed{k_2 = \\frac12(\\log r)^2} $$其期望为\n$$ \\mathbb{E}_Q[k_2] = \\mathbb{E}_Q\\left[\\frac12(\\log r)^2\\right] $$这是一个 f-divergence, 对应的 generator 为 $f_{k_2}(x)=1/2(\\log x)^2$, 而 $D_{KL}(Q\\parallel P)$ 对应的 generator 为 $f_{k_1}(x)=-\\log x$.\n当 $P$ 和 $Q$ 比较靠近时，我们记 $\\theta=r-1$， 对 $D_{f}(P\\parallel Q)$ 在 $x=1$ 处进行展开得到\n$$ \\begin{aligned} D_f(P\\parallel Q) \u0026= \\mathbb{E}_{x\\sim Q}\\left[ f(r)\\right]\\\\ \u0026= \\mathbb{E}_{x\\sim Q}\\left[ f(1) + f'(1)\\theta + \\frac{f''(1)}{2}f(1+\\lambda)\\theta^2+O(\\theta^3)\\right]\\\\ \u0026= \\frac{f''(1)}{2}F\\theta^2+O(\\theta^3) \\end{aligned} $$这里我们应用了 $f(1)=0$, $\\mathbb{E}[\\theta]=0$, $F=\\mathbb{E}[f(1+\\lambda\\theta)$ 是 Fisher information matrix.\n我们分别带入 $f_{k_1}(x)$ 和 $f_{k_2}(x)$ 得到 $f_{k_1}''(1)=f_{k_2}''(1)=1$, 即 $k_1$ 和 $k_2$ 在 $P$ 和 $Q$ 比较靠近时二阶近似是相同的。因此，**$k_2$ 表面上是一个二阶近似，在分布接近时有效，但本质上是在优化 另一个 f-divergence*\nJohn Schulman 还构造了第三种估计。回顾前面 f-divergence 的性质 2，即当 $f(x)=g(x)+c(x-1)$ 时，我们有 $D_f=D_g$, 因此我们可以选取合适的 $c$ 来降低估计的 variance. 注意到 $k_1$ 的主要问题在于存在负数的可能性，因此我们就构建一个对应的估计量来解决这个问题。注意到 $\\log x \\leq x -1$, 因此我们可以令 $c=1$, 此时就得到了新的估计\n$$ \\boxed{k_3 =(r-1)- \\log r } $$$k_3$ 继承了 $k_1$ 的无偏性，并且 $k_3$ 通过 f-divergence 等价类消除了负值，兼顾无偏与低方差，解决了 $k_1$ variance 过大的问题\nExperiments on Approximation 对于分布 $P=\\mathcal{N}(0,1)$ 以及 $Q=\\mathcal{N}(0.1, 1)$, 真实的 KV divergence 为 0.005, 三个 estimator 的误差如下表所示\nMethod Bias Std Dev $k_1$ 0.0001 20.0005 $k_2$ 0.0025 1.4175 $k_3$ 0.0000 1.4163 当 $P=\\mathcal{N}(1,1)$, $Q=\\mathcal{N}(0.1, 1)$ 时， 真实的 KV divergence 为 0.405, 三个 estimator 的误差如下表所示\nMethod Bias Std Dev $k_1$ -0.0000 2.2223 $k_2$ 0.2025 1.6762 $k_3$ 0.0000 1.6342 可以看到 $k_1$ 的 variance 非常大，$k_2$ 是一个有偏估计，$k_3$ 既满足了无偏又满足了 low variance.\nSummary 我们接下来总结 reverse KL $D_{KL}(Q\\parallel P)$ 的近似 $k_1$, $k_2$ 和 $k_3$ 的性质如下 ($r=P(x)/Q(x)$)\nestimation definition motivation bias variance $k_1$ $-\\log r$ naive estimation unbiased high $k_2$ $\\frac12(\\log r)^2$ f-divergence, taylor expansion biased low $k_3$ $(r-1)- \\log r$ f-divergence, non-negativity unbiased low Applications to ML Remark 本节内容主要参考了 KL Divergence for Machine Learning\n我们假设真实目标分布和近似的目标分布分别记为 $p_{data}(x)$ 和 $p_\\theta(x)$. 由于 KL divergence 的非对称性，因此我们需要考虑两种目标函数：\nforward KL: $\\arg\\min_\\theta D_{KL}(p_{data}\\parallel p_\\theta)$ reverse KL: $\\arg\\min_\\theta D_{KL}(p_\\theta \\parallel p_{data})$ 我们将会看到，这两种不同的目标函数导致的结果也不尽相同\nForward KL 对目标函数进行简化得到\n$$ \\arg\\min_\\theta D_{KL}(p_{data}\\parallel p_\\theta) = \\arg\\max_\\theta \\mathbb{E}_{x\\sim p_{data}}\\left[\\log p_\\theta(x)\\right] $$实际在计算时，我们会使用 Monte Carlo 的方式对真实分布进行采样然后进行估计。\nForward KL 其代表的含义为，我们从分布 $p_{data}$ 中进行采样，然后求 $p_\\theta$ 的最大似然估计。最终的结果满足：当 $p_{data}(x)$ 概率很高时，$p_\\theta(x)$ 的概率也需要很高. 这是一种 mean-seeking behavior, 因为 $p_\\theta$ 必须覆盖 $p_{data}$ 的所有 modes.\n一般来说，supervised learning 对应的就是 forward KL. 我们可以证明 forward KL divergence 和 MLE 是等价的。也就是说，最大似然估计得到的分布就是 KL divergence 最小的近似分布。我们将 $p_{data}(x)$ 和 $p_\\theta(x)$ 对应的 KL divergence 进行展开得到\n$$ \\begin{aligned} \\theta_{KL}^* \u0026= \\arg\\min_{\\theta}D_{KL}(p_{data}(x)\\parallel p_\\theta(x))\\\\ \u0026= \\arg\\min_{\\theta} \\int p_{data}(x)\\frac{p_{data}(x)}{p_\\theta(x)} dx\\\\ \u0026= \\arg\\min_{\\theta}\\int p_{data}(x)\\log p_{data}(x) dx - \\int p_{data}(x)\\log p_\\theta(x)dx \\\\ \u0026= \\arg\\min_{\\theta} - \\int p_{data}(x)\\log p_\\theta(x)dx \\\\ \u0026= \\arg\\max_{\\theta} \\int p_{data}(x)\\log p_\\theta(x)dx \\end{aligned} $$实际上，真实的数据分布 $p_{data}(x)$ 是未知的，我们只有从 $p_{data}(x)$ 采样得到的一批数据 $X=\\{x_1,\\dots,x_n\\}\\sim p_{data}(x)$. 基于大数定律，我们有\n$$ \\frac{1}{n}\\sum_{i=1}^n\\log p(\\theta_i\\mid \\theta)=\\mathbb{E}_{x\\sim p_{data}}[\\log p_\\theta(x)] = \\int p_{data}(x)\\log p_\\theta(x)dx, n\\to \\infty $$这样，最大似然估计就与最小化 KL divergence 构建起了联系：\n$$ \\begin{aligned} \\theta_{MLE}^*\u0026=\\arg\\max_{\\theta} \\sum_{i=1}^n \\log p(x_i\\mid \\theta)\\\\ \u0026= \\arg\\max_{\\theta} \\int p_{data}(x)\\log p_\\theta(x)dx\\\\ \u0026= \\theta_{KL}^*, n\\to\\infty. \\end{aligned} $$也就是说，当采样样本足够多的时候，最大似然估计和最小 KL divergence 是等价的。监督学习中，我们先从真实分布 $p_{data}(x,y)$ 中收集一个数据集 $\\mathcal{D}=\\{(x_i,y_i)\\}$, 然后我们会基于模型 $f_\\theta:\\mathcal{X}\\to\\mathcal{Y}$ 和损失函数 $\\mathcal{L}:\\mathcal{Y}\\times\\mathcal{Y}\\to\\mathbb{R}$ 来优化模型参数 $\\theta$:\n$$ \\arg\\min_\\theta \\mathbb{E}_{(x_i,y_i)\\sim\\mathcal{D}}[\\mathcal{L}(f_\\theta(x_i), y_i)] $$对于使用 cross-entropy loss 的分类问题以及 MSE loss 的回归问题，其目标函数实际上都是最小化 KL divergence.\nReverse KL 对目标函数进行简化，得到\n$$ \\arg\\min_\\theta D_{KL}(Q_\\theta\\parallel p_{data}) = \\arg\\max_\\theta \\mathbb{E}_{x\\sim Q_\\theta}\\left[\\log p_{data}(x)\\right] - \\mathbb{E}_{x\\sim Q_\\theta}\\left[\\log Q_\\theta(x)\\right] $$实际在计算时，我们需要知道真实概率分布在采样点上的概率值 $p_{data}(x)$.\nReverse KL 代表的含义为，我们从分布 $p_\\theta(x)$ 中进行采样，然后最大化采样点在 $p_{data}(x)$ 中的概率分布。entropy item 鼓励 $p_\\theta$ 尽可能均匀分布（覆盖广），从而最终结果满足：当 $p_\\theta(x)$ 概率很高时，$p_{data}(x)$ 的概率也需要很高。注意到与 forward KL 不同，Reverse KL 中包含 entropy 项，其避免了 $p_\\theta$ 收缩到 $p_{data}$ 的某一个 非常窄的 mode 上，最终结果是 $p_\\theta$ 会找到 $p_{data}$ 的一个 high probability 以及 wide support 的 mode, 然后进行覆盖。\n一般来说，reinforcement learning 对应的就是 reverse KL, 这是因为我们希望 policy model 不要离 reference model 太远，并不一定要 cover 所有的 mode.\nExperiments on forward and Reverse KL 我们通过概率分布来可视化 forward RL 与 reverse RL 的区别，验证 forward KL 与 reverse KL 不同的模式。\n我们假设 $p_{data}=w_1\\mathcal{N}(\\mu_1, \\sigma_1^2)+w_2\\mathcal{N}(\\mu_2, \\sigma_2^2)$, 然后我们用一个 normal distribution $p_\\theta=\\mathcal{N}(\\mu, \\sigma^2)$ 来近似 $p_{data}$, 这里 $\\theta=(\\mu, \\sigma^2)$. 对于 forward KL, 我们可以从理论上得出最优解，对应的 $\\mu=w_1\\mu_1+w_2\\mu_2$, 而 reverse KL 则只能通过优化的方式进行求解，并且解与初始化条件相关，下面是相关的实验结果\n首先我们令 $w_1=w_2=0.5$, $\\mu_1=\\mu_2=4.0$, $\\sigma_1=\\sigma_2=1$, reverse KL 的初始化条件为 $\\theta_0=(2,1)$, 对应的结果为\n接下来我们改变 reverse KL 的初始化条件为 $\\theta_0=(-2,1)$, 对应的结果为\n可以看到，与前面分析一致，使用 forward KL 时，最终得到的 $p_\\theta$ 会倾向于拟合分布的中心 (mean seeking), 即 $\\mu(p_\\theta)=\\mu(p_{data})$, 而使用 reverse KL 时，最终得到的 $P$ 会倾向于拟合分布的 mode (mode seeking).\nApplications to RL Remark 本节内容主要参考了 Understanding KL Divergence Estimators in RL: From Value Approximation to Gradient Estimation\n在本节中，我们将基于 RL 来推导 KL 的相关性质。为了统一，这里我们使用 RL 中常见的 notation 来进行计算\nnotation description $\\pi_\\theta$ policy model with parameter $\\theta$ $\\pi_{ref}$ reference model $\\pi_{old}$ behavior model to sample from $s_\\theta(x)=\\nabla_\\theta \\log \\pi_\\theta(x)$ score function $\\rho(x)=\\pi_\\theta(x)/\\pi_{old}(x)$ importance weight $\\mathrm{sg}(\\cdot)$ stop gradient operation 首先 score function 有一个期望为 0 的性质：\n$$ \\mathbb{E}_{x\\sim\\pi_\\theta}[s_\\theta(x)]=\\int_x \\pi_\\theta(x)\\nabla_\\theta \\log \\pi_\\theta(x)dx = \\int_x\\nabla_\\theta \\pi_\\theta(x)dx= \\nabla_\\theta\\int_x \\pi_\\theta(x)dx =\\nabla_\\theta1 = 0 $$接下来，我们分别推导 forward KL 和 reverse KL 的梯度。对于 forward KL, 我们有\n$$ \\nabla_\\theta D_{KL}(\\pi_{ref}\\parallel \\pi_\\theta) = -\\int \\pi_{ref}\\nabla_\\theta \\log \\pi_\\theta dx=-\\mathbb{E}_{\\pi_{ref}}[s_\\theta] = \\boxed{-\\mathbb{E}_{\\pi_\\theta}\\left[\\frac{\\pi_{ref}}{\\pi_\\theta}s_\\theta\\right]} $$对于 reverse KL,我们有\n$$ \\begin{aligned} \\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})\u0026 = \\int\\left[\\nabla_\\theta \\pi_\\theta\\cdot\\log\\frac{\\pi_\\theta}{\\pi_{ref}} + \\pi_\\theta \\nabla_\\theta\\log \\frac{\\pi_\\theta}{\\pi_{ref}}\\right]dx\\\\ \u0026= \\int \\pi_\\theta s_\\theta\\log \\frac{\\pi_\\theta}{\\pi_{ref}}dx + \\int \\pi_\\theta s_\\theta dx\\\\ \u0026= \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta\\log \\frac{\\pi_\\theta}{\\pi_{ref}}\\right]+\\mathbb{E}_{\\pi_\\theta}[s_\\theta]\\\\ \u0026= \\boxed{\\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta\\log \\frac{\\pi_\\theta}{\\pi_{ref}}\\right]} \\end{aligned} $$这里我们使用了 $\\nabla_\\theta\\pi_\\theta=\\pi_\\theta s_\\theta$ , $\\nabla_\\theta\\log\\pi_\\theta=s_\\theta$ 以及 前面推导的 $\\mathbb{E}_{\\pi_\\theta}[s_\\theta]=0$ 的结论.\nRL 的目标函数如下\n$$ \\mathcal{J}(\\theta) = \\mathbb{E}_{\\tau\\sim \\pi_\\theta}\\left[\\sum_{t=0}^T\\gamma^tr(s_t,a_t)\\right] - \\beta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref}) $$Ki as Loss 由于 KL divergcne 不能直接计算（或者计算难度较大），因此，基于前面对 KL divergence estimation 的分析，我们可以使用如下代理损失函数来优化我们的模型：\n$$ \\mathcal{J}_1(\\theta) = \\mathbb{E}_{\\tau\\sim \\pi_\\theta}\\left[\\sum_{t=0}^T\\gamma^tr(s_t,a_t)\\right] - \\beta k_i(\\pi_\\theta, \\pi_{ref}) $$这里 $i\\in\\{1,2,3\\}$ 代表了我们使用的估计。从直觉上来说，这样做是没问题的，但是我们将从数学分析上说明，$k_1,k_3$ 作为损失函数都存在问题。其核心问题在于\n$$ \\mathbb{E}[\\widehat{D_{KL}}]=D_{KL} \\nRightarrow \\mathbb{E}[\\nabla_\\theta \\widehat{D_{KL}}] =\\nabla_\\theta D_{KL} $$也就是说，KL divergence estimation 的无偏性不能推导出 KL divergence estimation gradient 的无偏性，这是因为我们在求期望时，对应的概率分布可能也与参数相关。实际上，我们有\n$$ \\begin{aligned} \\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref}) \u0026= \\nabla_\\theta \\mathbb{E}_{x\\sim\\pi_\\theta}[\\widehat{D_{KL}}(\\pi_\\theta\\parallel \\pi_{ref})]\\\\ \u0026= \\mathbb{E}_{x\\sim\\pi_\\theta}[\\nabla_\\theta \\widehat{D_{KL}}(\\pi_\\theta\\parallel \\pi_{ref})] + \\mathbb{E}_{x\\sim\\pi_\\theta}[\\widehat{D_{KL}}(\\pi_\\theta\\parallel \\pi_{ref})\\nabla_\\theta \\pi_\\theta(x)]\\\\ \u0026\\neq \\mathbb{E}_{x\\sim\\pi_\\theta}[\\nabla_\\theta \\widehat{D_{KL}}(\\pi_\\theta\\parallel \\pi_{ref})] \\end{aligned} $$因此 $\\nabla_\\theta \\widehat{D_{KL}}$ 是 $\\nabla_\\theta D_{KL}$ 的一个有偏估计。\n我们分别来分析一下 $k_1,k_2,k_3$ 梯度，\n$$ \\begin{aligned} \\nabla_\\theta k_1 \u0026= \\nabla_\\theta\\left[-\\log \\frac{\\pi_{ref}}{\\pi_\\theta}\\right] = s_\\theta\\\\ \\nabla_\\theta k_2 \u0026= \\nabla_\\theta\\left[\\frac12\\left(\\log \\frac{\\pi_{ref}}{\\pi_\\theta}\\right)^2\\right] = -\\log \\frac{\\pi_{ref}}{\\pi_\\theta}s_\\theta\\\\ \\nabla_\\theta k_3 \u0026= \\nabla_\\theta\\left[\\frac{\\pi_{ref}}{\\pi_\\theta}-1- \\log \\frac{\\pi_{ref}}{\\pi_\\theta}\\right] = \\left(1 - \\frac{\\pi_{ref}}{\\pi_\\theta}\\right)s_\\theta \\end{aligned} $$此时对应的梯度的期望为\n$$ \\begin{aligned} \\mathbb{E}_{\\pi_{\\theta}}[\\nabla_\\theta k_1] \u0026= \\mathbb{E}_{\\pi_{\\theta}}[s_\\theta]=0\\\\ \\mathbb{E}_{\\pi_{\\theta}}[\\nabla_\\theta k_2] \u0026= \\mathbb{E}_{\\pi_{\\theta}}\\left[-\\log \\frac{\\pi_{ref}}{\\pi_\\theta}s_\\theta\\right]=\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})\\\\ \\mathbb{E}_{\\pi_{\\theta}}[\\nabla_\\theta k_3] \u0026= \\mathbb{E}_{\\pi_{\\theta}}\\left[\\left(1 - \\frac{\\pi_{ref}}{\\pi_\\theta}\\right)s_\\theta\\right]=\\nabla_\\theta D_{KL}(\\pi_{ref}\\parallel \\pi_\\theta)\\\\ \\end{aligned} $$也就是说，$k_1$ 估计的梯度的期望为 0，对整体训练没有任何帮助，$k_3$ 估计的梯度的期望等价于优化 forward KL, **只有 $k_2$ 估计的梯度的期望等价于优化 reverse KL.\n在实际代码实现的时候，KL divergence 有两种不同的实现形式：\n第一种是根据定义将 KL divergence 作为损失函数的一部分，此时我们的 KL divergence 参与反向传播，对应的实现方式如下\n1 loss = -advantage * log_prob + beta * kl 第二种是只调整 reward, 而不参与反向传播（通过 $\\mathrm{sg}(\\cdot)$ 实现），对应的实现方式如下所示\n1 shaped_reward = reward - beta * kl.detach() 这两者对于模型的训练影响很大，下面我们分别来进行介绍\nKL as Loss 为了统一 on-policy 和 off-policy 两种形式，我们使用一个统一的表达形式，即\n$$ L=\\rho k_i $$此时对应的 RL 目标函数为\n$$ \\mathcal{J}_2(\\theta) = \\mathbb{E}_{\\tau\\sim \\pi_\\theta}\\left[\\sum_{t=0}^T\\gamma^tr(s_t,a_t)\\right] - \\beta\\rho k_i(\\pi_\\theta, \\pi_{ref}) $$这里\n$$ \\rho = \\frac{\\pi_\\theta}{\\mathrm{sg}(\\pi_{old})} $$是 importance weight,\n当算法为 on-policy 时，$\\pi_\\theta=\\pi_{old}$, $\\rho\\equiv1$. 当算法为 off-policy 时，$\\rho=\\pi_\\theta/\\pi_{old}$, $\\nabla_\\theta \\rho=\\rho s_\\theta$. 通过这种方式，我们使得参数分布本身不会对梯度计算产生影响，从而使得对期望进行求导和对导数求期望相等，即\n$$ \\nabla_\\theta\\mathbb{E}_{\\pi_{old}}[k] = \\int \\pi_{old}(x)\\nabla_\\theta kdx= \\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta k] $$接下来我们来计算对应估计的梯度的期望，即 $\\mathbb{E}[\\nabla_\\theta(\\rho k_i)]$, 首先我们计算对应的梯度\n$$ \\begin{aligned} \\nabla_\\theta (\\rho k_1) \u0026= \\rho s_\\theta k_1+r\\rho_\\theta=\\rho s_\\theta(k_1+1)\\\\ \\nabla_\\theta (\\rho k_2) \u0026= \\rho s_\\theta k_2+\\rho\\left(-\\log \\frac{\\pi_{ref}}{\\pi_\\theta}s_\\theta\\right)=\\rho s_\\theta(k_1+k_2)\\\\ \\nabla_\\theta (\\rho k_3) \u0026= \\rho s_\\theta k_3+\\rho\\left(1 - \\frac{\\pi_{ref}}{\\pi_\\theta}\\right)s_\\theta=\\rho s_\\theta\\left(k_3+1-\\frac{\\pi_{ref}}{\\pi_\\theta}\\right)=\\rho s_\\theta k_1 \\end{aligned} $$注意到 $\\mathbb{E}_{\\pi_{old}} [\\rho k_i]=\\mathbb{E}_{\\pi_{\\theta}}[k_i]$ 以及 $\\mathbb{E}_{\\pi_{\\theta}}[s_\\theta]=0$, 我们对上述梯度求期望得到\n$$ \\begin{aligned} \\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta (\\rho k_1)] \u0026= \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta(k_1+1)]=\\mathbb{E}_{\\pi_{\\theta}}[s_\\theta k_1]=\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})\\\\ \\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta (\\rho k_2)] \u0026= \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta(k_1+k_2)]=\\nabla_\\theta \\mathbb{E}_{\\pi_\\theta}[k_2]\\\\ \\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta (\\rho k_3)] \u0026= \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta k_1]=\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref}) \\end{aligned} $$这里在计算 $\\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta (\\rho k_2)]$ 时，我们使用了 Leibniz 乘法法则：\n$$ \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta(k_1+k_2)]= \\mathbb{E}_{\\pi_{\\theta}}[s_\\theta k_2]+\\mathbb{E}_{\\pi_{\\theta}}[\\nabla_\\theta k_2]=\\nabla_\\theta\\mathbb{E}_{\\pi_{\\theta}}[k_2] $$可以看到，$\\rho k_1$ 和 $\\rho k_3$ 都满足梯度与期望的可交换性，而 $\\rho k_2$ 不满足，为了解决这个问题，我们可以使用 stop gradient, 即 $\\mathrm{sg}(\\rho)l_2$, 此时，我们有\n$$ \\nabla_\\theta(\\mathrm{sg}(\\rho) k_2) = \\mathrm{sg}(\\rho)\\nabla_\\theta k_2 = \\rho s_\\theta k_1 $$对其求期望有\n$$ \\mathbb{E}_{\\pi_{old}}[\\nabla_\\theta(\\mathrm{sg}(\\rho) k_2)] = \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta k_1] = \\mathbb{E}_{\\pi_{\\theta}}[s_\\theta k_1]=\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref}) $$我们将如上结果总结为下表\nLoss gradient expected gradient objective $\\rho k_1$ $\\rho s_\\theta (k_1+1)$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL $\\rho k_2$ $\\rho s_\\theta (k_1+k_2)$ $\\nabla_\\theta\\mathbb{E}_{\\pi_{\\theta}}[k_2]$ f-divergence $\\mathrm{sg}(\\rho) k_2$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL $\\rho k_3$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL 接下来，我们就可以分析在 on-policy 和 off-policy 场景下分析不同 estimator 的性质了。\n如果说，我们显式加入 $\\rho$, 则根据上表我们可以使用上表的 $\\rho k_1$, $\\mathrm{sg}(\\rho) k_2$ 以及 $\\rho k_3$ 都可以作为损失函数的代替。\n注 实际上 on-policy 场景下使用 $k_2$ 也有用的原因在于 $\\nabla_\\theta k_2=s_\\theta k_1$, 也就是 $k_2$ 和 $\\rho k_3$ 的梯度相同，其本质上是一个等效梯度。但是其收敛得到的 policy 与 target optimal policy 不同\n接下来，我们来分析一下 $\\rho k_1, \\mathrm{sg}(\\rho)k_2, \\rho k_3$ 这三种估计的梯度的 variance, 为了避免混淆，【2】使用了 \u0026ldquo;projection variance in any direction\u0026rdquo; 的概念，即任意取一个向量 $u$, 然后计算 $\\rho k_1$ 和后两者之间对应的 variance 的差（由于 $\\mathrm{sg}(\\rho)k_2$ 的梯度与 $\\rho k_3$ 相同，因此这里我们仅计算 $\\rho k_3$），得到:\n$$ \\begin{aligned} \\mathrm{var}[\\nabla_\\theta (\\rho k_1)^Tu] - \\mathrm{var}[\\nabla_\\theta (\\rho k_3)^Tu] \u0026= (\\mathbb{E}_{\\pi_{old}}[(\\nabla_\\theta (\\rho k_1)^Tu)^2] -\\mathbb{E}_{\\pi_{old}}^2[\\nabla_\\theta (\\rho k_1)^Tu] ) - (\\mathbb{E}_{\\pi_{old}}[(\\nabla_\\theta (\\rho k_3)^Tu)^2] -\\mathbb{E}_{\\pi_{old}}^2[\\nabla_\\theta (\\rho k_3)^Tu] ) \\\\ \u0026= \\mathbb{E}_{\\pi_{old}}[(\\nabla_\\theta (\\rho k_1)^Tu)^2] - \\mathbb{E}_{\\pi_{old}}[(\\nabla_\\theta (\\rho k_3)^Tu)^2]\\\\ \u0026= \\mathbb{E}_{\\pi_{old}}[\\rho(x)^2(s(\\theta)(x)^Tu)^2(2k_1(x)+1)] \\end{aligned} $$当 $\\pi_\\theta$ 和 $\\pi_{ref}$ 比较接近时，我们有\n$$ \\frac{\\pi_{ref}(x)}{\\pi_\\theta(x)} = 1+\\epsilon(x), \\text{ where } |\\epsilon(x)| \u003c\u003c 1 $$此时\n$$ 2k_1(x) + 1 = 1-2\\log(1+\\epsilon(x))\\approx 1-2\\epsilon(x) \\geq 0 $$从而我们有\n$$ \\boxed{\\mathrm{var}[\\nabla_\\theta (\\rho k_1)]\\geq \\mathrm{var}[\\nabla_\\theta (\\rho k_3)]=\\mathrm{var}[\\nabla_\\theta (\\mathrm{sg}(\\rho)k_2)]} $$即当 $\\pi_\\theta$ 和 $\\pi_{ref}$ 比较接近时，$\\rho k_3$ 的 variance 比 $\\rho k_1$ 更小，这是由于 $\\rho s_\\theta (k_1+1)$ 额外包含了一个 期望为零的项，这导致了其 variance 比较高。在 DeepSeek-V3.2 中，作者就使用了 $\\rho k_3$ 来降低梯度的 variance, 提高训练的稳定性。\n【3】将相关的估计总结为了下表的形式\nType Loss Gradient Expected gradient Objective Biased Variance on/off-policy $\\rho k_1$ $\\rho s_\\theta (k_1+1)$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased high on/off-policy $\\rho k_2$ $\\rho s_\\theta (k_1+k_2)$ $\\nabla_\\theta\\mathbb{E}_{\\pi_{\\theta}}[k_2]$ f-divergence biased - on/off-policy $\\mathrm{sg}(\\rho) k_2$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased low on/off-policy $\\rho k_3$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased low 【3】还强调了一点就是我们的损失函数必须显式包含 $\\rho$, 在 on-policy 场景下，虽然 $\\rho\\equiv1$, 但是在反向传播时我们通过 $\\nabla_\\theta \\rho=s_\\theta$ 保留了采样信息从而避免了梯度估计期望的错配问题。\n对于 $\\rho k_1$ variance 比较高的特点，我们还可以采用 variance reduction 的方法来降低不同估计的 variance. 【TODO】\nanalytic gradient 当 action space 有限时，我们还可以使用解析梯度【TODO】\nAs a Reward Reshaping Item 接下来我们来探究一下第二种形式，即 KL divergence 只影响最终的 reward, 而不参与反向传播。对应的代理目标函数形式为\n$$ \\mathcal{J}_3(\\theta) = \\mathbb{E}_{\\tau\\sim \\pi_\\theta}\\left[R\\right] - \\beta\\ \\mathrm{sg}(k_i(\\pi_\\theta, \\pi_{ref})) $$这里 $R=\\sum_{t=0}^T\\gamma^tr(s_t,a_t)$ 为 accumulative reward\n首先，基于前面分析，我们可以得到原始目标函数的梯度为\n$$ \\begin{aligned} \\nabla_\\theta \\mathcal{J}(\\theta) \u0026= \\nabla_\\theta\\mathbb{E}_{\\pi_\\theta}\\left[R\\right] - \\beta \\nabla_\\theta D_{KL}(\\pi_\\theta, \\pi_{ref})\\\\ \u0026= \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta R\\right]-\\beta \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta\\log \\frac{\\pi_\\theta}{\\pi_{ref}}\\right]\\\\ \u0026= \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta(R-\\beta k_1) \\right] \\end{aligned} $$代理目标函数的梯度为\n$$ \\nabla_\\theta \\mathcal{J}_3(\\theta) = \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta(R-\\beta k_i) \\right] $$显然，当我们使用 $k_1$ 时，我们有 $\\nabla_\\theta \\mathcal{J}(\\theta)=\\nabla_\\theta \\mathcal{J}_3(\\theta)$.\n当我们使用 $k_2$ 时，带入 $k_2$ 表达式易知 $\\nabla_\\theta \\mathcal{J}_3(\\theta)\\neq \\nabla_\\theta \\mathcal{J}(\\theta)$,\n当我们使用 $k_3$ 时，\n$$ \\begin{aligned} \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta k_i \\right] \u0026= \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta \\left(\\frac{\\pi_{ref}}{\\pi_\\theta}-1- \\log \\frac{\\pi_{ref}}{\\pi_\\theta} \\right)\\right]\\\\ \u0026= \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta \\frac{\\pi_{ref}}{\\pi_\\theta}\\right] - \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta \\right] - \\mathbb{E}_{\\pi_\\theta}\\left[s_\\theta \\log \\frac{\\pi_{ref}}{\\pi_\\theta} \\right]\\\\ \u0026=s_\\theta k_1 -\\nabla_\\theta D_{KL}(\\pi_{ref}\\parallel \\pi_\\theta) \\end{aligned} $$此时，$\\nabla_\\theta \\mathcal{J}_3(\\theta)\\neq \\nabla_\\theta \\mathcal{J}(\\theta)$. 因此，在 on-policy 场景下，只有 $k_1$ 对应的梯度是无偏的\n在 off-policy 场景下，由于 Off-policy 只影响 $R$ 的计算，因此原始目标函数和代理目标函数的梯度仍然保持不变，on-policy 场景的结论也适用。\n总之，当我们将 KL divergence 作为 reward reshaping item 时，只有 $k_1$ 产生的梯度是无偏的。\nComparison of Two Paradigms 接下来我们来比较一下 KL divergence 作为 loss 和 reward shaping item 的异同之处。首先，两者对于梯度的贡献分别为\n$$ \\begin{align} \u0026\\rho s_\\theta k_1\\tag{loss}\\\\ \u0026 \\mathbb{E}_{\\pi_{old}}[\\rho s_\\theta k_1]\\tag{reward shaping} \\end{align} $$即两者在期望上时一致的。但是两者也存在不一致的地方，即 KL divergence 作为 loss 时不会影响 $R$, 而作为 reward shaping item 时会影响。因此这就导致两者的优化方向不一致。\nExperiments 首先，我们来验证前面的结论，我们构造一个包含 $100$ 个 arms 的 multi-arm bandits, 然后令\n$$ \\pi_{ref}=\\epsilon_1, \\pi= \\epsilon_1+\\epsilon_2 $$其中 $\\epsilon_1,\\epsilon_2\\sim\\mathcal{N}(0,1)$, 我们实验 100 次然后取平均值，然后分别计算 estimator 与真实 KL divergence 之间的 MSE 和 estimator gradient 与真实 kl divergence gradient 的 RMSE, 结果如下图所示\n可以看到，这验证了我们之前分析的结论，即 $k_1$ 和 $k_3$ 是无偏估计，而在计算梯度时，只有 $k_2$ 梯度的期望与真实 KL divergence 的梯度相同。\nOverview 我们在本节总结前面的分析，如下表所示\nType Loss Gradient Expected gradient Objective Biased Variance on-policy $k_1$ $s_\\theta$ $0$ constants biased - on-policy $k_2$ $-\\log r s_\\theta$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased on-policy $k_3$ $(1-r)s_\\theta$ $\\nabla_\\theta D_{KL}(\\pi_{ref}\\parallel \\pi_\\theta)$ forward KL biased - on/off-policy $\\rho k_1$ $\\rho s_\\theta (k_1+1)$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased high on/off-policy $\\rho k_2$ $\\rho s_\\theta (k_1+k_2)$ $\\nabla_\\theta\\mathbb{E}_{\\pi_{\\theta}}[k_2]$ f-divergence biased - on/off-policy $\\mathrm{sg}(\\rho) k_2$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased low on/off-policy $\\rho k_3$ $\\rho s_\\theta k_1$ $\\nabla_\\theta D_{KL}(\\pi_\\theta\\parallel \\pi_{ref})$ reverse KL unbiased low on/off-policy $\\rho\\mathrm{sg}(k_1)$ - - - unbiased - on/off-policy $\\rho \\mathrm{sg}(k_2)$ - - - biased - on/off-policy $\\rho \\mathrm{sg}(k_3)$ - - - biased - Conclusion 在本文中，我们详细介绍了 KL-divergence 的基本性质，相关估计方法以及在机器学习特别是 RL 领域中的应用。最终结论为：\n如果希望稳定可控，则将 KL divergence 作为 loss item; 如果希望更灵活，与奖励信号结合的话，则将其作为 reward shaping item. 使用 KL divergence 作为 loss item 时，on-policy 场景下使用 $k_2$ 近似 KL divergence 效果最好；off-policy 场景下，使用 $\\mathrm{sg}(\\rho)k_2, \\rho k_3$ 效果最好 使用 KL divergence 作为 reward shaping item 时，$k_1$ 的效果最好 References Approximating KL Divergence KL Divergence for Machine Learning Understanding KL Divergence Estimators in RL: From Value Approximation to Gradient Estimation On a few pitfalls in KL divergence gradient estimation for RL ","date":"2026-01-24T16:32:14+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-kl-divergence/","title":"Notes on KL divergence"},{"content":"Introduction 当前大语言模型在性能与效率上面临双重挑战：纯 Softmax 注意力计算成本高，而纯线性注意力则性能不足。Qwen3-Next 尝试通过混合注意力机制解决这一矛盾，同时结合 MoE 架构与多项训练优化策略，实现在保持高性能的同时大幅提升训练与推理效率。\nQwen3-Next 包含三个模型：\nQwen3-Next-80B-A3B-Base Qwen3-Next-80B-A3B-Instruct Qwen3-Next-80B-A3B-Thinking Method Architecture 模型架构如下图所示\nHybrid Attention 作者首先总结了 linear attention 和 softmax attention 各自的优缺点。\npros cons linear attention fast low performance softmax attention slow high performance 因此，作者的动机就是是结合 linear attention 与 softmax attention, 在局部利用 linear attention 的高效性来提高训练和推理效率，在关键部分使用 softmax attention 来提高模型的能力。 这种混合注意力机制之前也有很多模型采用，比如 MiniMax-01 等。最终 Qwen3-Next 使用了 Gated DeltaNet+Gated Attention 的混合注意力机制，模型的 transformer layers 按照 4 个为一组，前三层使用 Gated DeltaNet, 第四层使用 Gated Attention.\n下面是一些细节：\nGated DeltaNet 相比于 SWA 和 Mamba2, 其 in-context learning 能力更强 对于 softmax attention: 使用了 Gated Attention 提出的 gating 机制来解决 massive activation 和 attention sink 问题 将 attention head 的 dimension 从 128 提高到 256 使用了和 DeepSeek-V3 类似的 partial RoPE 机制，仅对前 $25\\%$ 的元素进行旋转 MoE 1 个共享专家，512 个路由专家，其中激活专家个数为 10 个。 对于 MoE router 的参数，作者还进行了 normalization 来保证每个专家被选择的概率相同。 与 Qwen3 一致，Qwen3-Next 也是用了 Global-batch load balancing 策略，在保持激活专家数不变的情况下，通过提高总专家个数来降低训练损失。 Normalization and Training 使用 Gemma 提出的 Zero-Centered RMSNorm 以及 weight decay 来避免过大的权重出现 为了提高数据使用效率，作者还使用了 MTP 策略来提高训练效率，模型表现以及 Speculative decoding 的接受率。 预训练时，Qwen3-Next 使用了15T token 进行训练，训练时间相比于 Qwen3-30B-A3B 有了大幅度的提升 Experiments Efficiency 下图是 Qwen3-Next 与 Qwen3-32B 模型的训练效率对比\n从结果可以看出，相比于 Qwen3-32B, Qwen3-Next 只用了 $9.3\\%$ 的算力就达到了更强的表现。\n并且，在 inference 阶段，由于使用了 linear attention, Qwen3-Next 的效率也更高，下面是 Qwen3-Next 相比于 Qwen3-32B 的效率提升\n4K 32K Prefilling $7\\times$ $10\\times$ Decoding $4\\times$ $10\\times$ Performance 下面是 Qwen3-Next-Base 的表现\n可以看到，Qwen3-Next-Base 在多个 Benchmark 上的表现仅次于 Qwen3-235B-A22B\nQwen3-Next-Instruct 的表现如下表所示\nBenchmark Qwen3-Next-80B-A3B-Instruct Qwen3-235B-A22B-Instruct-2507 Qwen3-32B Non-thinking Qwen3-30B-A3B-Instruct-2507 SuperGPQA 58.8 62.6 43.2 53.4 AIME25 69.5 70.3 20.2 61.3 LiveCodeBench v6 56.6 51.8 29.1 43.2 Arena-Hard v2 82.7 79.2 34.1 69.0 LiveBench 75.8 75.4 59.8 69.0 Qwen3-Next-Instruct 的长文本表现（RULER Benchmark）如下\nModel Avg. 4K 8K 16K 32K 64K 96k 128K 192k 256k 384k 512k 640k 768k 896k 1M Qwen3-30B-A3B-Instruct-2507 86.8 98.0 96.7 96.9 97.2 93.4 91.0 89.1 89.8 82.5 83.6 78.4 79.7 77.6 75.7 72.8 Qwen3-235B-A22B-Instruct-2507 92.5 98.5 97.6 96.9 97.3 95.8 94.9 93.9 94.5 91.0 92.2 90.9 87.8 84.8 86.5 84.5 Qwen3-Next-80B-A3B-Instruct 91.8 98.5 99.0 98.0 98.7 97.6 95.0 96.0 94.0 93.5 91.7 86.9 85.5 81.7 80.3 80.3 可以看到， Qwen3-Next-Instruct 在 1M 长度范围内保持稳定性能，整体平均得分 91.8，接近 Qwen3-235B（92.5）。\nQwen3-Next-Thinking 的表现如下表所示\nBenchmark Qwen3-Next-80B-A3B-Thinking Gemini-2.5-Flash Thinking Qwen3-32B Thinking Qwen3-30B-A3B-Thinking2507 SuperGPQA 60.8 57.8 54.1 56.8 AIME25 87.8 72.0 72.9 85.0 LiveCodeBench v6 68.7 61.2 60.6 66.0 Arena-Hard v2 62.3 56.7 48.4 56.0 LiveBench 76.6 74.3 74.9 76.8 可以看到，Qwen3-Next-Thinking 的表现在除了 Livebench 之外的三个 Benchmark 均达到了 SOTA\nConclusion Qwen3-Next 通过混合注意力架构与精细化 MoE 设计，在训练与推理效率上实现突破性提升。其仅以较小计算代价达到接近超大模型性能的表现，为下一代高效大语言模型的设计提供了重要参考。\nReferences Qwen3-Next: Towards Ultimate Training \u0026amp; Inference Efficiency ","date":"2026-01-23T10:29:56+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen3-next/","title":"Notes on Qwen3-Next"},{"content":"Introduction 随着模型参数变大，现有的 GPU 已经很难使用单一 GPU 来训练模型。对于多 GPU 训练场景，目前主要采用了 pipeline parallelism, 比如 GPipe 等，但是，这些策略需要我们对代码进行比较大的改动，这提高了开发成本。\n为了解决多 GPU 训练大规模 LLM 的效率，降低开发成本，目前主要使用了 model parallelism 策略，即对模型进行切分部署在多个 GPU 上。model parallelism 有两种范式：\npipeline parallelism (PP): 将模型按照 layer 进行切分，如 GPipe 等，这种方法的问题是需要额外的逻辑来处理通信以及存在 pipeline bubbles tensor parallelism (TP): 将模型的按照权重进行切分，部署在不同的 GPU 上。 作者在本文中基于 TP 策略来对 attention, FFN layer 进行简单改动来实现训练效率的提升。\n作者通过实现验证了 tensor parallelism 的有效性和高效率，结果发现在 512 张 GPU 的场景下，TP 可以达到 $76\\%$ 的 scaling efficiency (相比于 1 张 GPU 带来的性能提升)\nMethod 作者使用的 transformer 架构如下图所示\n本文中，作者探究了 BERT 和 GPT-2 两种架构。\n首先，我们假设 transformer layer 输入为 $X\\in\\mathbb{R}^{bs\\times d}$, 这里 $b, s$ 分别为 batch size, sequence length, 接下来我们介绍如何针对 FFN, attention 以及 embedding 构建 TP 策略\nFFN 论文中使用的 FFN 为 Linear-GeLU-Linear 的结构，对应第一层权重为 $W_1\\in\\mathbb{R}^{d\\times d_{ff}}$, 第二层权重为 $W_2\\in\\mathbb{R}^{d_{ff}\\times d}$, 对应数学表达式为\n$$ Y = \\mathrm{GeLU}(XW_1)W_2\\in\\mathbb{R}^{bs\\times d} $$我们首先对 $W_1$ 按照 column 进行切分，得到\n$$ W_1 = [W_{11}, W_{12}]\\in\\mathbb{R}^{d\\times d_{ff}}, \\text{ where } W_{11}\\in\\mathbb{R}^{d\\times d_1}, W_{12}\\in\\mathbb{R}^{d\\times d_2}, d_1+d_2=d_{ff} $$这里 $d_1, d_2$ 与我们并行的 GPU 数 (x-way TP) 相关，这样，我们就有\n$$ \\mathrm{GeLU}(XW_1) = \\mathrm{GeLU}(X[W_{11}, W_{12}]) = \\mathrm{GeLU}([XW_{11}, XW_{22}]) = [\\mathrm{GeLU}(XW_{11}), \\mathrm{GeLU}(XW_{12})] $$从而我们可以分别将 $W_{11}$ 和 $W_{12}$ 部署在两个 GPU 上，然后并行计算。\n论文中还介绍如果我们对 $W_1$ 按照 row 进行切分，则最终由于 $\\mathrm{GeLU}(A+B)\\neq \\mathrm{GeLU}(A)+\\mathrm{GeLU}(B)$ 计算时会产生一次额外的同步。\n接下来，对于 $W_2$, 我们按照 row 进行切分得到\n$$ W_2 = \\begin{bmatrix} W_{21}\\\\ W_{22} \\end{bmatrix}\\in\\mathbb{R}^{d_{ff}\\times d}, \\text{ where }W_{21}\\in\\mathbb{R}^{d_1\\times d}, W_{22}\\in\\mathbb{R}^{d_2\\times d}, d_1+d_2=d_{ff} $$计算时，我们有\n$$ \\mathrm{GeLU}(XW_1)W_2 = [\\mathrm{GeLU}(XW_{11}), \\mathrm{GeLU}(XW_{12})]\\begin{bmatrix} W_{21}\\\\ W_{22} \\end{bmatrix} = \\mathrm{GeLU}(XW_{11})W_{21} + \\mathrm{GeLU}(XW_{12})W_{22} $$可以看到，通过按照 row 进行切分，我们可以将 $W_{11}, W_{21}$ 部署在一个 GPU 上，将 $W_{12}, W_{22}$ 部署在另一个 GPU 上，分别计算出 $\\mathrm{GeLU}(XW_{11})W_{21}$ 和 $\\mathrm{GeLU}(XW_{12})W_{22}$ 之后，再通过一此 all-reduce 操作得到最终的输出结果。计算图如下所示\n这里 $f$ 和 $g$ 是两个对偶算子，代表了 TP 产生的额外通信开销\noperator forward backward $f$ identity all-reduce $g$ all-reduce identity 如果说我们使用的是 SwiGLU FFN, 即\n$$ Y = (XW_3\\odot \\mathrm{Swish}(XW_1))W_2 $$我们按照 column 对 $W_1, W_3$ 进行切分，按照 row 对 $W_2$ 进行切分（假设我们有 2 个 GPU），得到\n$$ \\begin{aligned} W_1 \u0026= [W_{11}, W_{12}]\\in\\mathbb{R}^{d\\times d_{ff}}, \\text{ where } W_{11}\\in\\mathbb{R}^{d\\times d_1}, W_{12}\\in\\mathbb{R}^{d\\times d_2}, d_1+d_2=d_{ff}\\\\ W_3 \u0026= [W_{31}, W_{32}]\\in\\mathbb{R}^{d\\times d_{ff}}, \\text{ where } W_{31}\\in\\mathbb{R}^{d\\times d_1}, W_{32}\\in\\mathbb{R}^{d\\times d_2}, d_1+d_2=d_{ff}\\\\ W_2 \u0026= \\begin{bmatrix} W_{21}\\\\ W_{22} \\end{bmatrix}\\in\\mathbb{R}^{d_{ff}\\times d}, \\text{ where }W_{21}\\in\\mathbb{R}^{d_1\\times d}, W_{22}\\in\\mathbb{R}^{d_2\\times d}, d_1+d_2=d_{ff} \\end{aligned} $$然后我们将 $W_{11}, W_{31}, W_{21}$ 放在第一个 GPU 上，将 $W_{12}, W_{32}, W_{22}$ 放在第二个 GPU 上，此时，\n$$ \\begin{aligned} \\mathrm{Swish}(XW_1) \u0026= \\mathrm{Swish}(X[W_{11}, W_{12}]) = \\mathrm{Swish}([XW_{11}, XW_{12}])=[\\mathrm{Swish}(XW_{11}, \\mathrm{Swish}(XW_{12}]\\\\ XW_3\\odot \\mathrm{Swish}(XW_1) \u0026= [XW_{31}, XW_{32}]\\mathrm{Swish}(XW_1) = [XW_{31}\\mathrm{Swish}(XW_{11}), XW_{32}\\mathrm{Swish}(XW_{12})]\\\\ Y = (XW_3\\odot \\mathrm{Swish}(XW_1))W_2\u0026=(XW_3\\odot \\mathrm{Swish}(XW_1))\\begin{bmatrix} W_{21}\\\\ W_{22} \\end{bmatrix} = XW_{31}\\mathrm{Swish}(XW_{11})W_{21}+ XW_{32}\\mathrm{Swish}(XW_{12})W_{22} \\end{aligned} $$这样我们通过一次 all-reduce 也可以完成 SwiGLU FFN 的 tensor parallelism, 示意图如下所示\nAttention Attention 的处理与 MLP 非常相似，论文中的做法就是将不同 head 部署到不同 gpu 上分别进行计算，最后在计算 output projection 时再通过一次 all-reduce 来合并输出，这里我们假设有 $h$ 个 heads, 每个 head 的 dimension 为 $d_h$, 我们先对 query, key, value layer 的 weight $W_Q, W_K, W_V\\in\\mathbb{R}^{d\\times hd_h}$ 进行切分\n$$ W_Q = [W_{Q1}, \\dots, W_{Qh}], W_K = [W_{K1}, \\dots, W_{kh}], W_V = [W_{V1}, \\dots, W_{Vh}] $$其中 $W_{Qi}, W_{Ki}, W_{Vi}\\in\\mathbb{R}^{d\\times d_h}$ 为每个 head 对应的 query, key, value weight. 我们将切分后的 $W_{Qi}, W_{Ki}, W_{Vi}$ 部署在一个 GPU 上（也可以将若干个 head 部署在一个 GPU 上），然后分别计算出每个 GPU 的 attention 结果，最后再进行汇总，如下所示\n$$ \\begin{aligned} o_i \u0026= \\mathrm{softmax}\\left(\\frac{(XW_{Qi})(XW_{Ki})^T}{d_h}\\right) XW_{Vi}, i=1,\\dots,h\\\\ O \u0026= [o_1,\\dots,o_h]W_O \\end{aligned} $$下面是 multi-head attention 对应的 TP 示意图\nEmbedding 对于 Input embedding, 作者将 embedding matrix $E\\in\\mathbb{E}^{V\\times d}$ 按照 row 进行切分（论文中使用了转置，因此是按照 column 进行切分），得到 $E=[E_1,E_2]^T$, 这里 $E_i\\in\\mathbb{R}^{d\\times V_i}$, $V_1+V_2=V$, 接下来我们把切分后的 embedding matrix 部署在不同的 GPU 上，由于每个 GPU 只有部分结果，因此我们还需要进行 all-reduce 来进行汇总。\n而对于 output embedding, 我们也可以使用类似的做法进行切分，每个 GPU 上计算完结果之后我们还需要一个 all-gather 来汇总结果。\n作者在这里还额外介绍了针对 output embedding 的优化方法，由于 embedding 的输出大小为 $[bs, V]$, 而 $V$ 通常比较大，因此，为了降低通信开销，作者将 cross-entropy-loss 与 output embedding kernel 进行融合，这样我们传输的数据量就减少到了 $bs$.\nExperiments 作者首先对 GPT-2 模型进行了修正，首先将 vocab_size 从 50257 提升到 128 的倍数，即 51200. 对于 model+data parallelism, 作者固定 global batch size 为 512. (64-way DP)\n配置如下表所示（head size 为 96）\nHidden size attention heads layers parameters (B) TP TP+DP 1536 16 40 1.2 1 64 1920 20 54 2.5 2 128 2304 24 64 4.2 4 256 3072 32 72 8.3 8 512 对应的 scaling （使用多卡训练后，每个 GPU 相对于单卡训练的利用率）如下表所示\nparallelism TP-1 TP-2 TP-4 TP-8 TP-1+DP-64 TP-2+DP-64 TP-4+DP-64 TP-8+DP-64 scaling 100% 95% 82% 77% 96% 83% 79% 74% Implementation 首先是 linear layer 的 TP 版本，如下所示\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 import torch import torch.nn as nn import torch.distributed as dist class ColumnParallelLinear(nn.Module): def __init__(self, in_features, out_features, bias=True): super().__init__() self.rank = dist.get_rank() self.world_size = dist.get_world_size() self.local_out = out_features // self.world_size self.weight = nn.Parameter(torch.empty(in_features, self.local_out)) def forward(self, x): out = x @ self.weight gather_list = [torch.empty_like(out) for _ in range(self.world_size)] dist.all_gather(gather_list, out) return torch.cat(gather_list, dim=-1) class RowParallelLinear(nn.Module): def __init__(self, in_features, out_features, bias=True): super().__init__() self.rank = dist.get_rank() self.world_size = dist.get_world_size() self.local_in = in_features // self.world_size self.weight = nn.Parameter(torch.empty(self.local_in, out_features)) def forward(self, x): x_local = torch.chunk(x, self.world_size, dim=-1)[self.rank] out = x_local @ self.weight dist.all_reduce(out, op=dist.ReduceOp.SUM) return out 接下来是针对 LLM 中使用的 SwiGLU FFN 进行的优化，基于前面的介绍，我们不需要对基于 column linear 进行 all-reduce, 代码如下所示\nSwiGLU\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 import torch from torch import nn import torch.nn.functional as F import torch.distributed as dist world_size = 1 rank = 0 class ColumnParallelLinear(nn.Module): def __init__(self, in_features: int, out_features: int, dtype = None): assert out_features % world_size == 0, f\u0026#34;Output features must be divisible by world size (world_size={world_size})\u0026#34; self.part_out_features = out_features // world_size self.weight = nn.Parameter(torch.empty(part_out_features, part_in_features, dtype=dtype)) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: y = x @ self.weight return y class RowParallelLinear(nn.Module): def __init__(self, in_features: int, out_features: int, dtype = None): assert in_features % world_size == 0, f\u0026#34;Input features must be divisible by world size (world_size={world_size})\u0026#34; self.part_in_features = in_features // world_size self.weight = nn.Parameter(torch.empty(out_features, part_in_features, dtype=dtype)) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: y = x @ self.weight if world_size \u0026gt; 1: dist.all_reduce(y) return y class MLP(nn.Module): def __init__(self, dim: int, inter_dim: int): super().__init__() self.w1 = ColumnParallelLinear(dim, inter_dim) self.w2 = RowParallelLinear(inter_dim, dim) self.w3 = ColumnParallelLinear(dim, inter_dim) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: return self.w2(F.silu(self.w1(x)) * self.w3(x)) attention\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 class TPMultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, head_dim: int = None): self.d_model = d_model self.num_heads = num_heads self.head_dim = head_dim if head_dim is not None else d_model // num_heads assert self.num_heads % world_size == 0, \u0026#34;num_heads must be divisible by world size\u0026#34; assert self.d_model == self.num_heads * self.head_dim, \u0026#34;d_model must equal num_heads * head_dim\u0026#34; # heads of different GPU self.local_num_heads = self.num_heads // world_size self.local_qkv_dim = self.local_num_heads * self.head_dim self.qkv_proj = ColumnParallelLinear(in_features=d_model, out_features=3 * self.local_qkv_dim) self.out_proj = RowParallelLinear(in_features=d_model, out_features=d_model) self.scale = 1.0 / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)) def _split_heads(self): batch_size, seq_len, _ = x.shape # [batch, seq_len, local_num_heads, head_dim] x = x.reshape(batch_size, seq_len, self.local_num_heads, self.head_dim) # [batch, local_num_heads, seq_len, head_dim] return x.transpose(1, 2) def forward(self, x: torch.Tensor, attn_mask: torch.Tensor = None): batch_size, seq_len, _ = x.shape # [batch, seq_len, local_qkv_dim * 3] qkv = self.qkv_proj(x) # [batch, seq_len, local_qkv_dim] q, k, v = torch.split(qkv, self.local_qkv_dim, dim=-1) # [batch, local_num_heads, seq_len, head_dim] q = self._split_heads(q) k = self._split_heads(k) v = self._split_heads(v) attn_scores = (q @ k.transpose(-2, -1)) * self.scale if attn_mask is not None: attn_scores = attn_scores + attn_mask attn_weights = F.softmax(attn_scores, dim=-1) # [batch, local_num_heads, seq_len, head_dim] attn_output = attn_weights @ v # [batch, seq_len, local_num_heads, head_dim] attn_output = attn_output.transpose(1, 2) # [batch, seq_len, local_qkv_dim] attn_output = attn_output.reshape(batch_size, seq_len, self.local_qkv_dim) gather_list = [torch.empty_like(attn_output) for _ in range(world_size)] dist.all_gather(gather_list, attn_output) # [batch, seq_len, d_model] full_attn_output = torch.cat(gather_list, dim=-1) out = self.out_proj(full_attn_output) return out Embedding\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class ParallelEmbedding(nn.Module): def __init__(self, vocab_size: int, dim: int): super().__init__() self.vocab_size = vocab_size self.dim = dim assert vocab_size % world_size == 0, f\u0026#34;Vocabulary size must be divisible by world size (world_size={world_size})\u0026#34; self.part_vocab_size = (vocab_size // world_size) self.vocab_start_idx = rank * self.part_vocab_size self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: if world_size \u0026gt; 1: mask = (x \u0026lt; self.vocab_start_idx) | (x \u0026gt;= self.vocab_end_idx) x = x - self.vocab_start_idx x[mask] = 0 y = F.embedding(x, self.weight) if world_size \u0026gt; 1: y[mask] = 0 dist.all_reduce(y) return y Conclusion 作者提出了针对 transformer 架构的 tensor parallelism 策略来提高整体的训练效率，通过在训练过程加入四次 all-reduce 通信我们就可以训练更大规模的模型。\nReferences Megatron-LM \u0026amp; Megatron Core Github Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism DeepSeek-V3 ","date":"2026-01-21T18:04:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/megatron-lm/","title":"megatron-lm"},{"content":"Introduction 现有的大部分模型都基于 Transformer 提出的 softmax attention (SDPA), 虽然也有相关的改进工作，但是主要集中于降低 attention 计算复杂度，提高 attention 在推理时的内存使用效率等。之前的工作提出了关于 attention 的两个问题：\nattention sink, 即模型的注意力会放在初始几个 token 上, 这限制了模型的上下文扩展能力 massive activation, 少部分 token 的 hidden states 会非常大，这限制了模型的训练稳定性 在本文中，作者通过在 attention 中加入 gating 机制来探索 gating 对模型表现和训练稳定性的影响。尽管 gating 并没有降低 attention 计算复杂度，但是 gating 提出了一个新的视角，即 sparity 与 attention sink 和 massive activation 息息相关，这为后面 sparse attention 的研究提供了 Insight.\n作者发现，对 Multi head attention 的输出进行 head-specific gating 的效果最好，并且这种方式还可以提高训练稳定性，模型的表达能力和长上下文能力。作者还进一步分析了这种 gating 方式更好的原因，发现有两点：\nnon-linearity: 通过 gating 可以有效提高 output projection layer 输入的秩，进而提高表达能力 sparsity: gating 可以降低 massive activation 和 attention sink 的影响 作者最终推荐使用 element-wise SDPA gating 方式来进行训练\nRelated Work 作者主要介绍了 gating 和 attention sink 这两部分的工作。\ngating 早在 LSTM 和 GRU 使其就得到了广泛的运用，在 transformer 之后，相关的现行注意力也有应用，比如 MiniMax-01 所使用的 Lightning Attention 等，但是这些工作没有系统性探究 gating 背后的机制。\n第二部分是 attention sink, attention sink 现象由 StreamingLLM 提出， 即模型会将相当一部分注意力权重方开始开始的几个 token 上。而本文提出的 gating 机制可以缓解 attention sink 现象。\nMethod 首先是标准 MHA 定义：\n$$ \\begin{aligned} Q \u0026= XW_Q, K=XW_K, V=XW_V\\\\ \\mathrm{Attn}_i(Q,K,V) \u0026= \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V, i=1,\\dots,h\\\\ \\mathrm{MHA}(Q, K, V) \u0026= \\mathrm{Concat}([\\mathrm{Attn}_1,\\dots,\\mathrm{Attn}_h])\\\\ O \u0026= \\mathrm{MHA}(Q, K, V) W_O \\end{aligned} $$这里 $X\\in\\mathbb{R}^{n\\times d}$ 是 transformer layer pre-normalization 的输出（或者 attention block 的输入）, $n$ 是 sequence length, $d$ 是 hidden size, $h$ 是 number of heads, $d_k$ 是 head dimension.\n接下来，作者介绍了不同的 gating 策略。这里作者用同一的公式来进行表示\n$$ Y' = g(Y,X,W_\\theta, \\sigma) = Y\\odot \\sigma(XW_\\theta) $$这里 $Y$ 是输入， $X$ 是 attention 的输入，$W_\\theta$ 是可学习权重\nPosition 首先是位置，作者考虑了如下几种变体：\n$$ \\begin{align} \\mathrm{MHA}(Q, K, V)' \u0026= \\mathrm{MHA}(Q, K, V)\\odot \\sigma\\left(X W_\\theta)\\right) \\tag{G1}\\\\ Q' \u0026= Q\\odot \\sigma\\left(XW_\\theta\\right) \\tag{G2}\\\\ K' \u0026= K\\odot \\sigma\\left(XW_\\theta\\right) \\tag{G3}\\\\ V' \u0026= V\\odot \\sigma\\left(XW_\\theta\\right) \\tag{G4}\\\\ O' \u0026= O\\odot \\sigma\\left(XW_\\theta\\right) \\tag{G5}\\\\ \\end{align} $$这里 $\\sigma$ 是激活函数，$W_\\theta$ 是激活函数的可学习参数，我们可以将其理解为一个 linear layer, 即当前模块的输出取决于输入 hidden sates 经过一个线性层和激活层之后的结果，相似的做法还有 MoE 中的 gating layer, NSA 中的 gating layer 等。对应的示意图如下所示\ngranularity 作者设计了不同粒度的 gating（假设输入为 $X\\in\\mathbb{R}^{n\\times h\\times d_k}$）：\nhead-shared: 不同 head 共享 gating score, Y'[i,h,k]=gate[i,k]*Y[i,h,k] head-wise: 同一个 head 共享 gating score, Y'[i,h,:]=gate[i,h]*Y[i,h,:] element-wise: 不同元素不共享 gating score, Y'[i,h,k]=gate[i,h,k]*Y[i,h,k] 从 attention 的角度看，不同 head 本身就承担不同的语义子空间，如果强行共享 gating，会破坏这种分工。\nformat 作者还构建了 multiplication 和 addition 两种形式：\nmultiplication: $Y'=Y\\odot \\sigma(XW_\\theta)$ addition: $Y'=Y+\\sigma(XW_\\theta)$ activation function 本文中作者使用了 SiLU 和 sigmoid 两种形式，即\n$$ \\sigma_{\\mathrm{sigmoid}}(x) = \\frac{1}{1+e^{-x}},\\quad \\sigma_{\\mathrm{SiLU}} = x*\\sigma_{\\mathrm{sigmoid}}(x)=\\frac{x}{1+e^{-x}} $$Experiments 作者构建了三个模型进行实验，模型配置如下表所示\nModel 1.7B-28 layers 1.7B-48 layers 15B-A2.4B MoE Layers 28 48 24 query heads 16 16 32 key/value heads 8 8 4 head dim 128 128 128 tie embedding yes yes no QK normalization yes yes yes hidden size 2048 1536 2048 ffn hidden size 6144 4608 768 experts - - 128 top-K - - 8 首先是不同 gating 方法对 MoE model 影响，结果如下图所示\n结论如下：\n对 SDPA 的输出 (G1) 或者 value (G2) 进行 gating 效果最好 head-specific gating 效果更好 multiplication 效果比 addition 效果更好 sigmoid 效果比 SiLU 效果更好 总的来说，position 对最终结果提升最明显，其次是 granularity 和 activation function.\n接下来是不同 gating 方法对 dense model 的影响，作者构建了两个 dense 模型，参数都是 1.7B, 这两个模型的 layers 和 FFN hidden size 不同（通过调整保持总参数一致）。作者对比了 G1 和 baseline 的表现， 结果如下图所示\n结论验证了 gating 机制可以有效提高模型的表现。作者还发现使用 gating 之后，模型的训练也更加稳定，训练的损失变化曲线如下图所示\nAnalysis 首先，作者对 multi head attention 进行了重写，得到如下形式\n$$ o_i^k = \\sum_{j=1}^i\\left(S_{ij}^k X_jW_V^k\\right)W_O^k = \\sum_{j=1}^i S_{ij}^k X_j(W_V^kW_O^k) $$也就是说，$W_K$ 和 $W_O$ 可以吸收到一起，由于 $W_V^j\\in\\mathbb{R}^{d\\times d_k}$, $W_O^k\\in\\mathbb{R}^{d_k\\times d}$, 从而 $\\mathrm{rank}(W_V^jW_O^k)\\leq \\max(\\mathrm{rank}(W_V^j), \\mathrm{rank}(W_O^k))\\leq d_k$. 对于 GQA 和 MQA, 最终的有效秩会进一步降低。\n而使用本文提到的 G1 和 G2 gating 策略之后，我们相当于是通过非线性机制提高了上面的秩，进而解决了 softmax attention 表达能力不足的问题, 实际上，StepFun 的 MFA 也是类似的思想。下面是 G1 和 G2 做的改进：\n$$ \\begin{align} o_i^k \u0026= \\sum_{j=1}^i\\left(S_{ij}^k \\mathrm{gating}(X_jW_V^k)\\right)W_O^k\\tag{G1}\\\\ o_i^k \u0026= \\mathrm{gating}\\left(\\sum_{j=1}^iS_{ij}^k X_jW_V^k\\right)W_O^k \\tag{G2} \\end{align} $$通过 gating 的非线性机制，我们提高的矩阵的秩，进而提高了模型的表达能力，而 G5 提升有限的原因也在于此。实验结果如下图所示\n可以看到，不同的 non-linearity 方法对模型表现都有提升，这验证了矩阵秩会影响模型表达能力的分析。\n接下来，作者探究了 gating 机制对 attention score distribution 的影响，结果如下图所示\n实验结果说明：\n有效的 gating 机制对应的 attention score 是非常稀疏的 head-specific sparsity 非常重要，当在不同的 head 共享 gating 时，模型表现会有所下降 gating 必须与 query 相关，与 G2 先比，G1 的表现更好，这说明 gating score 更依赖于 query. 作者认为基于当前 query token 构建 gating, 可以有效过滤历史 token 的噪音信息 non-sparse gating 效果比较差，作者构建了一个 non-sparse 版本的 sigmoid, 结果发现模型表现非常差，这说明了 attention score 应该是一个稀疏形式 通过前面的分析和实验结果，作者认为 gating 机制还可以缓解 attention sink 现象，作者对 baseline 以及 G1 两种方法的 attention 分布进行了可视化，结果如下图所示\n实验结果整理如下表所示\nmethod massive activation attention sink baseline high high input-independence high high head-shared gating low high head-specific gating low low 因此，作者的结论为，input-dependent, head-specific gating 可以提高 attention score distribution 的 sparsity, 进而减缓 attention sink. 并且引入 spaisity 之后，我们还可以避免 massive activation, 进而使用更低的精度进行训练。\n最后，作者探究了以下 gating 机制的上下文扩展能力，作者在已有的模型上基于 32k 上下文长度使用了 80B token 进行 continue pre-training, 然后使用 YARN 将模型上下文长度扩展到了 128K。 测试的结果如下图所示\nMethod 4k 8k 16k 32k 64k 128k Baseline 88.89 85.88 83.15 79.50 - - SDPA-Gate 90.56 87.11 84.61 79.77 - - YaRN Extended Baseline 82.90 (-6.0) 71.52 (-14.4) 61.23 (-21.9) 37.94 (-41.56) 37.51 31.65 SDPA-Gate 88.13 (-2.4) 80.01 (-7.1) 76.74 (-7.87) 72.88 (-6.89) 66.60 58.82 可以看到，对于短上下文，虽然两者表现都有所下降，但是本文提出的 gating 表现下降程度较小。而对于长上下文，本文提出的 gating 机制效果明显更好。作者分析原因认为这是由于 softmax attention 倾向于退化为对少数 token 的依赖， 而 gating 通过引入 token-level sparsity，避免了这种路径依赖。\nConclusion 在本文中，作者系统性探究了 attention 中的 gating 机制，包括 gating 对模型表现，训练稳定性以及训练动态的影响。作者发现，通过提高 non-linearity 和 sparsity 我们可以有效提高模型的上下文能力以及减缓 attention sink 现象。\n从更高层次看，本文的结果可以总结为一点：\nattention 的问题不在于 softmax 本身，而在于线性 aggregation 的表达上限与缺乏选择性。而 gating 提供了一种几乎零成本、却极其有效的方式来引入非线性与稀疏性。\nReferences Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free Appendix 作者在附录中还进一步分析了 massive activation 以及 attention sink.\nmassive activation 并不是 attention sink 产生的必要原因，并且 sparsity 可以减缓这一现象 head-specific gating 会提升 gating score 的值，因此不同的 head 需要安排不同的 sparsity 并不能通过 clipping 的方式来提高训练稳定性 在 continue pre-training 阶段加入 gating 机制并不能提高模型的表现 ","date":"2026-01-20T15:41:52+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gated-attention/","title":"Notes on Gated Attention"},{"content":"Introduction 现在的统一理解与生成多模态大模型面临的主要问题是难以解决语义特征和图像特征之间不匹配性。diffusion model 可以较好学习图像特征，但是缺乏了对于高阶语义特征的理解和推理能力，反之 LLM 可以比较好利用高阶语义特征，而代价则是难以处理图片细节。\n为了解决这个问题，目前的主流做法是分流，分流策略有两种：\n类似 BAGEL 在 transformer 架构上进行分流，分别处理图像模态以及文本模态特征，这也是自 BAGEL 以来比较常用的一种做法 输入输出端进行分流，代表性工作有 Transfusion, 这种策略使用 encoder 将图片映射到文本空间，然后对输出进行解码，通过使用图片生成的目标函数，我们可以保证图片生成的质量 尽管基于 transformer 架构分流策略的效果比较好，但是对应地，其增加了整体训练的代价，而且并没有完成深度模态统一的目标；而基于输入输出分流的方法则因为目标函数不同很容易导致训练不稳定或者损害模型本身的表现。\n在本文中，作者采取的策略为不使用分流策略，避免对模型架构产生比较大的修改。但是，这样就引入了一个新的问题，自回归生成模型的低效率性，目前主流的生成方式为 raster scan next-token prediction, 当图片非常大时，我们要生成的 token 非常多，从而整体的推理效率非常低。作者在这里举例提到 Emu3 生成一张 $1024\\times 1024$ 的图片需要 10 分钟。并且，自回归生成模型对应的 tokenizer 往往基于 reconstruction 的目标进行训练，而这种训练目标产生的 token 更关注图片细节，这与 LLM 更关注语义特征并不一致，因此其效果也更差。\n为了解决已有自回归生成模型的问题，作者提出了使用 VAR 提出的 next-scale prediction 策略，通过 next-scale prediction, 我们可以极大程度提高图片生成的效率，作者这里强调了 NextFlow 生成一张 $1024\\times 1024$ 的图片仅需要 5 秒，这个效率是 Emu3 的 120 倍。\n接下来就是作者针对提出的架构和方法进行的改进，主要包括：\n使用了 6T 数据进行训练 使用了基于 prefix-tuning 策略的 GRPO 方法来提高模型的 reasoning 能力 构建了一个基于 diffusion 的 decoder 来对输出图片进行优化 最后，作者通过实验验证了 NextFlow 的有效性，并且在效率上，相比于基于 MMDiT 架构的模型，NextFlow 使用了更少的算力进行推理。\nMethod Architecture NextFlow 的架构图如下所示\n模型架构包括 3 个部分：\ntokenizer: NextFlow 的 tokenizer 基于 TokenFlow, TokenFlow 通过使用两个 codebook 来分别提取对应的语义特征和图片特征进而提高 tokenizer 的表达能力 transformer: NextFlow 使用了 Qwen2.5-VL 7B 模型作为 base model, 作者将其 Vision Encoder 替换为了 TokenFlow tokenizer 用于提取视觉特征 optinal diffusion decoder: 用于进一步优化图片的细节。作者使用 TokenFLow tokenizer 的 token 表示，外接了一个 diffusion model 来优化最终的输出。 在数据格式上，作者使用了 \u0026lt;boi\u0026gt;, \u0026lt;eoi\u0026gt; 来标记图片 token, 然后每个图片通过 TokenFLow 表示为不同 scale 的 token, 如下图所示\n在位置编码上，作者采用了 Multi-Scale 3D RoPE 策略，第 $t$ 个 token 如果是文本 token, 则表示为 $(t,t,t)$, 如果是图片 token, 则表示为如下形式\n$$ (p_x, p_y, p_s) = \\left(\\frac{C}{\\sqrt{HW}}(i+0.5),\\frac{C}{\\sqrt{HW}}(j+0.5),s\\right) $$这里 $H, W, C, s$ 分别是 grid 对应的 size, constant range factor 以及 scale. 为了保持数据的一致性，作者还对所有的空间位置进行了归一化，避免在 SFT 时进行外推。\n与 VAR 一致，作者使用了 scale length positional embedding, 即在 scale 层面使用了 Transformer 提出的 Sinusoidal encoding, 作者认为通过显示编码 scale 可以提高模型对于图片精度的认知能力。\n在训练上，作者提出了两个优化策略。\n首先，作者发现，不同 scale 的 token 数量不一致，small scale 对应的 token 数量较少，随着生成 token 数增加，由于 attention 存在局部依赖性（见 NSA）模型对于早期 token 关注度较低，为了解决这个问题，作者的做法就是使用 scale-dependent weight, 不同 scale 对应 token 的损失函数权重为\n$$ k_s =\\frac{1}{(h_s\\times w_s)^\\alpha} $$这里 $\\alpha$ 是超参数，通过这种方式，作者保证了模型对于不同 scale token 能够一视同仁。\n接下来，就是训练与推理的不一致性，训练时，模型的上下文是无损的，但是在推理时，由于模型需要基于自己的输出来预测下一个 token, 因此很容易出现误差累积。为了解决这个问题，作者的做法是使用 [[infinity]] 提出的 self-correction 机制，具体做法就是，在 encode 当前 token 时，随机加入一些噪声，制造出有损的上下文然后让模型预测下一个 token, 也就是提高模型对于误差的稳健性。但是，使用这种策略并没有带来理想中的提升，作者分析原因发现是 VAR 的 feature 与输入 feature 存在较大的不一致性，作者的解决方法就是直接使用 residual feature. 通过这种方式模型的能力有了极大的提升\nAblation Study 作者在本节探讨了架构设计上的一些细节。\n首先，作者探究了 tokenizer 的训练，原始 TokenFLow tokenizer 使用了衣蛾预训练的 semantic encoder 以及一个从零开始训练的 pixel encoder, 作者发现这种不一致性会损害模型的表现。为了解决这个问题，作者分别从零开始训练 semantic encoder 和 pixel encoder, 确保两者都在一个水平上，然后再进行 TokenFlow 的训练，通过这种多阶段训练方式，我们可以有效提高训练效率以及模型的表现。训练时，作者还随机丢弃了 $50\\%$ 的 VAR scale 来提高模型的稳健性。作者进一步对比了 TokenFlow 这种 dual-codebook 和 single-codebook 之间的表现，如下图所示\n结果显示，尽管 single-codebook 的重建效果更好，但是在下游任务上其表现不如 TokenFlow.\n接下来，作者探讨了 output head, 也就是我们应该分别针对不同的模态设计不同的 output head 还是使用统一的 output head, 作者对比了两种架构，结果如下图所示\n实验结果显示，single head 的效果更好。因此作者使用了 single head 设计\nTraining 训练流程图如下所示\n在 alignment 时，作者使用 TokenFlow tokenizer 替换 Qwen2.5-VL 7B 的 ViT, 作者同时训练 connector (MLP) 和 output head, 冻结其他参数，这个阶段使用了 10M image-text pairs.\npre-training 时，作者采用了三阶段训练策略，图片数据的精度分别为 256, 512, 1024. 共使用了 6T token.\n256-level: 2B 文生图样本，纯文本数据，以及 147M 图文交错数据 512-level: 提高模型的细节生成能力，使用了 scale-dependent weight 策略 1024-level: 使用了 40M 高质量样本 continue pre-training 和 SFT 阶段，作者分别使用了美学相关的数据和对话数据来提高模型图片生成的质量和指令跟随能力。\nRL 阶段，由于早期的 token 对最终图片生成更加重要，作者使用了一个 prefix-tuning 策略，也就是我们只计算 image token sequence 中前 $m$ 个 scale 对应的 token, 通过这种方式我们提高训练速度以及保证模型生成的质量。训练目标函数如下所示\n$$ L_{GRPO}(\\theta) = \\mathbb{E}_{\\mathbf{c} \\sim \\mathcal{C},\\{\\mathbf{s}_t^i\\}_{i=1}^G \\sim \\pi_\\theta} \\frac{1}{G} \\sum_{t=1}^m k_t \\min\\left( \\frac{p_\\theta(\\mathbf{s}_{t+1}^i|\\mathbf{s}_t^i,\\mathbf{c})}{p_{\\theta_{old}}(\\mathbf{s}_{t+1}^i|\\mathbf{s}_t^i,\\mathbf{c})} A_i, \\, \\text{clip}\\left( \\frac{p_\\theta(\\mathbf{s}_{t+1}^i|\\mathbf{s}_t^i,\\mathbf{c})}{p_{\\theta_{old}}(\\mathbf{s}_{t+1}^i|\\mathbf{s}_t^i,\\mathbf{c})}, \\, 1-\\epsilon, \\, 1+\\epsilon \\right) A_i \\right) - \\beta D_{KL}(\\pi_\\theta, \\pi_{ref}) $$最终，training recipe 如下图所示\nExperiments 我们主要关注 NextFlow 和 Qwen-Image, Seedream 3.0 的表现对比情况，结果如下表所示\nBenchmark NextFlow NextFlow-RL Qwen-Image Seedream 3.0 DPG 86.00 88.32 88.32 88.27 GenEval 0.83 0.84 0.87 0.84 WISE 0.59 0.62 0.62 - PRISM-Bench 74.7 78.8 79.9 79.6 ImgEdit 4.44 4.49 4.27 - Conclusion 作者在本文中提出了 NextFlow, 一个基于自回归的统一理解与生成模型，NextFlow 使用了 TokenFlow 解决了单一 tokenizer 不能同时提取语义信息和像素信息的缺点，使用了 VAR 解决了自回归图片生成范式效率低的问题。\nReferences NextFlow: Unified Sequential Modeling Activates Multimodal Understanding and Generation ","date":"2026-01-17T17:31:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/nextflow-%E5%9F%BA%E4%BA%8Esingle-branch%E7%9A%84%E7%BB%9F%E4%B8%80%E7%90%86%E8%A7%A3%E4%B8%8E%E7%94%9F%E6%88%90%E5%A4%9A%E6%A8%A1%E6%80%81%E5%A4%A7%E6%A8%A1%E5%9E%8B/","title":"NextFlow 基于single-branch的统一理解与生成多模态大模型"},{"content":"介绍 本报告基于 OpenRouter 2024 年 11 月 - 2025 年 11 月 100T token 调用数据，从模型、任务、用户维度分析 AI 大模型使用特征，核心结论如下：\n模型维度\n市场格局：闭源模型占 70% token 使用量，主导高价值、高稳定性场景；开源模型占 30%，聚焦低成本、场景化需求，其中中国开源模型占比持续上升，但开源市场因竞争呈现碎片化（2025 年底无单一开源模型占比超 25%）。 模型偏好：不同闭源模型形成差异化优势（Anthropic 擅长复杂推理 / 代码、Google 偏向通用翻译 / 知识问答、xAI/Qwen 专注编程等）。 任务维度\n核心需求：编程类任务 token 占比超 50%，Anthropic 占该领域 60% 以上份额，且编程任务输入长度是其他类别 3 倍以上；但 90% 编程需求依赖闭源模型，开源模型存在能力短板。 场景特征：角色扮演是开源模型第一大任务（占比超 50%），其在该场景 token 占比（43%）接近闭源模型（42%）且持续上升；agentic inference（工具调用推理）场景 token 占比超 60%，Claude 系列主导该领域。 用户维度\n核心逻辑：能力优于成本，企业愿为强能力模型支付溢价（价格弹性低，降价 10% 仅带来 0.5~0.7% 使用率增长）；免费开源模型仅达 “可用” 水平，因无法落地到实际工作流程难以形成竞争力；小模型数量占比下降，反映市场对模型能力要求提升。 留存逻辑（水晶鞋效应）：用户留存由 “能力拐点”（首次解决未被满足的长尾需求）驱动，如 Gemini 2.5 Pro、Claude 4 Sonnet 实现能力突破后 5 个月留存率仍达 40%；缺乏能力突破的模型留存率极差，且 “水晶鞋时刻” 窗口狭窄。 Background 分析使用的 100T token 基于 OpenRouter 从2024 年 11 月到 2025 年 11 月的模型使用情况。\n对于模型，作者将模型分为三类， 分别是 Proprietary, Chinese Open Sourced (Chinese OSS), Rest-of-World (RoW) open source models\ncategory examples description Proprietary Claude, GPT-5 闭源商业模型 Chinese OSS DeepSeek-V3, Qwen 中文开源模型 RoW OSS mistral, LLaMA 其他国家开源模型 作者基于 metadata 和 GoogleTagClassifier 将任务分为以下类别\ncategory sub-category Programming - Computers \u0026amp; Electronics\n- Programming\n- Science\n- Computer Science Roleplay - Games\n- Roleplaying Games\n- Adult\n- Arts \u0026amp; Entertainment Translation - Reference\n- Language Resources General Q\u0026amp;A / Knowledge - Reference\n- General Reference\n- News Productivity/Writing - Computers \u0026amp; Electronics\n- Software\n- Business \u0026amp; Productivity Software\n- Business \u0026amp; Industrial\n- Business Services\n- Writing \u0026amp; Editing Services Education - Jobs \u0026amp; Education\n- Education Literature/Creative Writing - Books \u0026amp; Literature\n- narrative leaves under - Arts \u0026amp; Entertainment Model 首先是开源模型与闭源模型的对比，结果如下图所示\n结果显示，开源模型的 token 占比在 $30\\%$ 左右，并且中文开源模型 token 占比在持续上升。作者分析认为，闭源模型有着更好的表现以及稳定性，而开源模型的透明程度，成本控制以及可定制化更好。\n在闭源模型中，作者对比了不同模型在不同任务上的使用情况，得出了模型偏好如下表所示\nProvider preference description Anthropic - Programming\n- Technology 擅长推理，代码，复杂任务，领域专家 Google - Translation\n- Science\n- Technology\n- General Knowledge 各个任务都没有短板，通用信息引擎 xAI programming 专注编程，程序员专用 OpenAI - Science\n- Programming\n- Technology 介于 Anthropic 和 Google 之间，更人性化 DeepSeek - roleplay 日常对话，面向消费端 Qwen - programming 专注编程，适用范围比 Anthropic 广 在开源模型中，DeepSeek 占了 $42.51\\%$，模型 token 使用情况对比如下表所示\nProvider # Tokens (T) ratio (%) DeepSeek 14.37 42.51 Qwen 5.59 16.54 Meta LLaMA 3.96 11.72 Mistral AI 2.92 8.64 OpenAI 1.65 4.88 Minimax 1.26 3.73 Z-AI 1.18 3.49 TNGTech 1.13 3.34 MoonshotAI 0.92 2.72 Google 0.82 2.43 但是到 25 年底，因为竞争太强，已经不存在单一模型占比超过 $25\\%$, 下面是开源模型 token 使用随时间变化情况\n接下来，作者将模型分为 Large (\u0026gt; 70B), Medium (\u0026gt; 15B, \u0026lt; 70B) 以及 Small (\u0026lt; 15B) 三个区间，分析了各自的使用情况。结果发现，整体趋势为，Small 区间的模型数量占比正在减少，如下图所示\ntoken 使用这方面，由于小模型本地部署较多，因此报告结果存在一定偏差性。\n作者还对比了不同模型的价格与 token 使用情况，作者根据中位数 $0.73\\$$ per 1M tokens 来将模型分为了四类：\nPremium Workloads (high-cost, high-usage): technology, science 等任务 Mass-Market Volum Drivers (low-cost, high-usage): programming, roleplay 等任务 Specialized Experts (high-cost, low-usage): finance, academia, health 等任务 Niche Utilities (low-cost, low-usage): translation, legal 等任务 图中红线是拟合出来的结果，红线说明，降低 $10\\%$ 的价格只会带来 $0.5\\sim0.7\\%$ 的 token 使用率增长。上面这幅图说明了闭源模型主要解决高价值的任务，而开原模型则是解决 low-cost 的任务\n作者对比不同模型，给出了一些例子，如下表所示\n这里的关键结论有几点：\n宏观需求无弹性，微观行为分化：企业愿意支付成本使用更强模型，而个人用户则对成本比较敏感 Jevons Paradox: 模型成本下降之后，使用率反而会上升 能力优于成本：用户愿意为更强的模型付出更高的成本 低价不能成为竞争力：免费开源模型不能落地的原因是仅仅达到“可用”水平，无法部署到实际工作流程中 Task 首先是所有模型在不同任务上的 token 使用情况，可以看到，programming 的占比最近已经升到了 $50\\%$ 以上\n并且对于 programming, Anthropic 拥有 $60\\%$ 以上的份额，如下表所示\n接下来是开源模型在不同任务上的 token 使用情况。\n排名前二的任务分别为 roleplay 以及 programming, 前者占比 $50\\%$ 以上， 而后者占比在 $15\\sim20\\%$ 。中文的 OSS model 主要也集中在这两个任务上，但是 roleplay 占比下降到了 $33\\%$. 而 programming+technology 的占比为 $39\\%$.\n对于 programming，目前 $90\\%$ 的 token 都基于闭源商业模型，对于开源模型，目前中文开源模型的占比已经超过了其他开源模型\n对于 roleplay, 其他开源模型与闭源商业模型的占比分别为 $43\\%$, $42\\%$ , 且开源模型占比持续上升\n开源模型的几个关键用途为：\nroleplay and creative dialogue: 写作，虚拟人物等 programming: 编程，开发 translation: 多语种任务 general QA: 日常问答 作者还对 token 长度进行了分析，结果显示，目前输入长度区间为 $[1.5K, 6K]$, 输出长度区间为 $[150, 400]$, 这说明不同于早期简单问答，现在用户倾向于输入更丰富的上下文或者材料来让模型解决相应问题。而且，programming 相关的输入长度大约是其他 category 输入的 3 倍以上，下图是不同 category 输入长度随时间的变化情况\n作者还展示了不同子任务的占比情况（只列出大于 $10\\%$ 的部分）\ncategory sub-category ratio (%) Roleplay Games/Roleplaying Games 57.9 Books \u0026amp; Literature/Wrters Resources 16.4 Adult 15.0 Programming general programming 66.1 developing tools 26.4 Translation general 51 Foreign Language Resources 49 Science Machine Learning \u0026amp; AI 80.4 Technology personal assistences 31.3 software 10.9 health general 25.0 research 11.6 finance currencies 19.2 stocks 15.9 investing 15.5 accounting 13.0 academia educational 42.7 writing 36.1 management 14.8 legal Government 42.1 Legal 18.9 trivia - 96.2 marketing marketing 66.3 sales 16.0 seo seo 100 接下来是模型在 agentic inference 场景下的 token 使用情况，从下图可以看出，Reasoning token 占比已经超过了 $60\\%$. 这代表了用户对于使用工具解决复杂能力的需求\n这 agentic inference 场景下, 使用最多的模型有 Grok Code Fast1, Gemini 2.5 Pro/flash 等，下图是不同模型的占比情况\n可以看到，Claude 系列占了大部分份额\nUser 用户这方面，中国和北美的使用占了近 $80\\%$, 其中中文用户最近占 $31\\%$\n语言上，token 的语言占比情况如下表所示\nLanguage Token Share (%) English 82.87 Chinese (Simplified) 4.95 Russian 2.47 Spanish 1.43 Thai 1.03 Other (combined) 7.25 作者还分了用户留存行为，发现一小部分用户的留存率非常高，作者将这个现象称之为水晶鞋效应 (Glass Slipper effect). 作者分析原因有以下几点：\n市场存在未被满足的需求：尽管 AI 模型层出不穷，但是始终有些任务连续多代模型都无法解决，也就是 \u0026quot; 水晶鞋“” 新模型发布是一个“试穿水晶鞋的过程”：每一代模型都会被用于测试是否能解决这些未解决需求 “灰姑娘时刻”：一旦某一个模型解决这个未解决的需求，这个模型就会吸引相当一部分用户从维持比较高的留存率 作者举例说明，Gemini 2.5 Pro 和 Claude 4 Sonnet 在 5 个月以后用户留存率还有 $40\\%$, 而 Gemini 2.0 Flash 和 LLaMA 4 Maverick 的用户留存率非常差。作者因此得出三个关键结论：\n首次解决是持久竞争的核心优势：也就是先发制人 用户留存率是能力拐点的信号：因为模型实现了从不可能到可能得跨越，才能保留一批早期用户 “水晶鞋时刻”的窗口很窄：一旦抓不住机会，很可能就会丢掉一大批用户群体 References state of AI ","date":"2026-01-17T17:04:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/state-of-ai--%E4%BB%8Eopenrouter-100t-token%E4%BD%BF%E7%94%A8%E6%83%85%E5%86%B5%E4%BA%86%E8%A7%A3ai-%E5%A4%A7%E6%A8%A1%E5%9E%8B%E8%83%BD%E5%8A%9B%E5%88%86%E5%B1%82%E7%AB%9E%E4%BA%89%E9%80%BB%E8%BE%91/","title":"State of AI--从OpenRouter 100T token使用情况了解AI 大模型能力分层竞争逻辑"},{"content":"Motivation 我们从一个简单的问题开始\n假如我有一张 80GB 的显卡，我想训练/推理一个 4B 的模型，我应该设置多大的 batch size 和 sequence length?\n在这个 tutorial 中，我们将基于这个问题来进行思考和分析。我们将考虑更一般的问题形式：\nMotivation: 在训练和推理时 LLM 所需要的内存是多少？如何进行优化内存占用？\n为了回答以上问题，我们先介绍训练/推理阶段的内存计算，再针对可优化部分进行分析并介绍相应优化算法。\nBackground Transformer Architecture 以 Qwen3 为例，现代 LLM 的架构包含多层 Transformer Block，其中具体的模块不同的模型可能有改动。下图是对应的模型架构\nNotation 与参数量、FLOPs 计算所用记号一致；参数量 \\(P\\) 的推导见 LLM parameter analysis。\n变量 含义 \\(P\\) number of parameters \\(L\\) layers \\(V\\) vocabulary size \\(d\\) hidden size \\(d_{\\text{ff}}\\) FFN hidden size \\(s\\) sequence length \\(b\\) batch size \\(h\\) number of attention heads \\(d_h\\) attention head dimension Assumptions 若无特别说明，使用 BF16/FP16，每个参数 2 byte。 不使用 dropout（与现代大模型设定一致）。 Attention 基于原始 multi-head attention. FFN 基于 SwiGLU. Training Memory Analysis Training Memory Components 训练部分的内存占用由四部分组成：\n$$ \\text{training\\_memory} = \\text{weight} + \\text{activation} + \\text{optimizer} + \\text{gradient} $$ Weights: \\(\\boxed{2P}\\) Gradients（与权重同精度）: \\(\\boxed{2P}\\) Optimizer States AdamW 优化器需要维护两个动量状态：\n一阶动量 \\(m_t\\)：\\(2P\\) 二阶动量 \\(v_t\\)：\\(2P\\) 合计：\\(2 \\times 2P = \\boxed{4P}\\) AdamW [2] 的更新规则如下：\n$$ \\begin{aligned} m_t \u0026\\leftarrow \\beta_1 \\cdot m_{t-1} + (1 - \\beta_1) \\cdot g_t \\\\ v_t \u0026\\leftarrow \\beta_2 \\cdot v_{t-1} + (1 - \\beta_2) \\cdot g_t^2 \\\\ \\hat{m}_t \u0026\\leftarrow \\frac{m_t}{1 - \\beta_1^t}, \\quad \\hat{v}_t \\leftarrow \\frac{v_t}{1 - \\beta_2^t} \\\\ \\theta_t \u0026\\leftarrow \\theta_{t-1} - \\alpha \\left( \\frac{\\hat{m}_t}{\\sqrt{\\hat{v}_t} + \\epsilon} + \\lambda \\theta_{t-1} \\right) \\end{aligned} $$Activation 激活值是前向传播过程中计算得到的中间结果，用于在反向传播时计算梯度。\n我们仅针对 linear layer 进行推导：\n$$ \\begin{aligned} \\text{forward:} \\quad \u0026 \\mathbf{z}_\\ell = W_\\ell \\mathbf{a}_{\\ell-1} + b_\\ell, \\quad \\mathbf{a}_\\ell = \\phi(\\mathbf{z}_\\ell) \\\\ \\text{backward:} \\quad \u0026 \\frac{\\partial \\mathcal{L}}{\\partial W_\\ell} = \\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{z}_\\ell} \\cdot \\frac{\\partial \\mathbf{z}_\\ell}{\\partial W_\\ell} = \\frac{\\partial \\mathcal{L}}{\\partial \\mathbf{z}_\\ell} \\cdot \\boxed{\\mathbf{a}_{\\ell-1}} \\end{aligned} $$可以看到，计算第 \\(\\ell\\) 层关于 \\(W_\\ell\\) 的梯度时需要其输入 \\(\\mathbf{a}_{\\ell-1}\\)，因此训练时需保存每个模块对应的输入，也就是激活值 (activation)。\nActivation — Attention 按计算图（无优化）可得需保存的激活：\nQ/K/V 投影：共享输入 → \\(2bsd\\) \\(Q^\\top K\\)：Q, K 均需保存 → \\(2 \\times 2bsd = 4bsd\\) softmax 输入：\\(2bhs^2\\) weighted sum 输入：\\(2bhs^2 + 2bsd\\) output projection 输入：\\(2bsd\\) Attention 合计： \\(\\boxed{10bsd + 4bhs^2}\\)\nActivation — FFN \u0026amp; LayerNorm FFN（SwiGLU，assume \\(d_{\\text{ff}} = 4d\\)）：\n第一层输入：\\(2bsd\\) SwiGLU 输入：\\(2 \\times d_{\\text{ff}} \\times d = 8bsd\\) 第二层输入：\\(2 \\times d_{\\text{ff}} \\times d = 8bsd\\) 合计：\\(\\boxed{18bsd}\\) LayerNorm：保存输入 → \\(\\boxed{2bsd}\\)\nActivation — Output Output 包含以下组成部分：\nFinalNorm 输入：\\(2bsd\\) lm_head 输入：\\(2bsd\\) Loss 输入：\\(2bsV\\) 合计：\\(\\boxed{4bsd + 2bsV}\\)\nActivation — Total 将上面的结果汇总在一起，得到：\n$$ \\begin{aligned} \\text{activation} \u0026= L \\cdot \\text{transformer\\_block} + \\text{output} \\\\ \u0026= L \\cdot (\\text{Pre\\_Norm} + \\textcolor{red}{\\text{Attention}} + \\text{Post\\_Norm} + \\text{FFN}) + \\text{output} \\\\ \u0026= \\boxed{bs(32dL + \\textcolor{red}{4hsL} + 4d + 2V)} \\\\ \u0026\\approx bsL(32d + \\textcolor{red}{4hs}) \\\\ \u0026\\approx \\textcolor{red}{4bs^2hL} \\end{aligned} $$ 注：在 Qwen3 中，\\(2V / 32dL \\approx 6\\%\\)，\\(32dL / 4hsL \\approx 2.5\\%\\)。\n可以看到，未优化的情况下，\\(\\text{activation} \\propto bs^2\\)。这里 \\(s^2\\) 主要由 attention 部分产生，后续 Flash Attention 就针对这一点进行了优化。\nTotal Training Memory 将上面的结果进行汇总：\n$$ \\begin{aligned} \\text{training\\_memory} \u0026= \\text{weight} + \\text{activation} + \\text{optimizer} + \\text{gradient} \\\\ \u0026= 2P + bs(32dL + 4hsL + 4d + 2V) + 4P + 2P \\\\ \u0026= 8P + bs(32dL + 4hsL + 4d + 2V) \\quad \\text{(exact)} \\\\ \u0026\\approx 8P + 4bs^2hL \\end{aligned} $$可以看到，训练阶段的内存占用分为固定部分 (\\(8P\\)) 和动态部分 (\\(4bs^2hL\\))，动态部分主要是 attention 的缓存。\nExperiments 我们分别针对 80GB 的显卡计算 Qwen3 系列模型的最高配置：\nModel \\(P\\) \\(L\\) \\(h\\) \\(s\\) predicted \\(b\\) actual \\(b\\) Qwen3-0.6B 0.6 28 16 512 68 34 Qwen3-1.7B 1.7 28 16 512 42 28 Qwen3-4B 4 36 32 512 16 12 Qwen3-8B 8.1 36 32 512 4 2 其中 predicted \\(b\\) 基于前面的准确公式计算得到；actual \\(b\\) 通过实验验证得到。注意我们这里的 prediction 没有考虑任何优化手段与其他内存开销，因此与实际值有出入。\nCase Study 我们分别使用 Qwen3-4B 和 Qwen3-8B 来进行实验（\\(b=1\\), \\(s=512\\)）。参考 PyTorch 显存可视化与 Snapshot 数据分析。\nInference Memory Analysis Inference Components Inference 阶段内存占用主要与模型参数、KV cache 两部分相关：\n$$ \\text{Inference\\_Memory} = \\text{weight} + \\text{activation} + \\text{KV cache} $$ Weights: \\(\\boxed{2P}\\) Activations（经验值，batch size=1）: \\(\\approx 0.4P\\)（参见 transformer-math） KV cache：与序列长度相关，见后续分析。 KV Cache Mechanism LLM 推理中为避免重复计算历史 token 的 key/value 而使用的空间换时间的缓存机制。\n自回归时逐 token 生成，每步 attention 形式为（\\(\\mathbf{q}_t\\) 当前 query，\\(\\mathbf{k}_{:,t}\\) / \\(\\mathbf{v}_{:,t}\\) 历史 K/V）：\n$$ \\mathbf{q}_t = W_Q \\mathbf{x}_t, \\quad \\mathbf{k}_{:,t} = W_K[\\mathbf{x}_1, \\ldots, \\mathbf{x}_t], \\quad \\mathbf{v}_{:,t} = W_V[\\mathbf{x}_1, \\ldots, \\mathbf{x}_t] $$处理下一 token \\(\\mathbf{x}_{t+1}\\) 时只需在已有结果后追加当前步：\n$$ \\mathbf{k}_{:,t+1} = [\\mathbf{k}_{:,t},\\, W_K \\mathbf{x}_{t+1}], \\quad \\mathbf{v}_{:,t+1} = [\\mathbf{v}_{:,t},\\, W_V \\mathbf{x}_{t+1}] $$ 缓存前：每生成一个 token 都重新计算 → 总计算量 \\(\\sum_{t=1}^{s} \\mathcal{O}(t) = \\mathcal{O}(s^2)\\) 缓存后：每步只算当前 token \\(W_K \\mathbf{x}_{t+1}\\)、\\(W_V \\mathbf{x}_{t+1}\\) → 计算量 \\(\\mathcal{O}(s)\\)、空间占用 \\(\\mathcal{O}(s)\\) —— 以空间换时间 KV Cache Memory 对于 multi-head attention，KV cache 的显存占用为：\n$$ \\text{Memory}(\\text{KV cache}) = s \\times 2 \\times 2 \\times L \\times h \\times d_h = \\boxed{4sLhd_h} $$因子含义： \\(s\\) 序列长，第一个 \\(2\\) 为 K+V，第二个 \\(2\\) 为 BF16 的 2 bytes，\\(L\\) 层、\\(h\\) 头、\\(d_h\\) 头维度。\nRemark:\nKV 占用与模型配置（\\(L, h, d_h\\)）和序列长度 \\(s\\) 都有关，token 越多占用越高。 实际中因 page granularity、padding、fragmentation 往往略高于理论值。 长输出时 KV 占比会超过权重，成为推理瓶颈 → 见后续「KV Cache 优化」。 Total Inference Memory 综合前面分析，推理阶段的总内存为：\n$$ \\boxed{\\text{Inference\\_Memory} \\approx 2.4P + 4sLhd_h} $$可以看到，推理阶段也由固定部分（参数量，activation）以及动态部分（KV cache）组成。\nDynamic vs. Static 由于 Qwen3 的 KV cache 计算为 \\(4sLh_{kv}d_h\\)，而不同模型只有 \\(L\\) 不一样，因此对于更大的模型，KV cache 显存占用超过模型权重的上下文长度更高。\nOptimization Overview 阶段 核心方法 典型技术 Training 显存与效率提升 Activation Checkpointing, Mixed Precision Training, Flash Attention, ZeRO, Pipeline/Model/Data Parallelism Inference 长序列与速度优化 KV Cache Optimization, Paged/Radix Attention, Faster Attention, Quantization Mixed Precision Training 计算量大的部分用低精度，计算量小的部分用高精度。低精度参与运算，高精度避免 Overflow/Underflow。\n下表是 DeepSeek-V3 [3] 使用的混合精度训练框架的显存分析：\nPrecision BF16 FP32 BF16 AMP No Yes Yes Weights BF16 (2) FP32 (4) BF16 (2) Master weights - - FP32 (4) Gradients BF16 (2) FP32 (4) BF16 (2) Adam m BF16 (2) FP32 (4) FP32 (4) Adam v BF16 (2) FP32 (4) FP32 (4) Static total (bytes/param) 8 16 16 Remark:\nBF16 (w/ AMP) 与 FP32 (w/ AMP) 的静态显存占用相同，但 BF16 (w/ AMP) 的动态显存占用更低。 主流框架基本都使用了 BF16/FP8 (w/ AMP) 的训练方式。 ZeRO 将 optimizer states / gradients / weights 按不同 GPU 切片存储，需要参与计算时再 all-gather 整合成完整参数。这样每张卡只需维护自己负责的一部分，大幅降低单卡显存需求 [4]。\nZeRO Stages:\nZeRO-1：shard optimizer states $$ \\text{Training\\_Memory} = \\text{weight} + \\text{activation} + \\frac{\\text{optimizer}}{\\#\\text{GPUs}} + \\text{gradient} $$ ZeRO-2：shard optimizer states + gradients $$ \\text{Training\\_Memory} = \\text{weight} + \\text{activation} + \\frac{\\text{optimizer} + \\text{gradient}}{\\#\\text{GPUs}} $$ ZeRO-3：shard all $$ \\text{Training\\_Memory} = \\text{activation} + \\frac{\\text{weight} + \\text{optimizer} + \\text{gradient}}{\\#\\text{GPUs}} $$ ZeRO-3 可极大降低单卡显存上限，但通信量也会提高。\nModel Parallelism 将模型切分到不同的 GPU 上，计算时，先 dispatch，再执行计算，最后通过 all-gather 等操作得到最终结果 [5]。切分方式包括 PP (Pipeline Parallelism)、TP (Tensor Parallelism)、EP (Expert Parallelism) 等。\n$$ \\text{Training\\_Memory} = \\frac{\\text{Memory}(\\text{weight})}{\\text{PP degree} \\times \\text{TP degree}} $$结合 ZeRO-1 与 Model Parallelism 时（activation 中与 TP 相关的部分按 TP degree 缩减）：\n$$ \\text{Memory}_{\\text{train}} \\approx \\frac{\\text{weight}}{\\text{PP} \\times \\text{TP}} + \\frac{\\text{activation}}{\\text{TP}} + \\frac{\\text{optimizer}}{\\#\\text{GPUs}} + \\frac{\\text{gradient}}{\\text{PP}} $$Activation Checkpointing 在反向传播时，重新计算所需的输入，来达到以时间换空间的目的 [6]。\nNo ckpt Selective ckpt Full ckpt memory 很高 中等 很低 \\(\\sim 2bsd\\) extra compute 无 中等 很高 \\(\\sim 2Pbs\\) 实践中常结合 model parallelism 与 selective checkpointing 来实现 trade-off。\nFlash Attention 通过将 Attention 的计算进行分块，来提高内存访问效率以及降低反向传播时所需要的 activation 大小 [7]。\nFlash Attention 通过 tiling 与 online-softmax 降低该部分显存并提升效率（详见 notes on Flash Attention）。这样 attention 部分的显存就由 \\(\\text{activation} \\propto bs^2\\) 降低到了 \\(\\text{activation} \\propto bs\\)。\nTheorem: Flash Attention 输出 \\(O = \\text{softmax}(QK^T)V\\)（correctness）。其时间复杂度为 \\(\\mathcal{O}(s^2 d)\\)，空间复杂度为 \\(\\mathcal{O}(s)\\)（memory savings）。\nKV Cache Optimization $$ \\text{Memory}(\\text{KV cache}) = s \\times 2 \\times 2 \\times L \\times h \\times d_h $$针对公式中各因子的优化方向 [8]：\n\\(s\\): KV cache compression, eviction, selection \\(2\\) (bytes): KV cache quantization \\(2\\) (K+V): key-value sharing, MLA [9] \\(h \\times d_h\\): MQA [10], GQA [11], MLA Weight Quantization 使用低精度来表示高精度数值的方法，来减少内存占用/提高计算效率。\n量化时机 代表性工作 训练后量化 (PTQ) GPTQ [12], AWQ [13], SmoothQuant [14], GGUF [15] 量化感知训练 (QAT) LLM-QAT [16], PEQA [17] Activation Offloading 将一部分参数/优化器状态/激活值等存储到 CPU 上，需要的时候再加载到 GPU 上。\nOffloading 场景 代表性工作 训练阶段 Offloading ZeRO-Offload [18] / ZeRO-Infinity, FSDP [19] CPU Offload 推理阶段 Offloading FlexGen [20], vLLM [21] KV Cache Offload MoE Offloading KTransformers [22], DeepSpeed-MoE [23] Real Systems Training Frameworks Framework Memory Optimizations Megatron-LM [5] TP, SP, Activation Checkpointing DeepSpeed [24] ZeRO-1/2/3, CPU/NVMe Offload, Activation Checkpointing FSDP [19] Full parameter sharding, Gradient sharding, CPU Offload Colossal-AI [25] ZeRO, TP, PP, Activation Checkpointing Inference Frameworks Framework Key Techniques Memory Optimizations vLLM [21] Paged Attention KV cache paging, Continuous batching SGLang [26] Radix Attention KV cache reuse, Efficient scheduling TensorRT-LLM [27] Kernel fusion Weight quantization, KV cache optimization Conclusion Takeaway Components Training Inference Optimization weights \\(2P\\) \\(2P\\) quantization optimizer states \\(4P\\) 0 ZeRO, offloading gradients \\(2P\\) 0 ZeRO activations \\(\\sim 4Lhs^2b\\) \\(\\sim 0.4P\\) ckpt, offloading, flash attention KV cache 0 \\(4sLhd_h\\) KV optimization, attention TOTAL \\(8P + 4Lhs^2b\\) \\(2.4P + 4sLhd_h\\) 训练：瓶颈主要在激活值（随 batch size / 序列长度线性增长）。 推理：瓶颈主要在 KV Cache（随序列长度增长）。 Future Directions More efficient architecture (attention, MoE). Scalable training/inference framework. Software-hardware co-design algorithms. References An Yang et al., \u0026ldquo;Qwen3 Technical Report,\u0026rdquo; arXiv:2505.09388, 2025. Ilya Loshchilov and Frank Hutter, \u0026ldquo;Decoupled Weight Decay Regularization,\u0026rdquo; arXiv:1711.05101, 2019. DeepSeek-AI, \u0026ldquo;DeepSeek-V3 Technical Report,\u0026rdquo; arXiv:2412.19437, 2025. Samyam Rajbhandari et al., \u0026ldquo;ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,\u0026rdquo; arXiv:1910.02054, 2020. Mohammad Shoeybi et al., \u0026ldquo;Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,\u0026rdquo; arXiv:1909.08053, 2020. Vijay Korthikanti et al., \u0026ldquo;Reducing Activation Recomputation in Large Transformer Models,\u0026rdquo; arXiv:2205.05198, 2022. Tri Dao et al., \u0026ldquo;FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,\u0026rdquo; arXiv:2205.14135, 2022. Haoyang Li et al., \u0026ldquo;A Survey on Large Language Model Acceleration based on KV Cache Management,\u0026rdquo; arXiv:2412.19442, 2025. DeepSeek-AI, \u0026ldquo;DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model,\u0026rdquo; arXiv:2405.04434, 2024. Noam Shazeer, \u0026ldquo;Fast Transformer Decoding: One Write-Head is All You Need,\u0026rdquo; arXiv:1911.02150, 2019. Joshua Ainslie et al., \u0026ldquo;GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,\u0026rdquo; arXiv:2305.13245, 2023. Elias Frantar et al., \u0026ldquo;GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers,\u0026rdquo; arXiv:2210.17323, 2023. Ji Lin et al., \u0026ldquo;AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration,\u0026rdquo; arXiv:2306.00978, 2024. Guangxuan Xiao et al., \u0026ldquo;SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models,\u0026rdquo; arXiv:2211.10438, 2024. Georgi Gerganov, \u0026ldquo;ggml: Tensor library for machine learning,\u0026rdquo; GitHub, 2023. Zechun Liu et al., \u0026ldquo;LLM-QAT: Data-Free Quantization Aware Training for Large Language Models,\u0026rdquo; arXiv:2305.17888, 2023. Jeonghoon Kim et al., \u0026ldquo;Memory-Efficient Fine-Tuning of Compressed Large Language Models via sub-4-bit Integer Quantization,\u0026rdquo; arXiv:2305.14152, 2023. Jie Ren et al., \u0026ldquo;ZeRO-Offload: Democratizing Billion-Scale Model Training,\u0026rdquo; arXiv:2101.06840, 2021. Yanli Zhao et al., \u0026ldquo;PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,\u0026rdquo; arXiv:2304.11277, 2023. Ying Sheng et al., \u0026ldquo;FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU,\u0026rdquo; arXiv:2303.06865, 2023. Woosuk Kwon et al., \u0026ldquo;Efficient Memory Management for Large Language Model Serving with PagedAttention,\u0026rdquo; SOSP, 2023. Hongtao Chen et al., \u0026ldquo;KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models,\u0026rdquo; SOSP, 2025. Samyam Rajbhandari et al., \u0026ldquo;DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale,\u0026rdquo; arXiv:2201.05596, 2022. Jeff Rasley et al., \u0026ldquo;DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters,\u0026rdquo; KDD, 2020. Shenggui Li et al., \u0026ldquo;Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training,\u0026rdquo; ICPP, 2023. Lianmin Zheng et al., \u0026ldquo;SGLang: Efficient Execution of Structured Language Model Programs,\u0026rdquo; NeurIPS, 2024. NVIDIA Corporation, \u0026ldquo;TensorRT-LLM: A TensorRT Toolset for Optimizing LLM Inference,\u0026rdquo; GitHub, 2023. ","date":"2026-01-17T10:04:32+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/llm-memory-computation/","title":"LLM Memory Computation"},{"content":"V100 V100 关键改进 Volta architecture SM architecture: 支持深度学习 2nd NVIDIA NVLink HBM2 memory Volta Multi-process Service V100 技术规格 Tesla Product Tesla K40 Tesla M40 Tesla P100 Tesla V100 GPU GK180 (Kepler) GM200 (Maxwell) GP100 (Pascal) GV100 (Volta) SMs 15 24 56 80 TPCs 15 24 28 40 FP32 Cores / GPU 2880 3072 3584 5120 FP64 Cores / GPU 960 96 1792 2560 Tensor Cores / GPU NA NA NA 640 GPU Boost Clock 810/875 MHz 1114 MHz 1480 MHz 1530 MHz Peak FP32 TFLOPS² 5 6.8 10.6 15.7 Peak FP64 TFLOPS² 1.7 .21 5.3 7.8 Peak Tensor TFLOPS² NA NA NA 125 Memory Size Up to 12 GB Up to 24 GB 16 GB 16 GB Memory Interface 384-bit GDDR5 384-bit GDDR5 4096-bit HBM2 4096-bit HBM2 TDP 235 Watts 250 Watts 300 Watts 300 Watts Manufacturing Process 28 nm 28 nm 16 nm FinFET+ 12 nm FFN 内存规格\nGPU Kepler GK180 Maxwell GM200 Pascal GP100 Volta GV100 Compute Capability 3.5 5.2 6.0 7.0 Threads / Warp 32 32 32 32 Max Warps / SM 64 64 64 64 Max Threads / SM 2048 2048 2048 2048 Max Thread Blocks / SM 32 32 32 32 Max 32-bit Registers / SM 65536 65536 65536 65536 Max Registers / Block 65536 65536 65536 65536 Max Registers / Thread 255 255 255 255 Max Thread Block Size 1024 1024 1024 1024 FP32 Cores / SM 192 128 64 64 Ratio of SM Registers to FP32 Cores 341 512 1024 1024 Shared Memory Size / SM 16 KB/32 KB/ 48 KB 96 KB 64 KB Configurable up to 96 KB 系统规格\nSpecification DGX-1 (Tesla P100) DGX-1 (Tesla V100) GPU 8x Tesla P100 GPUs 8x Tesla V100 GPUs TFLOPS 170 (GPU FP16) + 3 (CPU FP32) 1 (GPU Tensor PFLOP) GPU Memory 16 GB per GPU / 128 GB per DGX-1 Node 16 GB or 32 GB per GPU / 128-256 GB per DGX-1 Node CPU Dual 20-core Intel® Xeon® E5-2698 v4 Dual 20-core Intel® Xeon® E5-2698 v4 FP32 CUDA Cores 28,672 Cores 40,960 Cores System Memory Up to 512 GB 2133 MHz DDR4 LRDIMM Up to 512 GB 2133 MHz DDR4 LRDIMM Storage 4x 1.92 TB SSD RAID 0 4x 1.92 TB SSD RAID 0 Network Interconnect Dual 10 GbE, 4 IB EDR Dual 10 GbE, 4 IB EDR System Dimensions 866 D x 444 W x 131 H (mm) 866 D x 444 W x 131 H (mm) System Weight 80 lbs 80 lbs Max Power TDP 3200 W 3200 W Operating Temp 10 - 35°C 10 - 35°C A100 A100 关键改进 Ampere 架构：使用 MIG 来将 A100 切分为更小的实例或者链接更多 GPU Tensor Cores: 312 TFLOPs/s NVLink: 更高的 throughput MIG (multi-instance GPU): 一个 A100 可以切分为至多 7 个硬件层面隔离的实例 HBM2e: 更大的 HBM, 更快的 bandwidth, 更高的 DRAM 使用效率 structure sparsity: 稀疏运算可以带来 2 倍的算力提升 A100 技术规格 A100 80GB PCIe A100 80GB SXM FP64 9.7 TFLOPS 9.7 TFLOPS FP64 Tensor Core 19.5 TFLOPS 19.5 TFLOPS FP32 19.5 TFLOPS 19.5 TFLOPS Tensor Float 32 (TF32) 156 TFLOPS | 312 TFLOPS 156 TFLOPS | 312 TFLOPS* BFLOAT16 Tensor Core 312 TFLOPS | 624 TFLOPS* 312 TFLOPS | 624 TFLOPS* FP16 Tensor Core 312 TFLOPS | 624 TFLOPS* 312 TFLOPS | 624 TFLOPS* INT8 Tensor Core 624 TOPS | 1248 TOPS* 624 TOPS | 1248 TOPS* GPU Memory 80GB HBM2e 80GB HBM2e GPU Memory Bandwidth 1,935 GB/s 2,039 GB/s Max Thermal Design Power (TDP) 300W 400W *** Multi-Instance GPU Up to 7 MIGs @ 10GB Up to 7 MIGs @ 10GB Form Factor PCIe Dual-slot air-cooled or single-slot liquid-cooled SXM Interconnect NVIDIA® NVLink® Bridge for 2 GPUs: 600 GB/s ** PCIe Gen4: 64 GB/s NVLink: 600 GB/s PCIe Gen4: 64 GB/s Server Options Partner and NVIDIA-Certified Systems™ with 1-8 GPUs NVIDIA HGX™ A100-Partner and NVIDIA-Certified Systems with 4,8, or 16 GPUs NVIDIA DGX™ A100 with 8 GPUs H100 H100 关键改进 Hopper 架构 Tensor Core: 更强的 tensor core transformer engine: 加速基于 transformer 架构模型的训练 NVLink: 900GB/s 的 bandwidth 2nd MIG: 支持 multi-tenant, multi-user 使用 DPX: 基于 DPX 指令集加速动态规划算法 H100 技术规格 H100 SXM H100 NVL FP64 34 teraFLOPS 30 teraFLOPs FP64 Tensor Core 67 teraFLOPS 60 teraFLOPs FP32 67 teraFLOPS 60 teraFLOPs TF32 Tensor Core* 989 teraFLOPS 835 teraFLOPs BFLOAT16 Tensor Core* 1,979 teraFLOPS 1,671 teraFLOPS FP16 Tensor Core* 1,979 teraFLOPS 1,671 teraFLOPS FP8 Tensor Core* 3,958 teraFLOPS 3,341 teraFLOPS INT8 Tensor Core* 3,958 teraFLOPS 3,341 teraFLOPS GPU Memory 80GB 94GB GPU Memory Bandwidth 3.35TB/s 3.9TB/s Decoders 7 NVDEC 7 JPEG 7 NVDEC 7 JPEG Max Thermal Design Power (TDP) Up to 700W (configurable) 350-400W (configurable) Multi-Instance GPUs Up to 7 MIGS @ 10GB each Up to 7 MIGS @ 12GB each Form Factor SXM PCIe dual-slot air-cooled Interconnect NVIDIA NVLink™: 900GB/s PCIe Gen5: 128GB/s NVIDIA NVLink: 600GB/s PCIe Gen5: 128GB/s Server Options NVIDIA HGX H100 Partner and NVIDIA- Certified Systems™ with 4 or 8 GPUs NVIDIA DGX H100 with 8 GPUs Partner and NVIDIA-Certified Systems with 1–8 GPUs NVIDIA AI Enterprise Add-on Included H200 H200 关键改进 更高的 HBM 内存和带宽 更高的 LLM inference 速度 H200 技术规格 H200 SXM H200 NVL FP64 34 teraFLOPS 30 teraFLOPs FP64 Tensor Core 67 teraFLOPS 60 teraFLOPs FP32 67 teraFLOPS 60 teraFLOPs TF32 Tensor Core* 989 teraFLOPS 835 teraFLOPs BFLOAT16 Tensor Core* 1,979 teraFLOPS 1,671 teraFLOPS FP16 Tensor Core* 1,979 teraFLOPS 1,671 teraFLOPS FP8 Tensor Core* 3,958 teraFLOPS 3,341 teraFLOPS INT8 Tensor Core* 3,958 teraFLOPS 3,341 teraFLOPS GPU Memory 141GB 141GB GPU Memory Bandwidth 4.8TB/s 4.8TB/s Decoders 7 NVDEC 7 JPEG 7 NVDEC 7 JPEG Confidential Computing Supported Supported Max Thermal Design Power (TDP) Up to 700W (configurable) Up to 600W (configurable) Multi-Instance GPUs Up to 7 MIGS @ 18GB each Up to 7 MIGS @ 18GB each Form Factor SXM PCIe dual-slot air-cooled Interconnect NVIDIA NVLink™: 900GB/s PCIe Gen5: 128GB/s 2- or 4-way NVIDIA NVLink bridge: ** 900GB/s** per GPU\nPCIe Gen5: 128GB/s Server Options NVIDIA HGX H200 Partner and NVIDIA- Certified Systems™ with 4 or 8 GPUs NVIDIA MGX™ H200 NVL partner and NVIDIA-Certified Systems with up to 8 GPUs NVIDIA AI Enterprise Add-on Included 相比于 H100, H200 升级了 HBM 和 bandwidth\nB200 B200 关键改进 blackwell 架构： GPU 之间的通信效率大幅度提升 Grace CPU: GPU 可以与 Grace CPu 之间达到 900GB/s 的 bidirectional bandwidth 5th NVIDIA NVLink: 可以链接 576 块 GPU 来支持计算，NVlink 的带宽可以达到 130TB/s RAS engine: 自动识别故障来提高效率 NVIDIA networking B2100 技术规格 system specification 如下\nSpecification GB200 NVL72 GB200 NVL4 HGX B200 NVIDIA Blackwell GPUs | Grace CPUs 72 | 36 4 | 2 8 | 0 CPU Cores 2,592 Arm® Neoverse V2 Cores 144 Arm Neoverse V2 Cores - Total NVFP4 Tensor Core² 1,440 | 720 PFLOPS 80 | 40 PFLOPS 144 | 72 PFLOPS Total FP8/FP6 Tensor Core² 720 PFLOPS 40 PFLOPS 72 PFLOPS Total Fast Memory 31 TB 1.8 TB 1.4 TB Total Memory Bandwidth 576 TB/s 32 TB/s 62 TB/s Total NVLink Bandwidth 130 TB/s 7.2 TB/s 14.4 TB/s individual specification 如下\nSpecification GB200 NVL72 GB200 NVL4 HGX B200 FP4 Tensor Core 20 PFLOPS 20 PFLOPS 18 PFLOPS FP8/FP6 Tensor Core² 10 PFLOPS 10 PFLOPS 9 PFLOPS INT8 Tensor Core² 10 POPS 10 POPS 9 POPS FP16/BF16 Tensor Core² 5 PFLOPS 5 PFLOPS 4.5 PFLOPS TF32 Tensor Core² 2.5 PFLOPS 2.5 PFLOPS 2.2 PFLOPS FP32 80 TFLOPS 80 TFLOPS 75 TFLOPS FP64 / FP64 Tensor Core 40 TFLOPS 40 TFLOPS 37 TFLOPS GPU Memory Bandwidth 186 GB HBM3E 8 TB/s 186 GB HBM3E 8 TB/s 180 GB HBM3E 7.7 TB/s Multi-Instance GPU (MIG) - 7 - Decompression Engine - Yes - Decoders - 7 NVDEC³ 7 nvJPEG - Max Thermal Design Power (TDP) Configurable up to 1,200 W Configurable up to 1,200 W Configurable up to 1,000 W Interconnect - Fifth-generation NVLink: 1.8 TB/s PCIe Gen5: 128 GB/s - Server Options NVIDIA GB200 NVL72 partner and NVIDIA-Certified Systems™ with 72 GPUs NVIDIA MGX partner and NVIDIA-Certified Systems NVIDIA HGX B200 partner and NVIDIA-Certified Systems with 8 GPUs B300 B300 关键改进 Blackwell 架构 AI reasoning inference: 支持 test-time scaling, 对 attention layer 和 FLOPs 都有加速 HBM3e: 支持更大的 batch size 和 throughput ConnectX-8 SuperNIC, 一个 host2 个 ConnectX-8 设备，支持 800Gb/s 的 GPU 之间通信 Grace-CPU: 更强的表现和带宽 5th NVIDIA NVLink: 更高的通信效率 B3100 技术规格 system specification 如下\nGB300 NVL72 HGX B300 Blackwell Ultra GPUs| Grace CPUs 72 | 36 8 | 0 CPU Cores 2,592 Arm Neoverse V2 Cores - Total FP4 Tensor Core 1 1,440 PFLOPS | 1,080 PFLOPS 144 PFLOPS | 108 PFLOPS Total FP8/FP6 Tensor Core 2 720 PFLOPS 72 PFLOPS Total Fast Memory 37 TB 2.1 TB Total Memory Bandwidth 576 TB/s 62 TB/s Total NVLink Switch Bandwidth 130 TB/s 14.4 TB/s individual specification 如下\nGB300 NVL72 HGX B300 FP4 Tensor Core 20 PFLOPS | 15 PFLOPS 18 PFLOPS | 14 PFLOPS FP8/FP6 Tensor Core2 10 PFLOPS 9 PFLOPS INT8 Tensor Core2 330 TOPS 307 TOPS FP16/BF16 Tensor Core 5 PFLOPS 4.5 PLFOPS TF32 Tensor Core2 2.5 PFLOPS 2.2 PFLOPS FP32 80 TFLOPS 75 TFLOPS FP64/FP64 Tensor Core 1.3 TFLOPS 1.2 TFLOPS GPU Memory | Bandwidth 279 GB HBM3E | 8 TB/s 270 GB HBM3E | 7.7 TB/s Multi-Instance GPU (MIG) 7 7 Decompression Engine Yes Yes Decoders 7 NVDEC3 7 nvJPEG 7 NVDEC3 7 nvJPEG Max Thermal Design Power (TDP) Configurable up to 1,400 W Configurable up to 1,100 W Interconnect Fifth-Generation NVLink: 1.8 TB/s PCIe Gen6: 256 GB/s Fifth-Generation NVLink: 1.8 TB/s PCIe Gen6: 256 GB/s Server Options NVIDIA GB300 NVL72 partner and NVIDIA-Certified Systems™ NVIDIA HGX B300 partner and NVIDIA-Certified Systems References V100 white paper A100 Hopper Architecture H100 H200 B200 B300 blackwell ","date":"2026-01-14T11:09:19+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/nvidia-gpu-specs/","title":"Nvidia-GPU specs"},{"content":"Google 在 2022 年 8 提出了 GLaM，一个基于 MoE 架构的大语言模型系列，模型超过了 GPT-3 的表现\nIntroduction 作者在本文中证明了基于 MoE 架构的大语言模型可以达到与 dense 模型相同的性能，且计算更加高效。作者构建了 GLaM 大语言模型系列，包括 1.9B-A0.1B, 105B-A1.9B, 143B-A9.8B 以及 1.2T-A96.6B 等模型，尽管只激活了不到 $8\\%$ 的参数，模型的 zero-shot, one-shot 和 few-shot 表现均超过了 GPT-3.\n作者认为，在相同的算力下，MoE 模型比 dense 模型的表现更好，MoE 模型是一个非常具有前景的方向。\nMethod 训练数据集包括 1.6T token, 具体分布如下表所示\nDataset Tokens (B) Weight in mixture Filtered Webpages 143 0.42 Wikipedia 3 0.06 Conversations 174 0.28 Forums 247 0.02 Books 390 0.20 News 650 0.02 模型架构如下图所示\nGLaM 的模型架构与 GShard 基本相同，作者将 transformer block 按照两个为一组，一组中一个 block 为 dense FFN, 另一个为 MoE layer, 交替进行。MoE layer 中，总专家个数为 64 个，激活专家个数为 2 个。\n模型配置如下表所示\nGLaM Model Type $n_{\\text{params}}$ $n_{\\text{act-params}}$ $L$ $M$ $H$ $n_{\\text{heads}}$ $d_{\\text{head}}$ $E$ 0.1B Dense 130M 130M 12 768 3,072 12 64 $-$ 0.1B/64E MoE 1.9B 145M 12 768 3,072 12 64 64 1.7B Dense 1.7B 1.700B $-$ 1.7B/32E MoE 20B 1.878B 32 1.7B/64E MoE 27B 1.879B 24 2,048 8,192 16 128 64 1.7B/128E MoE 53B 1.881B 128 1.7B/256E MoE 105B 1.886B 256 8B Dense 8.7B 8.7B 32 4,096 16,384 32 128 $-$ 8B/64E MoE 143B 9.8B 32 4,096 16,384 32 128 64 137B Dense 137B 137B 64 8,192 65,536 128 128 $-$ 64B/64E MoE 1.2T 96.6B 64 8,192 32,768 128 128 64 作者用 GLaM(8B/64E) 来表示一个 8B 参数的 dense model 中每隔一层被转换为 MoE layer\nExperiments 训练时，作者设置模型的上下文为 1024 token, batch size 为 1M, 优化器为 Adafactor, 作者还使用了 GShard 中提出来的 load balancing loss, tokenizer 为 SentencePiece\nGLaM 与 GPT-3 的对比如下表所示\nGPT-3 GLAM relative cost FLOPs / token (G) 350 180 -48.6% Train energy (MWh) 1287 456 -64.6% accuracy on average Zero-shot 56.9 62.7 +10.2% One-shot 61.6 65.5 +6.3% Few-shot 65.2 68.1 +4.4% 实验结果显示，GLaM 的训练效率更高，且表现更好\n作者还探究了提升 expert 个数对最终表现的影响，结果如下图所示\n可以看到，随着专家个数提升，模型的表现逐渐增强，但是当专家个数超过 64 个之后，模型的表现反而有所下降。\nConclusion 作者在本文中提出了 GLaM 大语言模型系列，验证了 MoE 模型的有效性和效率，结果发现，MoE 模型可以在相同的算力下达到更好的表现，并且学习效率更高。\nReferences arxiv ","date":"2026-01-06T18:07:29+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-glam/","title":"Notes on GLaM"},{"content":"MiniMax-01 是一个基于 hybrid attention 架构的大模型系列，包含 MiniMax-Text-01 和 MiniMax-VL-01 两个模型，其中 MiniMax-Text-01 推理时支持 4M 的上下文长度，MiniMax-VL-01 支持 512B 的上下文长度\nIntroduction 现有大部分模型的上下文长度为 32K-256K, 但实际上我们对于长上下文的需求已经超过了这个范围。已有的模型主要基于 self-attention, 这是一个平方复杂度的算法。\n为了解决 self-attention 的问题，相关工作如 sparse attention, linear attention, state space models 等都提出了对应的解决办法。但是已有的这些解决方法的问题就是表现不是很强劲。\n因此，作者在本文中就提出了一个基于 hybrid attention 架构的大语言模型系列 MiniMax-01. 作者主要从架构，数据和 infra 三个方面进行了改进。\n在架构上，作者使用了基于 Lightning Attention 的混合架构。作者还基于实际部署来决定模型的参数。为了最大化模型的表现，作者使用了 MoE 架构。\n已有的 infra 主要是针对基于 softmax 的 attention 进行优化的，MiniMax-01 包含 softmax attention, linear attention 和 MoE, 架构比较复杂，因此作者实现了 expert parallel 和 expert tensor parallel 来提高整体的计算效率和 GPU 之间的通信效率。作者还实现了 varlen ring attention 来减少计算冗余。最终，作者发现模型在 NVIDIA-H20 上的 MFU 超过了 75%.\n基于前面提到的架构设计，作者训练得到了 MiniMax-Text-01, 模型总参数为 456B, 激活参数为 45.9B, 专家个数为 32 个，激活专家个数为 2 个。 作者首先构建了高质量的数据集，然后构建了一个三阶段的训练 pipeline。\n基于 MiniMax-Text-01, 作者扩展得到了 MiniMax-VL-01,训练使用了512B token\n作者总结本文的贡献如下：\n作者构建了一个领先的大模型系列，支持超过 4M 上下文 作者构建了第一个大规模的基于 linear attention 的大语言模型系列 作者详细介绍了使用的数据，模型和训练策略 作者开源了模型 Architecture 模型架构如下图所示\n相比于原始的 transformer, MiniMax-01 做出了以下改进：\n将 attention block 按照 8 个 block 为一组，每组里只有最后 1 个 block 使用 softmax attention, 其余 7 个 block 使用 lightning attention 使用 MoE 替换 FFN, MoE 总专家个数为 32 个，激活专家个数为 2 个 softmax attention 使用了 GQA, 来提高内存加载效率, group size 为 8 使用了 RoPE 作为 position embedding 使用了 RMSNorm 替换了 LayerNorm FFN: MoE (32 个专家，激活 2 个专家)\ntransformer: 8 个 block 为一组，一组里前 7 个使用 Lightning attention，第 8 个使用 softmax attention，80 layers\nAttention ： GQA，group size=8，64 heads，\nMoE 在训练基于 MoE 的 LLM 时，有两种策略，分别是 token-drop 和 dropless, 前者保证每个专家处理的 token 个数差不多，可以提高效率，缺点是某些 token 不会被任何专家处理，某些 token 会被多个专家处理；后者是保证每个 token 都会被处理。\n在本文中，作者采取了 token-drop 的方式，作者为每个专家设置一个 capacity limit，超过这个 limit 之后该专家就不再处理新的 token.\n为了评估 MoE 模型的有效性，作者对比了 MoE 和 dense 模型的表现，结果如下图所示：\n结果发现，相同的算力下，MoE 模型比 dense 模型好。\n但是，当 scaling up 到更大的模型是，作者发现训练产生了 routing collapse 的情况，这是因为 routing 的分布过于集中。为了解决这个问题，作者使用了 auxiliary loss 以及 global router 两个方法\nAuxiliary loss 作者首先构建了 auxiliary loss 来提高负载均衡, 也就是\n$$ \\mathcal{L}_{B} = K\\sum_{i=1}^K f_i P_i $$这里 $f_i$ 是分配给第 $i$ 个专家的 token 比例， $P_i$ 是第 $i$ 个专家的平均 routing 概率。\nGlobal router 作者还基于 GShard 构建了一个 global routing 策略，由于 GPU memory 限制，对于每一个 micro batch size, token 分布不仅在一个 EP group 内分布不平衡，在不同 EP group 之间可能也不平衡。因此，作者实现了一个 global token dispatching 策略。具体来说，在 token 分发到不同的 EP group 之前，作者使用 allgather 来计算每个专家需要处理的 token. 这样就可以基于全局信息智能分配 token, 避免某些专家过载。\nLinear Attention 本文中使用的 linear attention 由 Lightning Attention 提出，其表达式为\n$$ O = \\mathrm{Norm}((QK^T)V) $$这里 $Q,K,V\\in\\mathbb{R}^{n\\times n}$ 分别是 query, key 和 value, $n$ 和 $d$ 分别是序列长度和 hidden size.上式可以改变运算顺序，得到\n$$ O = \\mathrm{Norm}(Q(K^TV)) $$这样，attention 的计算复杂毒就从 $O(n^2d)$ 变成了 $O(nd^2)$,\nLightning Attention 当我们不考虑 attention mask 的时候，我们很轻松可以降低 attention 计算的复杂度。但实际上，LLM 会使用 causal mask, 也就是每个 token 智能看到其前面 token 的信息，这样我们的 attention 计算实际上是\n$$ O = \\mathrm{Norm}[QK^T\\odot M]V) $$这里 $M_{ij}=\\mathbb{1}(i\\geq j)$ 是 attention mask.\n见 Lightning Attention\nEffectiveness of Lightning Attention 作者接下来分析了一下 softmax attention, lightning attention 和 hybrid attention 之间的效率\n首先，作者计算了一下三种架构的参数量以及 FLOPS. 作者分别使用 $l, d, h, b, n$ 来代表 Layer 数，hidden dimension, number of attention heads, batch size 和 sequence length.\n最终计算结果如下\nArchitecture Parameter count FLOPs count Softmax $12ld^2$ $72bnld^2 (1 + \\frac{n}{6d} + \\frac{5}{18d} )$ Lightning $12ld^2+2ld^2/h$ $72bnld^2(1+\\frac{1}{2h}+\\frac{5}{18d})$ Hybrid $12ld^2+7ld^2/4h$ $72bnld^2(1+\\frac{n}{48d}+\\frac{7}{16h}+\\frac{5}{18d})$ [!todo] TODO compute these results\nScaling law 作者接下来分别针对三个架构设计了 70M, 160M, 410M, 1B, 3B, 7B 系列模型，模型使用 300B token 进行训练，上下文长度为 8192, 对于每种架构，作者将 batch size 设置为 4M tokens. 与 Chinchilla scaling law 类似， 作者探究了以下模型的 scaling law, 结果如下图所示\n从实验结果可以看到，在相同的算力下，lighting attention 倾向于使用更多的参数和 token, 但是其表现相比于 softmax attention 更好。\nPerformance 作者还对比了一下三种 attention 在 public benchmark 上的表现，结果如下图\n是检验结果发现，lightning attention 和 softmax attention 除了在 retrieval 任务（Needle in a Haystack）上之外，表现都差不多。与之相对的是，hybrid attention 弥补了这一问题，大幅度提升了模型在 retrieval 任务上的表现\n[!tip] Observation Lightning attention 与 softmax attention 的效果差不多，但是其在长上下文任务上的表现比较差。Hybrid attention 可以解决这个问题。\n作者还评估了一下三种 attention 的速度，结果如下图所示\n实验结果显示，softmax attention 随序列长度上升其速度急剧下降，Lightning attention 的速度基本没有太大变化。而 Hybrid attention 的速度介于两者之间。\n作者进一步对比了以下两种不同的变体：hybrid-cosformer2 以及 hybrid-hgrn2.这两个模型替换的逻辑与 minimax-01 的结果一致，都是 8 个 block 为一组，每组中前面 7 个 block 将 softmax attention 更换为对应的模块，最后一层不做改动。三种 hybrid attention 机制的表现如下表所示\n实验结果显示，hybrid-lightning 的表现最好。\n作者最后对比了以下 hybrid-lightning 和 hybrid-window, hybrid window 在每个 group 的前 7 个 block 使用了 window attention, window size 为别为 256, 512, 1024.实验结果如下图所示\nDiscussion 作者最后总结认为，虽然 linear attention 的效率很高，但是它们在 retrieval 相关的任务表现很差，而 retrieval 对于 In-context learning 来说是至关重要的。因此，作者采取了 hybrid 架构，来兼顾模型的效率以及表现。\n作者给出了一个解释。\nAblation Study 作者主要进行了两个消融实验：\nHybrid-lightning 与 softmax attention 的对比：作者训练了一个总参数 28B, 激活参数 5B 的 MoE 模型，然后作者使用 MiniMax-01 的方式替换每个 group 的 softmax attention, 并使用了 1T 的 token 进行训练 pre-layer normalization 与 Post-layer normalization 的对比： 现有的 LLaMA 和 Qwen 等系列模型采用的都是 PreNorm 的方式。PreNorm 可以让 gradient 通过 residual connection 传播更加直接，但是这也减少了模型的有效深度。反之，PostNorm 则可以保留模型的有效深度，但是其问题是会导致梯度消失或爆炸。作者构建了一个总参数为 60B, 激活参数为 9.3B 的 MoE 模型，包含 48 个 block, 模型训练使用 500B token. 模型有两个变体，一个使用 PreNorm, 另一个使用 PostNorm, 对于 PostNorm, 作者使用的是 DeepNorm. 最终表现如下图所示\n实验结果显示，hybrid lightning 的表现与 softmax 的表现相当，甚至超过了 softmax 的表现。\n另一方面，PostNorm 的表现也超过了 PreNorm 的表现。\nModel Spec 基于已有的模型设计，作者探究了如何决定模型的参数。作者的目标在 performance 和 inference efficiency 之间达到一个平衡。\n作者将模型的参数限制在 500B 以下，要求能够在 $8\\times 80G$ 的服务器上和 8-bit 的量化下面，支持 1M 的上下文长度。建模的问题如下：\n$$ \\min_{P_{all}, P_{act}}\\mathcal{L}(P_{all}, P_{act}, T), \\quad \\mathrm{s.t.}\\ C_{compute}(P_{all}, P_{act}, T)","date":"2026-01-06T17:38:01+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-minimax-01/","title":"Notes on MiniMax-01"},{"content":"Introduction 作者首先回顾了开源模型如 MiniMax-01, Kimi-k2, Qwen3, GLM-4.5 和闭源模型的进展，作者指出，现在的开源模型和闭源模型在表现上仍然存在较大差距。作者认为这种差距主要是由于三个原因：\nTransformer 提出的 softmax attention 在处理长文本时效率非常低 已有的开源模型在 post-training 阶段使用的算力不够 开原模型的泛化和指令跟随能力不如闭源模型 基于这三个问题，DeepSeek-V3.2 分别进行了改进：\n在架构上，作者提出了 DSA，一个高效的稀疏注意力机制，用于降低计算复杂度 在 post-training 阶段，作者使用了比 pre-training 阶段高 $10\\%$ 的算力，用于提高模型的能力 作者提出了一个 pipeline 用于提高模型在工具调用场景下的 reasoning 能力 通过实验作者发现，模型达到了和 Kimi-k2 以及 GPT-5 差不多的 reasoning 表现。\nMethod Architecture DeepSeek-V3.2 与 DeepSeek-V3.1 不同之处在于使用了 DeepSeek Sparse Attention (DSA). 架构如下图所示\nDSA 包含两个模块：\nlightning indexer fine-grained token selection mechanism 其中，lightning indexer 负责计算 query token $h_t\\in\\mathbb{R}^d$ 和一个 preceding token $h_s\\in\\mathbb{R}^d$ 之间的 index score $I_{t,s}$ 来决定 query token 选择的 token:\n$$ I_{t,s} = \\sum_{j=1}^{H_I}w_{t,j}^I \\mathrm{ReLU}(q_{t,j}^I\\cdot k_s^I) $$其中， $H^I$ 代表 indexer heads 的个数，$q_{t,j}^I\\in\\mathbb{R}^{d^I}$ 和 $w_{t,j}^I$ 由 query token $h_t$ 得到，$k_s^I\\in\\mathbb{R}^{d^I}$ 由 preceding token $h_s$ 得到\n给定 query token $h_t$ 对应的 index score $\\{I_{t,s}\\}$, fine-grained token selection mechanism 负责选取 top-K index score 对应的 key-value entries $\\{c_s\\}$, 然后 attention 的输出由 query token 个选取的 key value entries 得到：\n$$ u_t = \\mathrm{Attn}(h_t,\\{c_s\\mid I_{t,s}\\in \\mathrm{TopK}(I_{t,:})\\}) $$ [!Recall] MoBA 也提出了类似的方法，但是 MoBA 是一个无需训练的策略\n受 NSA 启发，作者实现了基于 MQA 模式的 MLA, 其中 latent vector 对于 query token 所有的 query heads 都是共享的。示意图如下所示\nContinue Pre-training 作者在 DeepSeek-V3.1 的基础上进行了 continue pre-training. Continue pre-training 包含两个阶段：\nDense Warm-up stage 这个阶段用于训练 lightning indexer, 作者冻结除 lightning indexer 之外的参数，为了对齐 indexer output 和 main attention distribution, 对于第 $t$ 个 query token, 作者首先计算所有 attention heads 的 main attention score 之和，然后在 sequence 层面进行 L1-normalization 得到 $p_{t,:}\\in\\mathbb{R}^t$, 最后计算 lightning indexer 输出与 $p_{t,:}$ 之间的 KL divergence:\n$$ \\mathcal{L}^I = \\sum_{t}\\mathcal{D}_{KL}(p_{t,:}\\ \\Vert\\ \\mathrm{Softmax}(I_{t,:})) $$这个阶段训练一共使用了 2.1B 的 token, lr 为 1e-3, 训练的步数为 1000 steps, batch size 为 16.\nSparse training stage 这个阶段模型所有的参数都参与训练，该阶段的目的是让模型学习到 DSA 的 sparse pattern. 训练时，作者让 lightning indexer 的输出与 $p_{t,S_t}$ 之间的输出进行对齐，其中 $S_t=\\{s\\mid I_{t,s}\\in\\mathrm{TopK}(I_{t,:})\\}$:\n$$ \\mathcal{L}^I = \\sum_{t}\\mathcal{D}_{KL}(p_{t,S_t}\\ \\Vert\\ \\mathrm{Softmax}(I_{t,:})) $$实际训练时，lightning indexer 仅接受 $\\mathcal{L}^I$ 的反向传播，而 LLM 则仅接受 next-token prediction loss. 这个阶段模型一共使用了943.7B token, 其中 $K$ 设置为 $2048$. 学习率为 $7.3\\times 1e-6$, 训练步数为 15,000 steps, batch size 为 480.\nPost-training post-training 与 DeepSeek-V3.1 一致：\nSpecialist Distillation 作者基于 DeepSeek-V3.2 base 构建了不同领域的 specialized model, 这些领域包括：\nmath competitive programming general logical reasoning agentic coding agentic search 每个 specialized model 都使用 RL 进行训练，训练数据包括 long CoT reasoning 数据以及 direct response generation 数据，specialized model 训练完毕之后，就被用于生产 domain-specific data, 作者通过实验发现，基于这种蒸馏方法，模型的表现仅比 specialized model 低一点，并且这个 gap 可以被后续的 RL 训练所抵消。GLM-4.5 也采取了类似的做法\nMixed RL Training 作者使用了 GRPO 算法进行训练，与 DeepSeek-V3 不同，作者将 reasoning, agent 以及 human alignment 的 RL 训练合并为了一个阶段，作者认为这种方法可以平衡模型在多个 domain 上的表现，并且可以防止 multi-stage training 带来的灾难性遗忘问题。对于 reasoning 和 agent 任务，作者使用了 rule-basd outcome reward, length penalty 以及 language consistency reward. 对于通用任务，作者使用了 generative reward model, 每个 prompt 都有对应的 rubris 用于 evaluation. 作者构建 reward 时主要考虑了:\nlength versus accuracy language consistency versus accuracy DeepSeek-V3.2-Speciale 除了 DeepSeek-V3.2 之外，作者还训练了 DeepSeek-V3.2-Speciale 模型，该模型仅使用 reasoning 数据进行训练，reasoning 数据包含了 DeepSeek-Math-V2 的训练数据以及 reward 方法。训练时，作者降低了 length penalty 的惩罚系数，最终 DeepSeek-V3.2-Speciale 模型拥有更强的 reasoning 能力\nScaling GRPO 作者在 GRPO 的基础上对 KL estimate 进行了改进（见 KL divergence），使用了 importance sampling 对 K3 estimator 进行修正:\n$$ \\mathcal{D}_{\\mathrm{KL}}(\\pi_\\theta(o_{i,t})\\Vert\\pi_{\\mathrm{ref}}(o_{i,t})) = \\frac{\\pi_\\theta(o_{i,t}\\mid q, o_{i,","date":"2026-01-06T17:30:40+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseek-v3.2/","title":"Notes on DeepSeek-V3.2"},{"content":"Gemini 3.0 是是 Google 新一代最强模型，model card 介绍了 Gemini 3.0 系列的评估结果以及基本能力\nIntroduction Gemini 3.0 系列包含\nGemini 3.0 Pro Gemini 3.0 Flash Gemini 3.0 Pro Image 三个模型 Gemini 3.0 Pro 拥有原生多模态以及 reasoning 能力，可以处理 text, audio, images, video 以及 code repositories 等模态。\nmodalities context input text, images, audio, video 1M output text 64K Gemini 3.0 Flash 与 Gemini 3.0 Pro 基本一致，与 Gemini2.5 相同，应该是采取了蒸馏的方式来实现更高的吞吐速度以及效率\nGemini 3.0 Pro Image 基于 Gemini 3.0 Pro 开发，是一个支持 text, image prompt 的图片生成模型\nMethod 模型从零开始训练，使用了 MoE 架构和 Transformer 架构\n模型使用 TPU 进行训练，训练架构为 JAX 和 ML Pathways.\nExperiments Gemini 3.0 Pro 对比了 Gemini2.5 , Claude Sonnet 4.5 和 GPT-5.1\nGemini 3.0 Flash 对比了 Gemini 3.0 Pro, Gemini 2.5 Flash, Gemini 2.5 Pro, Claude Sonnet 4.5, GPT-5.2 和 Grok 4.1 Fast.\nGemini 3.0 Pro Image 对比了 Gemini 2.5 Flash Image, GPT-Image 1, Seedream v4, Flux Pro Kontext Max\nReferences Gemini 3.0 Pro Model Card Gemini 3.0 Flash Model Card Gemini 3.0 Pro Image Model Card ","date":"2026-01-06T10:26:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gemini3.0/","title":"Notes on Gemini3.0"},{"content":"Introduction softmax 函数用于将 $K$ 个实数转换为一个 $K$ 维概率分布。其具体做法是先对所有元素指数化，即求 $e^x$, 然后每个元素除以所有指数的和。即\n$$ \\begin{aligned} \\mathrm{softmax}:\\mathbb{R}^K\u0026\\to (0,1)^K\\\\ \\mathrm{softmax}(\\mathbf{z}) \u0026=\\left[\\frac{e^{z_1}}{\\sum_{j=1}^Ke^{z_j}},\\dots,\\frac{e^{z_K}}{\\sum_{j=1}^Ke^{z_j}}\\right] \\end{aligned} $$Analysis Properties softmax 的第一个性质是 shift invariance, 即\n$$ \\mathrm{softmax}(\\mathbf{z}+c) = \\mathrm{softmax}(\\mathbf{z}) $$证明比较容易：\n$$ \\mathrm{softmax}(\\mathbf{z}+c)_i = \\frac{e^{z_i+c}}{\\sum_{j=1}^Ke^{z_j+c}} = \\frac{e^ce^{z_i}}{e^c\\sum_{j=1}^Ke^{z_j}} = \\frac{e^{z_i}}{\\sum_{j=1}^Ke^{z_j}}=\\mathrm{softmax}(\\mathbf{z})_i,\\ i=1,\\dots,K $$Gradient 向量输入下 Softmax 函数的 Jacobian 矩阵推导\n设输入为向量 $\\mathbf{z} = [z_1, z_2, \\dots, z_d]^\\top \\in \\mathbb{R}^d$，Softmax 函数的输出为向量 $\\mathbf{a} = [a_1, a_2, \\dots, a_d]^\\top \\in \\mathbb{R}^d$，其中每个元素定义为：\n$$ a_j = \\text{softmax}(\\mathbf{z})_j = \\frac{e^{z_j}}{\\sum_{k=1}^d e^{z_k}} $$记分母（归一化因子）为 $S = \\sum_{k=1}^d e^{z_k}$，则 $a_j = e^{z_j}/S$.\n我们分两种情况计算 $\\frac{\\partial a_j}{\\partial z_k}$：\n当 $j = k$ 时， 此时求 $a_j$ 对自身输入 $z_j$ 的偏导数：\n$$ \\frac{\\partial a_j}{\\partial z_j} = \\frac{\\partial}{\\partial z_j} \\left( \\frac{e^{z_j}}{S} \\right) = \\frac{e^{z_j}S-e^{z_j}e^{z_j}}{S^2}=\\frac{e^{z_j}}{S}\\left(1-\\frac{e^{z_j}}{S}\\right)=a_j(1-a_j) $$当 $j \\neq k$ 时， 此时求 $a_j$ 对输入 $z_k$ 的偏导数有：\n$$ \\frac{\\partial a_j}{\\partial z_k} = \\frac{\\partial}{\\partial z_k} \\left( \\frac{e^{z_j}}{S} \\right) = \\frac{0\\cdot S-e^{z_j}e^{z_k}}{S^2}=-\\frac{e^{z_j}e^{z_k}}{S}=-a_ja_j $$综合以上两种情况，Jacobian 矩阵 $\\mathbf{J}$ 可表示为：\n$$ \\mathbf{J} = \\text{diag}(\\mathbf{a}) - \\mathbf{a} \\mathbf{a}^\\top $$Interpretation Soft Argmax softmax 是 argmax 的 smooth approximation, 所以实际上 softmax 指的是 “soft argmax\u0026quot;. 为了证明这一点，我们首先定义如下函数\n$$ \\mathrm{softmax}(\\mathbf{z};\\tau) =\\mathrm{softmax}(\\mathbf{z}/\\tau)=\\left[\\frac{e^{z_1/\\tau}}{\\sum_{j=1}^Ke^{z_j/\\tau}},\\dots,\\frac{e^{z_K/\\tau}}{\\sum_{j=1}^Ke^{z_j/\\tau}}\\right] $$易知， $\\mathrm{softmax}(\\mathbf{z})=\\mathrm{softmax}(\\mathbf{z};1)$. 并且，$\\mathrm{softmax}$ 还是一个光滑函数\n我们定义 smooth approximation 为\nDefinition 如果 $\\lim_{\\tau\\to0^+}\\mathrm{softmax}(\\mathbf{z};\\tau)=\\mathbb{1}_{\\arg\\max(\\mathbf{z})}$, 则我们说 $\\mathrm{softmax}(\\cdot;\\tau)$ 是 $\\arg\\max$ 的光滑近似，特别地，$\\mathrm{softmax}(\\cdot)$ 是 $\\arg\\max$ 的光滑近似。 这里 $\\arg\\max(\\mathbf{z})=\\arg\\max_k z_k$ 是最大值的索引， $\\mathbb{1}\\in\\{0,1\\}^K$ 是示性函数 (indicator function), 即 $\\mathbb{1}_{\\arg\\max(\\mathbf{z})}[i]=1$ 当且仅当 $z_i=\\max_jz_j$.\n我们下面来进行证明。我们不妨假设最大值唯一，其 index 为 $m$, 即 $z_m = \\max_i z_i$. 由前面的性质，我们有：\n$$ \\mathrm{softmax}(\\mathbf{z};\\tau) = \\mathrm{softmax}(\\mathbf{z}-z_m;\\tau) =\\left[\\frac{e^{(z_1-z_m)/\\tau}}{\\sum_{j=1}^Ke^{(z_j-z_m)/\\tau}},\\dots,\\frac{e^{(z_K-z_m)/\\tau}}{\\sum_{j=1}^Ke^{(z_j-z_m)/\\tau}}\\right] $$此时，我们有\n$$ \\lim_{\\tau\\to0^+}\\mathrm{softmax}(\\mathbf{z};\\tau)_i = \\begin{cases} 1, \u0026\\text{if }i = m\\\\ 0, \u0026\\text{otherwise} \\end{cases} $$当最大值不唯一的时候，我们记 $\\mathcal{I} = \\{i\\in[K]\\mid z_i=\\max_j z_j\\}$, 与上面方法类似，最终 $\\mathrm{softmax}(\\cdot;\\tau)$ 的结果为\n$$ \\lim_{\\tau\\to0^+}\\mathrm{softmax}(\\mathbf{z};\\tau)_i = \\begin{cases} 1/|\\mathcal{I}|, \u0026\\text{if }i \\in \\mathcal{I}\\\\ 0, \u0026\\text{otherwise} \\end{cases} $$因此，我们就证明了 softmax 是 argmax 函数的 smooth approximation.\nStatistical Mechanics Temperature 我们前面介绍了 $\\mathrm{softmax}(\\mathbf{z};\\tau)$ 函数，这里的 $\\tau$ 实际上被称为温度 (temperature), 它控制了输入的 variance, $T$ 越大，输入的 variance 越低，输出就倾向于均匀分布，而 $T$ 越小，则说明输入的 variance 越高，输出就倾向于 one-hot 分布。\n我们前面已经证明了后者，现在我们来证明一下前者，证明思路也很简单，$T\\to+\\infty$ 时，$e^{x/T}\\to 1$, 因而\n$$ \\lim_{\\tau\\to+\\infty}\\mathrm{softmax}(\\mathbf{z};\\tau)_i =\\frac1K,\\ i=1,\\dots,K $$下面是可视化的代码以及结果\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import make_interp_spline def softmax(x): e_x = np.exp(x - np.max(x)) return e_x / e_x.sum() num_elements = 15 indices = np.arange(num_elements) logits = np.linspace(-3.5, 3.5, num_elements) scales = [0.01, 0.1, 1.0, 5.0, 10.0, 100.0] plt.figure(figsize=(10, 6)) for s in scales: probs = softmax(logits * s) x_smooth = np.linspace(indices.min(), indices.max(), 300) spl = make_interp_spline(indices, probs, k=3) y_smooth = np.clip(spl(x_smooth), 0, None) # Clip to ensure no negative artifacts plt.plot(x_smooth, y_smooth, label=f\u0026#39;Scale = {s}\u0026#39;, linewidth=2) uniform_prob = 1.0 / len(indices) plt.axhline(y=uniform_prob, color=\u0026#39;black\u0026#39;, linestyle=\u0026#39;:\u0026#39;, alpha=0.6, label=f\u0026#39;Uniform distribution\u0026#39;) plt.xticks(indices) plt.xlabel(\u0026#39;Logit Index\u0026#39;, fontsize=12) plt.ylabel(\u0026#39;Softmax Probability\u0026#39;, fontsize=12) plt.title(\u0026#39;Impact of Variance Scaling on Softmax Distribution\u0026#39;, fontsize=14) plt.legend(title=\u0026#34;Variance Scale\u0026#34;) plt.grid(True, linestyle=\u0026#39;--\u0026#39;, alpha=0.5) plt.tight_layout() plt.show() 可以看到，当 variance 比较小的时候，输出的分布接近于均匀分布，而 variance 越大，输出的分布越接近 One-hot 分布。\n在 attention 的计算过程中，我们也有 softmax 函数，为了在 softmax 过程中避免 variance 的影响，现在会在计算 softmax 之前加入 normalization layer 来提前进行归一化。见 QK-norm.\nAlgorithms Implementation 由于 $e^x$ 在实际计算时，非常容易溢出，因此在实现的时候，我们往往会考虑其数值稳定性。实际上，现在的 softmax 函数基本由 logsumexp 实现，logsumexp 函数定义如下\n$$ \\mathrm{logsumexp}(\\mathbf{z}) = \\log \\left(\\sum_{i=1}^K e^{z_i}\\right) $$softmax 函数与 logsumexp 函数的关系如下\n$$ \\begin{aligned} \\mathrm{softmax}(\\mathbf{z}) \u0026=\\exp\\log\\left(\\frac{e^{\\mathbf{z}}}{\\sum_{j=1}^Ke^{z_j}}\\right)\\\\ \u0026= \\exp\\left(\\mathbf{z} - \\log\\left(\\sum_{i=1}^K e^{z_i}\\right)\\right)\\\\ \u0026= \\exp(\\mathbf{z} - \\mathrm{logsumexp}(\\mathbf{z})) \\end{aligned} $$考虑前面提到的 $e^x$ 数值溢出的问题，我们的输入会先经过 shift, 减掉最大值。此时我们有\n$$ \\mathrm{softmax}(\\mathbf{z}) = \\mathrm{softmax}(\\mathbf{z}-c) = \\exp((\\mathbf{z}-c) - \\mathrm{logsumexp}(\\mathbf{z}-c)) $$这里我们使用了前面推导出来的 shift invariance 性质。对应的代码实现如下：\n1 2 3 4 def softmax(x: torch.Tensor, dim: int = -1) -\u0026gt; torch.Tensor: x = x - x.max(dim=dim, keepdim=True).values log_sum_exp = torch.log(torch.sum(torch.exp(x), dim=dim, keepdim=True)) return torch.exp(x - log_sum_exp) Gumbel-softmax Reparametrization Trick TODO\nOnline Softmax 注意到我们在计算 softmax 时，需要加载 $\\mathbf{z}$ 的全部信息，如果 $\\mathbf{z}$ 非常大的话，会产生频繁的内存读写进而影响整体效率。因此 flash attention 中提出了 online softmax 算法来减少内存访问开销。\n其具体做法是假设我们的输入被分为若干个 block, 即 $\\mathbf{z}=[\\mathbf{z}^1;\\dots,\\mathbf{z}^n]\\in\\mathbb{R}^K$, 这里 $\\mathbf{z}^i\\in\\mathbb{R}^{K/n}$ ($K\\mod n=0$).\n对于 $\\mathbf{z}\\in\\mathbb{R}^K$, flash attention 定义如下结果\n$$ m(\\mathbf{z}) = \\max_i z_i,\\ f(\\mathbf{z}) = [e^{z_1-m(\\mathbf{z})},\\dots,e^{z_K-m(\\mathbf{z})}], \\ \\ell(\\mathbf{z})=\\sum_if(z)_i, \\ \\mathrm{softmax}(\\mathbf{z}) = \\frac{f(\\mathbf{z})}{\\ell(\\mathbf{z})} $$对于 $\\mathbf{z}=[\\mathbf{z}^1;\\dots,\\mathbf{z}^n]\\in\\mathbb{R}^K$, 我们现在的计算方式为\n$$ \\begin{aligned} m_i(\\mathbf{z}) \u0026= \\max([\\mathbf{z}^1;\\dots;\\mathbf{z}^i]) = \\max(m_{i-1}(\\mathbf{z}),m(\\mathbf{z}^i))\\\\ \\ell_i(\\mathbf{z}) \u0026= \\sum_{j=1}^if(\\mathbf{z}^j) = \\exp(m_{i-1}(\\mathbf{z}) - m_i(\\mathbf{z}))\\ell(\\mathbf{z}^{i-1}) + \\exp(\\mathbf{z}^i-m_i(\\mathbf{z})) \\end{aligned} $$因此，如果我们额外记录 $m(x)$ 以及 $\\ell(x)$ 这两个量，那么我们可以每次仅计算 softmax 的一个 block. 计算完毕之后，$m_i(\\mathbf{z})$ 和 $\\ell_i(\\mathbf{z})$ 就分别代表了 global max 和 global denominator.\nConclusion 我们回顾了机器学习中 softmax function 的基本定义与性质\nReferences Softmax function Softmax to the Max Online normalizer calculation for softmax ","date":"2025-12-27T16:39:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-softmax/","title":"Notes on Softmax"},{"content":"Introduction 【参考文献 1】中系统性对比了 AliBi, RoPE , T5 提出的 T5 bias 以及 Transformer 提出的绝对位置编码 (APE).\n作者发现，常用的方法在 length generalization 上表现并不是最好的，而 NoPE 不需要额外的计算开销反而效果最好。\n【参考文献 2】 进一步探究了 NoPE 长度外推的泛化性。作者有三点发现：\nNoPE 相比于 RoPE, 其长度外推泛化能力更强 对于 NoPE 来说，模型会在还没有到达预训练上下文长度之前，表现就出现下降的情况 通过调整 softmax 的温度超参数，我们可以提高 NoPE 的长度外推泛化性能力。 Method 【参考文献 1】对比了不同 position encoding 的相似度，结果如下图所示\n实验结果表明，NoPE 与 T5 提出的 T5 bias 最相似。\n作者在理论上推导出了 NoPE 的两个性质：\nTheorem 1 (Absolute Encoding) Let $x$ be an input sequence of length $T + 1$ to the model. Then, the first layer of $f_θ$ can recover absolute positions $[1, . . . , T + 1]$ in the hidden state $H^{(1)}$. That is, there exist $W_Q, W_K , W_V , W_O, W_1$, and $W_2$ such that the self-attention and feedforward operations in the first layer compute absolute positions and write it to the next hidden state.\nTheorem 2 (Relative Encoding) Suppose that the hidden state $H^{(1)}$ contains absolute positional information, as stated in Theorem 1, and assume that it is not overwritten by any subsequent layers. Then, the self-attention in all subsequent layers can implement a relative positional encoding: there exists a parameterization of fθ such that, for $\\ell ≥ 2$, the attention dot product between query $q_n$ and key $k_m$ at positions n and m can be expressed as:\n$$ \\langle q_n, k_m\\rangle = f_{cnt}(q, k) + f_{rel}(n − m) $$where $f_{cnt}$ is a function of their content, and $f_{rel}$ is a function of their relative distance.\n【参考文献 2】探究了 softmax 中 normalization factor 对模型表现的影响，作者定义 attention 为\n$$ \\mathrm{Attn}(q,k,v) = \\mathrm{softmax}\\left(\\lambda q^Tk\\right)v $$实验结果如下图所示\n结果说明，通过调整 $\\lambda$ 我们可以有效提高 NoPE 的上下文扩展泛化能力\nConclusion NoPE 说明在 transformer 中我们可以不需要加入位置编码模块，这两篇论文均验证了 NoPE 的有效性。\nReferences The Impact of Positional Encoding on Length Generalization in Transformers Length Generalization of Causal Transformers without Position Encoding ","date":"2025-12-24T15:19:42+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-nope/","title":"Notes on NoPE"},{"content":"meta 等提出了 ALiBi, 一个通过 linear biases 来实现位置编码的方法来提高 LLM 在推理阶段的外推能力。\nIntroduction 当下，有若干种位置编码的方式：\nSinusoidal position embeddings: Transformer 提出的正弦位置编码 RoPE: RoPE 提出的旋转位置编码 T5 bias: T5 提出的相对位置编码 作者通过实验对比了不同的位置编码方法，发现这些方法在推理阶段的外推能力都比较差。\n为了解决这个问题，作者提出了 ALiBi (attention with linear biases), 一个几乎不增加计算和内存开销的位置编码方法，来提高 LLM 在推理阶段的外推能力。\nMethod 作者将外推能力定义为\na model’s ability to continue performing well as the number of input tokens during validation increases beyond the number of tokens on which the the model was trained.\n计 $L$ 为训练阶段的上下文长度， $L_{valid}$ 为推理阶段的上下文长度。\n作者首先对比了不同的位置编码方法的外推能力，结果如下图所示\n结果显示，不同位置编码在推理阶段扩展模型的上下文能力均有限。\nContext Length $L$ $L_{valid}$ Sinusoidal 512 50 1024 50 RoPE 512 200 1024 100 T5 bias 512 600 1024 800 ALiBi 512 - 1024 - 为了解决这个问题，作者提出了 AliBi, 其表达式为\n$$ \\mathrm{softmax}(q_iK^T+m\\cdot [-(i-1),\\dots,-2,-1,0]) $$其中 $m$ 是一个和 heads 相关的超参数。如果我们有 8 个 heads, 则对应的 scaling 值分别为 $[1/2^1,1/2^2,\\dots,1/2^8]$, 如果我们有 16 个 heads, 则我们对 8 个 heads 的结果进行插值，得到 $[1/2^{0.5},1/2^1,\\dots,1/2^8]$. ALiBi 的示意图如下所示\nALiBi 通过 bias 惩罚了较远的 query-key pairs, 并且不同的 heads 的惩罚项也不同，从而每个 head 对距离的信息敏感度也不尽相同。\nExperiments ALiBi 在 WikiText-103 上的实验结果如下图所示\nConclusion 作者分析了已有的 position embedding 方法，发现已有的方法在推理阶段均不能有效扩展模型的上下文长度。因此，作者提出了 AliBi, 一个通过 linear bias 来增加位置信息的方法，作者通过实验验证了 ALiBi 的有效性。\nReferences Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation ","date":"2025-12-24T15:10:55+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-alibi/","title":"Notes on ALiBi"},{"content":"google 在 2020 年发表了 T5 (Text-to-Text Transfer Transformer), 一个使用统一框架来将所有 NLP 任务转换为 text-to-text 格式的迁移学习框架。\nIntroduction 作者首先回顾了迁移学习和 pre-training, 迁移学习是提高模型在下游任务上表现的一类方法，但是目前还没有一个能够对比各种方法的框架。pre-training 通过在大量数据上进行预训练然后再进行微调，可以有效提高模型在下游任务上的表现。\n为了解决这两个问题，作者首先将所有的文本处理任务统一为 \u0026ldquo;text-to-text\u0026rdquo; 的形式，这样我们就可以对比不同架构，训练方式以及数据对模型表现的影响\n作者提到，本文并不是提供一个新的方法，而是详细对比不同方法，为后续研究提供基础。\nMethod Model 在架构上，作者使用了 Transformer 的 encoder-decoder 架构，但是作者做了几点修改\n作者提出了 T5 bias, 一个用于替换原始 transformer 绝对位置编码的相对位置编码形式 作者使用了 RMSNorm 替换了 Transformer 中的 LayerNorm. Data 作者基于 Common Crawl 构建训练数据集，作者对数据进行了清洗，最终数据集大小为 750GB. 作者将这个数据集记为 C4 (Clean Crawled Corpus).\nDownstream Tasks 下游任务包括：\ntext classification: GLUE, SuperGLUE abstractive summarization: CNN/Daily Mail question answering: SQuAD translation: WMT English to German, French and Romanian Input and Output Format 所有任务的输入输出都被转换为 text-to-text 格式。\nExperiments 作者使用的 baseline 模型是一个基于 encoder-decoder 架构的 transformer 模型，其大小以及 configuration 与 BERT base 差不多，最终模型参数量为 220M。\nfield num layers hidden size MLP hidden size num heads head dim dropout seq len value 12 768 3072 12 64 0.1 512 训练时，作者使用了 AdaFactor 优化器，batch size 为 512， 训练使用了 34B token. 学习率作者使用了 inverse square root learning schedule: $1/\\sqrt{\\max(n,k)}$, $n$ 和 $k$ 分别代表当前 step 和 warming up steps.\n作者基于 sentencepiece (见 LLM tokenizer) 构建了 Tokenizer, 覆盖 English, German, French 和 Romanian 四种语言。\n模型训练的目标函数为 BERT 使用的 \u0026ldquo;masked language modeling\u0026rdquo;, 格式如下所示\n1 2 3 4 5 6 7 8 # original text Thank you for inviting me to your party last week. # inputs Thank you \u0026lt;X\u0026gt; me to your party \u0026lt;Y\u0026gt; week. # targets \u0026lt;X\u0026gt; for inviting \u0026lt;Y\u0026gt; last \u0026lt;Z\u0026gt; 首先，作者对比了不同的架构。作者对比了如下三种 transformer 的变体：\n实验结果如下图所示\n结果显示，encoder-decoder 架构，denoising 训练目标的效果最好。并且，当 layers 减少一半之后，模型的表现大幅度下降。共享参数的 encoder-decoder 架构表现比 prefix LM 效果更好\n接下来作者针对 denoising 的配置进行了测试，实验结果发现 BERT-style 的训练目标效果最好，并且 corruption 比例对模型的表现影响有限，作者使用了 BERT 的配置，即 $15\\%$ 的 token 被 masked 掉。对于 span length, 作者通过实验发现不同的 span length 对结果影响不大。因此，作者将 span length 设置为 $3$.\n在数据上，作者发现：\n对数据进行过滤可以提高模型的表现 使用 in-domain 的数据可以提高模型在该 domain 上的表现，但是问题在于 In-domain 的数据往往比较少 数据量过少时，模型会出现 memorization，也就是过拟合的情况 Overall 作者总结前面的发现，构建了 5 个 size 的模型：\nModel layers hidden size FFN hidden size num heads head dimension Small 6 512 2048 8 64 Base 12 768 3072 12 64 Large 24 1024 4096 16 64 3B 24 1024 16384 32 128 11B 24 1024 65536 128 128 Conclusion 作者在本文中提出了 T5, 一个统一所有文本处理任务的迁移学习框架，作者系统性探究了架构，数据以及训练对模型最终表现的影响。最终作者基于 encoder-decoder transformer 架构以及 denoising training objective 训练得到了 T5 系列大语言模型。\nReferences Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer ","date":"2025-12-24T15:07:08+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-t5/","title":"Notes on T5"},{"content":"google 在 2018 年提出了 GPipe, 一个使用 pipeline parallelism 来训练大规模神经网络的并行策略\nIntroduction 大规模神经网络已经在计算机视觉和自然语言处理等任务上取得了突破性进展。但是目前训练大规模神经网络存在的问题时，我们无法在单一 GPU 上训练我们的模型。基于多 GPU 训练模型需要考虑模型的切分以及通信优化。\n为了解决这个问题，作者提出了 GPipe, 一个用于将大规模性模型分割部署到不同设备上的并行计算策略。\nMethod Notation 作者首先定义 notation 如下表所示\nnotation description $L$ number of layers $w_i$ weights of a layer $f_i$ forward function of a layer $F_k=f_j\\circ\\cdots\\circ f_i$ forward of a partition $B_k$ backward of a partition $K$ number of partitions $N$ batch size $M$ micro batch size GPipe 首先是 naive pipeline parallelism (naive PP), 我们的输入为一个 batch, 然后我们依次计算 $F_1$, 通信传输，计算 $F_2$, 计算完成之后，我们再进行反向传播，更新参数。最后继续下一个 batch 的计算。\n总体的过程如下图所示\n下面是一个按照时间轴给出的例子\nnaive PP 的问题在于，每个时刻只有一个 GPU 在工作，GPU 的利用效率很低。因此，GPipe 的做法在于将一个 batch 切分为 $M$ 个更小的 micro-batch, 下面是一个 $M=4$ 的例子\n通过切分更小的 batch，我们可以提高 GPU 的利用率\nAnalysis Bubble 接下来作者分析了 GPipe 的 bubble 情况，bubble 指的是 PP 过程中的 GPU idle time.\n对于 naive PP 来说，一个 GPU 工作时，其余 GPU 都处于空闲状态，因此其 bubble 为\n$$ T_{bubble} = (K-1)(F+B) $$总的计算时间为\n$$ T_{total} = K(F+B) $$从而 bubble rate 为\n$$ Bubble_{naive} = \\frac{T_{bubble}}{T_{total}} = \\frac{K-1}{K} $$当 $K=8$ 时，我们有 $Bubble_{naive}=87.5\\%$, 也就是说，当前训练的 GPU 空闲率为 $87.5\\%$.\n对于 GPipe 来说，由于我们将一个 batch 拆分为了更小的 batch, 我们可以提高 GPU 的利用率。\n此时，我们的 bubble time 仍然是 $T_{bubble} = (K-1)(F+B)$. 但是，现在同一时刻工作的 GPU 变多了，从上面的示意图可以看到，前向过程所需要的时间为第一个 micro batch 运行的时间加上 $M-1$ 个 batch 运行所需要的时间，反向同理，因此，GPipe 的总计算时间为\n$$ T_{total} = (M+K-1)(F+B) $$从而 GPipe 的 bubble rate 为\n$$ Bubble_{naive} = \\frac{T_{bubble}}{T_{total}} = \\frac{K-1}{M+K-1} $$当我们令 $M=8, K=8$ 时，我们有 $Bubble_{naive}=46.7\\%$, 可以看到，通过提高 micro batch 数量，我们可以显著降低 bubble rate.\nActivation Memory 对于 naive PP 来说，我们需要缓存每一层的输入，因此 activation memory 为 $\\mathcal{O}(N\\times L/K)$, 而使用 activation checkpointing 之后，我们现在的 activation memory 为\n$$ \\mathcal{O}(N + \\frac{L}{K}\\times\\frac{N}{M}) $$其中第一项代表了 boundary activation, 第二项代表了 Internal activation.\nExperiments 作者在 image classification, machine translation 任务上进行了实验。\n作者还进一步分析了影响 GPipe 性能的因素，结果如下图所示\n可以看到，activation checkpointing 是 GPipe 的主要开销来源。\nConclusion 作者在本文中提出了 GPipe, 一个针对大规模神经网络训练的并行策略。通过将模型切分部署在不同的设备上以及使用 micro batch, 我们可以显著提高硬件的利用效率以及训练稳定性。\nReferences GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism Pipeline-Parallelism: Distributed Training via Model Partitioning ","date":"2025-12-23T16:49:25+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/gpipe/","title":"GPipe"},{"content":"Introduction RoPE 已经成为了大多数 LLM 使用的 position encoding 范式，但是，RoPE 与 LLM long context 之间的关系还没有被探索清楚。在本文中，作者就探究了 base frequency 与 LLM context capability 之间的关系，并给出了一个达到指定上下文长度所需要的 base frequency 的 lower bound.\nMethod 首先作者回顾了 attention 与 RoPE 的定义， 关键就是 RoPE 这部分，如下所示\n$$ A_{ij} = (R_{i,\\theta}q_i)^T(R_{j,\\theta}k_j) = q_i^TR_{i-j,\\theta}k_j $$这里 $\\theta$ 就是 base frequency, 作者总结不同模型的 base frequency 配置如下表所示\nModel Llama-7B Llama2-7B Llama3-8B Mistral-7B Baichuan2-7B Base frequency 10,000 10,000 500,000 1,000,000 10,000 Context length 2048 4096 8192 32,768 4,096 接下来，作者回顾了 YARN. 其核心思想在于，预训练阶段所有可能的 $\\cos(t-s)\\theta_i$ 都见过，才能保证模型的 OOD 表现\n作者认为 base frequency 的设置应该满足两个条件：\nThe closer token gets more attention: 当前的 token 应该给邻近的 token 更高的注意力 The similar token gets more attention: 当前的 token 应该给相似的 token 更高的注意力 在 RoPE 中，作者已经给出了 $A_{ij}$ 与相对距离 $|i-j|$ 之间的关系。因此第一个性质已经满足了。\n接下来，作者分析了相似 token 的性质，作者定义 token 的相似性如下：\n$$ \\mathbb{E}_{q,k^*}[q^TR_{m,\\theta}k^*] - \\mathbb{E}_{q,k}[q^TR_{m,\\theta}k] $$这里 $k^*=q+\\epsilon$ 代表了相似的 token, 而 $k$ 是一个随机 token. 作者给出的结论如下\nTheorem 假设 $q,k\\in\\mathbb{R}^d$ 独立同分布，它们的标准差为 $\\sigma\\in\\mathbb{R}$, 则对于 $k^*=q+\\epsilon$, $\\epsilon$ 是一个随机变量满足 $\\mathbb{E}[\\epsilon]=0$, 则我们有\n$$ \\mathbb{E}_{q,k^*}[q^TR_{m,\\theta}k^*] - \\mathbb{E}_{q,k}[q^TR_{m,\\theta}k] = 2\\sigma^2\\sum_{i=0}^{d/2-1}\\cos(m\\theta_i) $$作者定义 $B_{m,\\theta}=\\sum_{i=0}^{d/2-1}\\cos(m\\theta_i)$, 作者认为给定 $\\theta$, 模型的上下文长度 $L_\\theta$ 满足\n$$ L_\\theta = \\sup\\{L\\mid B_{m,\\theta}\\geq 0, \\forall m\\in[L]\\} $$也就是说，base frequency 决定了 LLM 的上下文长度。作者给出了不同的上下文长度对应的 base frequency 如下表所示\nContext Len. 1k 2k 4k 8k 16k 32k 64k 128k 256k 512k 1M Lower Bound 4.3e3 1.6e4 2.7e4 8.4e4 3.1e5 6.4e5 2.1e6 7.8e6 3.6e7 6.4e7 5.1e8 总的来说，远距离衰减性保证了模型会更关注邻近的 token, 而相似 token 保证了模型能够区分出真正有意义的 token.\nExperiments 作者首先分析了 base frequency 在 fine-tuning 阶段对模型上下文能力的影响，实验结果如下图所示\n从实验结果可以看到，当 base frequency 低于阈值时，模型的表现急剧下降。\n作者进一步探讨了 base frequency 对于模型 pre-training 阶段的影响，结果也是一样的，即非常小的 base frequency 会限制模型的 context 能力，结果下图所示 （三行分别代表了 base frequency 为 1e2, 1e4 和 1e6 的情况）\n可以看到，尽管 perplexity 都差不多，但是使用更大的 base frequency 其长上下文能力明显更好。\n作者进一步分析了为什么较小的 base frequency 会影响模型的长上下文能力。作者认为较小的 base frequency 会导致 $B_{m,\\theta}$ 接近于 0， 从而模型难以区分随机 token 和相似 token, 这样模型只能依赖于邻近 token 进行学习，这样就限制了模型的长上下文能力\n作者还进一步对比了提高 base frequency 与 Interpolation 两种做法，实验结果如下表所示\n实验结果说明，Interpolation 在上下文超过 30K 之后，其 $B_{m,\\theta}\\leq0$ 的 次数显著增加，表明了其和上下文能力之间的关系。\nConclusion 作者在本文中，探究了 RoPE 中 base frequency 与 LLM 上下文能力之间的关系，发现了提高模型的上下文能力需要关注 RoPE 的 base frequency 超参数，并给出了对应的 lower bound. 作者通过实验验证了这个观点。\nReferences arxiv ","date":"2025-12-22T11:34:42+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/base-of-rope-bounds-context-length/","title":"Base of RoPE Bounds Context Length"},{"content":"DeepSeek 在 25 年 1 月提出了 Natively trainable Sparse Attention (NSA), 一个软硬件结合的稀疏注意力机制，NSA 可以在提高模型推理效率的同时提高计算效率。\nIntroduction 现有的大模型主要是基于 Transformer 提出的 softmax attention, 其主要问题在于随上下文长度增加，其 latency 也上升更快。理论估计，对于 64k 上下文长度的输出，softmax attention 部分的计算占 $70\\%\\sim80\\%$ 的 latency.\n为了解决 softmax 的 high latency 问题，，一个做法就是使用稀疏注意力机制，如 MInference 等，但是这些系数注意力机制大多没有实际部署，且它们一般只在 inference 阶段使用\n作者认为解决这个问题有两个挑战：\nHardware-aligned inference speedup: 降低 inference latency 需要算法与硬件结合，不能只关注算法层面的改进 Training-aware algorithm design: 需要在训练阶段也支持算法，从而可以降低训练的算力消耗并且保持模型的表现 为了解决这两个问题，作者就提出了 natively trainable sparse attention (NSA) 架构。NSA 通过将 key 和 value 分割为不同的 block, 然后基于三种 path: compressed coarse-grained tokens, selectively retrained fine-grained tokens 以及 sliding windows for local contextual information 来进行处理和过滤。NSA 提出了两点观点改进：\nHardware-aligned system: 优化了 blockwise sparse attention 来平衡 arithmetic intensity. Training-aware design: 支持端到端的训练和部署 Method Overview 作者首先回顾了 attention 的定义如下：\n$$ \\mathbf{o}_t=\\mathrm{Attn}(\\mathbf{q_t},\\mathbf{k}_{:,t}, \\mathbf{v}_{:,t})=\\sum_{i=1}^t \\frac{\\alpha_{t,i}}{\\sum_{t,i}\\alpha_{t,i}}\\mathbf{v_i},\\ \\alpha_{t,i} = \\exp\\left(\\frac{\\mathbf{q_t}^T\\mathbf{k}_{i}}{\\sqrt{d_k}}\\right) $$其中 $\\mathbf{q_t}\\in\\mathbb{R}^{d_k}$.\n接下来是 Arithmetic Intensity. Arithmetic intensity 指的是 FLOPs 与内存访问次数之比。由于现在的 GPU 都是计算密集型设备，理想情况下应该是 Arithmetic intensity 越高越好。\n对于 causal self-attention 来说，在训练以及 prefilling 阶段，由于 batch 较大，因此整体的 Arithmetic intensity 较高，因而这两个阶段是 computer-bound. 但是在 decoding 阶段，由于其 token-by-token generation 的性质，每次生成新的 token 时都需要重新加载 KV cache, 因而是 memory-bound.\n从而我们的优化目标也变得不一致：在训练阶段，我们希望降低计算消耗，而在推理 (decodng) 阶段，我们希望降低内存访问次数。\n基于这两个目标，作者提出了使用 $\\mathbf{k}_{:,t}, \\mathbf{v}_{:,t}$ 的子集 $\\tilde{K}_t, \\tilde{V}_t$ 来参与计算，其对应的 attention 如下所示\n$$ \\tilde{K}_t=f_K(\\mathbf{q_t},\\mathbf{k}_{:,t}, \\mathbf{v}_{:,t}), \\tilde{V}_t=f_V(\\mathbf{q_t},\\mathbf{k}_{:,t}, \\mathbf{v}_{:,t}), \\mathbf{o}_t=\\mathrm{Attn}(\\mathbf{q_t},\\tilde{K}_t, \\tilde{V}_t) $$我们还可以结合不同的方法来进行组合：\n$$ \\mathbf{o}_t^*=\\sum_{c\\in\\mathcal{C}}g_t^c\\mathrm{Attn}(\\mathbf{q_t},\\tilde{K}_t^c, \\tilde{V}_t^c) $$作者在本文中使用了三种方法 $\\mathcal{C}=\\{\\mathrm{cmp},\\mathrm{slc},\\mathrm{win}\\}$, 分别代表了 compression, selection 以及 sliding window, $g_t^c\\in[0,1]$ 代表了不同方法对应的 gating score, 类似于 MoE 的 gating layer, $g_t^c$ 由一个 MLP 和一个 sigmoid activation 生成。最终 NSA 的架构如下图所示\n作者定义 $N_t$ 代表参与计算的 KV 的总个数：\n$$ N_t = \\sum_{c\\in\\mathcal{C}} \\mathrm{size}[\\tilde{K}_t^c]. $$作者使用了一个较高的 sparsity ratio 来保证 $N_t\u003c","date":"2025-12-15T17:39:16+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-nsa/","title":"Notes on NSA"},{"content":"本 blog 详细介绍了 MoE 模型的一些关键设计与相关实验结果，为 MoE 模型的学习提供基础。\nIntroduction Motivation 现有大部分大语言模型均是基于 Transformer 架构，Kaplan scaling law 通过实验说明，大语言模型的表现与算力，数据，模型参数量息息相关。但是，对于 dense 模型来说，我们提高模型参数量时，必须同时提高所使用的算力。这就限制了大模型的 scaling law.\n而 MoE 模型的解决方法为在计算时只激活部分参数，这样，我们就可以在同等激活参数量/算力下训练更大参数量的模型，从而达到更好地表现。\n因此，MoE 模型的核心思想在于\n使用相同的激活参数量/算力，提高模型总参数量，从而达到更好的表现。\nDefinition MoE 模型和 dense 模型的示意图如下，图源 olmoe\n一个 MoE layer 包括两个模块：\nRouter：Router 负责为 token 指定合适的专家 Expert：Expert 负责处理 token 对于输入 $x\\in\\mathbb{R}^d$, 我们假设有 $N$ 个 Expert，router 一般是一个 linear layer 再加上一个 gating function (softmax 或者 sigmoid， 我们本文中使用 softmax), 其构建了 $\\mathbb{R}^d\\to\\mathbb{R}^N$ 的映射，定义为：\n$$ G(x) =[G_1(x),\\dots,G_N(x)] = \\mathrm{softmax}(W_gx + b)\\in\\mathbb{R}^N $$其中 $W_g\\in\\mathbb{R}^{N\\times d}$, $b\\in\\mathbb{R}^N$ 是可学习的参数。$G_{i}(x)$ 代表了当前 token $x$ 选择第 $i$ 个 Expert 的概率。\n一般来说，Expert 会使用和 dense 模型一样的 MLP, 我们记为\n$$ E_i(x) = \\mathrm{FFN}(x), \\quad i = 1,\\dots,N $$接下来，基于 $G(x)$ 和 $E(x)$, 我们会使用合适的方法来挑选 $K","date":"2025-12-13T16:04:04+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/moe-tutorial/","title":"MoE tutorial"},{"content":"Introduction 目前已经有了针对 dense LLM 的 scaling law, 如 Kaplan scaling law 和 Chinchilla scaling law.\n但是，对于 MoE 模型，目前还缺乏一个比较系统的 scaling law.\n为了解决这个问题，作者提出了 efficiency leverage (EL), 用于衡量 MoE 模型的效率，其定义为\n$$ EL(\\mathcal{X}_{\\mathrm{MoE}}\\mid \\mathrm{Dense})=\\frac{C_{\\mathrm{Dense}}}{C_{{\\mathrm{MoE}}}} $$其中 $C_{\\mathrm{Dense}}, C_{{\\mathrm{MoE}}}$ 分别代表了训练模型所需要的算力。EL 衡量了 moe 模型达到对应 dense 模型表现所需要的算力，EL 值越大，说明 MoE 模型效率越高。EL 的可视化如下图所示\n作者通过训练多个模型，探究了 MoE 架构与 EL 之间的关系。作者发现，MoE 模型的表现主要与专家激活比例以及算力相关。基于 scaling law, 作者训练了 Ling-mini-beta, 一个 17.5B-A0.85B 的 MoE 模型，其表现超过了 6.1B dense 模型的表现。\nPreliminary Notation description $N$ total parameters $N_a$ active parameters $E$ routed experts $E_a$ activated experts $E_s$ shared experts 作者定义 activation ratio 如下:\n$$ A = \\frac{E_a+E_s}{E+E_s} $$定义 sharing ratio 如下\n$$ S=\\frac{E_s}{E_a+E_s} $$定义 expert granularity 如下\n$$ G = \\frac{d_{\\mathrm{model}}}{d_{\\mathrm{Expert}}} $$与 DeepSeek-LLM 一样，作者使用 $C=MD$ 来表示算力，non-embedding FLOPs $M$ 和训练 token 数 $D$ 之间的关系。\nHyper Parameters 作者首先探究了针对 MoE 模型的超参数配置，最终你和出来的结果如下图\n实验结果说明，相比于 dense model, MoE model 需要更大的 batch size.\n作者基于这个 scaling law 进行了验证，结果说明这个 scaling law 比较准确。\nParameters and Dataset Size 接下来作者探究了对于模型参数量以及训练 token 个数之间的 scaling law, 求解的问题如下\n$$ (M^{opt}, D^{opt}) = \\arg\\min_{M,D}\\mathcal{L}(M,D;C,A,G,S)\\quad s.t.\\ C=MD $$实验结果如下图所示\n结果说明，不同架构对应的系数接近 $0.5$, 说明我们应该将算力均衡分配到数据和 model size 上。领一方面，MoE 模型可以通过使用更多的数据来达到更优的表现。\nEfficiency Leverage 作者将 Efficiency Leverage (EL) 定义为给定算力 $C_{target}$ 和一个 MoE 模型 $\\mathcal{X}_{MoE}$, 对应 dense 模型达到 $\\mathcal{X}_{MoE}$ 相同的表现所需要的算力 $C_{dense}$, 即\n$$ \\begin{aligned} \u0026EL(\\mathcal{X}_{\\mathrm{MoE}}\\mid \\mathrm{Dense};C_{target})=\\frac{C_{\\mathrm{Dense}}}{C_{{\\mathrm{MoE}}}}\\\\ s.t.\\ \u0026 |\\mathcal{L}(C_{MoE}, \\mathcal{X}_{\\mathrm{MoE}})-\\mathcal{L}(C_{dense}, \\mathcal{X}_{\\mathrm{dense}})|\\leq \\epsilon (\\epsilon\\to 0) \\end{aligned} $$EL 值越高说明 MoE 模型越有效。为了公平起见，dense 模型和 MoE 模型的架构参数基本相同，作者只改变 $d_{model}, d_{ffn}, d_{expert}$ 以及 $n_{layer}$.\n接下来，作者就探究了给定算力的情况下，最优的 MoE 配置，即\n$$ (A^{opt}, G^{opt}, S^{opt}) = \\arg\\min_{(A,G,S)\\in\\mathcal{X}_{\\mathrm{MoE}}}EL(\\mathcal{X}_{\\mathrm{MoE}}\\mid \\mathrm{Dense};C) $$Scaling Law 首先，坐着探究了最优的 activation ratio $A$, 即\n$$ A^{opt} = \\arg\\min_{A}\\mathcal{L}(A;C,M,G,S) $$拟合的结果如下图所示\n实验结果表明：\n模型表现随激活比例降低 er 提高 更系数的模型对于算力的提升其效率也提升更快 然后，作者探究了最优的 granularity ratio, 即\n$$ G^{opt} = \\arg\\min_{G}\\mathcal{L}(G;C,M,A,S) $$拟合的结果如下图所示\n可以看到，无限制提升 granularity 并不会提高模型的表现。并且，不同的算力对应的最优 granularity 处于一个固定的范围\n接下来，作者探究了最优的 shared expert ratio, 即\n$$ S^{opt} = \\arg\\min_{S}\\mathcal{L}(S;C,M,A,G) $$拟合的结果如下图所示\n结果说明 shared expert 的比例也不是越多越好，其存在最优值。并且给定算力的情况下，非零最小值的 shared expert 表现最好。因此作者认为，一个 shared expert 的效果最好。\n作者还探究了其他可能的因素，结论如下：\n与 DeepSeek-V3 一样，将 early layer 替换为 dense layer 可以避免 routing imbalance, 并且不会损害模型的表现 attention 应该占 $30\\%\\sim40\\%$ 左右的算力才能保证模型的表现和效率。进一步提升 attention 的算力占比虽然会提升表现但是会降低推理效率。 通过前面的发现，作者将 shared expert 设置为 1 个，然后探究 EL 与 activation ratio $A$, granularity $G$, FLOPs $C$ 之间的关系。\n首先，作者分别假设 $EL$ 与 $A$, $G$, $C$ 之间存在如下关系：\n$$ \\begin{aligned} \\log EL_{C,G}(\\hat{A}) \u0026= a_A\\log\\hat{A}, \\text{ where }\\frac{1}{\\hat{A}}=\\frac{1}{A+(1/A_{start}-1/A_{\\max})^{-1}}+\\frac{1}{A_{\\max}}\\\\ \\log EL_{C,A}(\\hat{G}) \u0026= a_G+b_G(\\log G(\\log G+c_G))\\\\ \\log EL_{A,G}(C) \u0026= a_C\\log C+c_C \\end{aligned} $$拟合的结果如下图所示\n结果显示：\n提升算力以及降低 activation ratio 都可以提高 EL granularity 对 EL 的影响在不同算力的情况下都是一致的 对于 MoE 模型，提升算力可以提高 EL 作者的结论为，activation ratio 是影响 MoE EL 的核心因素。并且随着算力的提升，MoE EL 会越来越明显。\n作者因此构建了一个统一的公式来统一三个因素\n$$ EL(A,G,C) = \\hat{A}^{\\alpha+\\gamma(\\log G)^2+\\beta \\log G} $$其中 $\\alpha=a+d\\log C$ 代表了 EL 和 activation ratio 之间的关系。拟合出来的参数如下表所示\n$\\alpha$ $d$ $\\gamma$ $\\beta$ $A_{start}$ $A_{\\max}$ 1.23 -7.61e-2 1.67e-2 -1.17e-1 1.63e-2 5.28e+16 基于这个结果，作者发现在 1e22 FLOPs 的算力下，一个 activation ratio 为 $3.1\\%$, granularity 为 $12$ 的 MoE 模型，其 EL 为 $7$.\nLing-mini-beta 基于上一节的发现，作者构建了 Ling-mini-beta, 一个 17.5B 总参数，激活参数为 0.85B 的 MoE 模型。训练使用了 1T token, 模型参数如下表所示\nModel $n_{\\text{layers}}$ $d_{\\text{model}}$ $d_{\\text{ffn}}$ $d_{\\text{expert}}$ $n_{\\text{heads}}$ $n_{\\text{kv\\_head}}$ $E$ $E_a$ $E_s$ $N$ $N_a$ Dense 6.1B 28 4096 14336 - 32 8 - - - 6.11B 6.11B Ling-mini-beta (A0.8B) 20 2048 5120 384 16 4 384 12 1 17.5B 0.85B 训练的损失变化情况如下图所示\n从图中我们可以看出，dense model 一开始的损失下降比较快，但是其最终表现不如 moe 模型。\nConclusion 在本文中，作者提出了 Efficiency Leverage, 一个衡量 MoE 模型相对于 dense 模型计算效率的 metric, 作者构建了针对 MoE 模型的 scaling law. scaling 揭示了两个主要影响 MoE 模型效率的因素：算力与激活参数比例。基于 scaling law, 作者构建了 Ling-mini-beta, 一个 17B-A0.8B 的 MoE 模型，其效率超过了对应 dense 模型的 7 倍。\nReferences arxiv ","date":"2025-12-13T15:58:51+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-ling-mini-beta/","title":"Notes on Ling-mini-beta"},{"content":"我们在本文中探讨关于 load balancing loss 的定义，性质和推广\nIntroduction 我们在 MoE tutorial 中已经介绍了 MoE 模块，MoE 尽管可以在相同的算力下扩大模型的 size, 但是其问题在于训练时容易出现负载不均衡，也就是说只有少数几个专家被激活，其他专家处于闲置状态，从而导致模型性能下降\n为了解决这个问题，一个通用的做法是使用 load balancing loss. load balancing loss 可以有效实现负载均衡，让各个专家被激活的概率差不多。\nLoad balancing loss 一共经历了如下几个阶段，发展路线如下图所示\n1 2 3 4 5 6 flowchart TD; C[Bengio et.at 2015] --\u0026gt; A[Shazeer et al., 2017]; A --\u0026gt; B[GShard]; B --\u0026gt; F[Switch Transformer]; F --\u0026gt; D[Loss-Free Balancing]; F --\u0026gt; E[Global-batch load balancing]; Method Notation 我们假设输入的 batch size 为 $B$, 总专家个数为 $N$, 激活专家个数为 $K$, $G_i(x_j)$ 代表 gating layer 预测的第 $i$ 个专家对于 token $x_j$ 的重要性程度。\nCoefficient of Variation 第一个阶段的 load balancing loss 针对广义的 MoE 模型。\nBengio et.at 2015 作者提出了一个让 batch 里每个专家平均激活概率分布更均匀的损失函数，其定义如下所示\n$$ L_b = \\sum_{i=1}^N\\left\\Vert \\frac1B\\sum_{j=1}^BG_i(x_j) - \\frac1N\\right\\Vert_2 $$其主要目标是让每个专家的平均激活概率接近 $1/N$, 即均匀分布.\nNoam et.al 2017 作者对 (Bengio et.at 2015) 进行了改进，提出了变异系数 $\\mathrm{CV}(\\cdot)$ 来保证负载均衡，其定义如下：\n$$ L_{importance} = \\mathrm{CV}(\\mathrm{Importance}(X))^2 $$其中\n$$ \\mathrm{Importance}_i(X) = \\sum_{j=1}^B G_i(x_j),\\ \\mathrm{CV}(X)= \\frac{\\mathrm{var}[X]}{\\mathbb{E}[X]},\\ i=1,\\dots,N $$注意到 $\\mathrm{var}[X]\\geq0$, 当且仅当 $X$ 为均匀分布时 $\\mathrm{var}[X]=0$, 因此 important loss 可以让每个专家在一个 batch 中的激活的平均概率尽可能一致。\n但是平均概率一致不代表各个专家处理的 token 个数一致，比如一个专家可能处理比较少的 token, 但是每个 token 的权重都很大。因此，作者额外加入了 load balancing loss\n$$ \\mathcal{L}_{\\mathrm{load}} = \\mathrm{CV}(\\mathrm{Load}(X))^2 $$其中\n$$ \\begin{aligned} \\mathrm{Load}_i(X) \u0026= \\mathrm{soft\\_estimation}\\left(\\sum_{j=1}^B\\mathbb{1}(\\text{token }j \\text{ selects expert }i)\\right)\\\\ \u0026=\\mathrm{soft\\_estimation}\\left(\\sum_{j=1}^B\\mathbb{1}\\{i\\in\\mathrm{TopK}(\\{G_i(x_j)\\}_{i=1}^{E}, K)\\}\\right),\\ i=1,\\dots,N \\end{aligned} $$这里 $\\mathrm{soft\\_estimation}$ 是作者针对离散变量进行的一个光滑化处理，以方便反向传播。\nGShard GShard 进一步对 (Noam et.al 2017) 提出的 loss 进行了简化，首先注意到 $\\sum_{i=1}^N\\mathrm{Load}_i(X)=BK$, 因此我们有\n$$ \\begin{aligned} \\mathcal{L}_{\\mathrm{load}} \u0026= \\mathrm{CV}(\\mathrm{Load}(X))^2 \\\\ \u0026= \\left(\\frac{\\mathrm{var}[\\mathrm{Load}(X)]}{\\mathbb{E}[\\mathrm{Load}(X)]}\\right)^2\\\\ \u0026= \\left(C_1[\\mathrm{Load}(X)]\\right)^2\\\\ \u0026= C_1\\mathbb{E}[\\mathrm{Load}(X)^2]-C_2 \\end{aligned} $$这里 $C_1,C_2$ 均为常数。GShard 并没有使用 soft estimation, 因此这里的 $\\mathrm{Load}_i(X)$ 定义为\n$$ \\mathrm{Load}_i(X)=\\sum_{j=1}^B\\mathbb{1}(\\text{token }j \\text{ selects expert }i),\\ i=1,\\dots,N $$其代表了一个 batch 里专家 $i$ 被激活的次数，理想情况下，所有专家被激活的次数应该是一致的，从而我们就实现了负载均衡。\n但是现在的问题是，$\\mathrm{Load}_i(X)$ 是一个离散变量，为了解决这个问题，作者使用了重要性来进行近似，即\n$$ \\begin{aligned} \\mathcal{L}_{\\mathrm{load}} \u0026= \\mathbb{E}[\\mathrm{Load}(X)^2]\\\\ \u0026\\approx \\mathbb{E}[\\mathrm{Load}(X)\\cdot\\mathrm{Importance}(X)] \\end{aligned} $$Switch Transformer Switch Transformer 进一步对 GShard 提出 load balancing loss 进行了规范化，也是现在大部分 load balancing loss 使用的形式，其定义如下：\n$$ \\mathcal{L}_{\\mathrm{load}} = \\sum_{i=1}^N f_i P_i $$其中，\n$$ f_i = \\frac{1}{B}\\mathrm{Load}_i(X)=\\frac{1}{B}\\sum_{j=1}^B\\mathbb{1}(\\text{token }j \\text{ selects expert }i),\\ i=1,\\dots,N $$代表了分配给专家 $i$ 的 token 比例，\n$$ P_i = \\frac{1}{B}\\mathrm{Importance}_i(X) = \\frac{1}{B}\\sum_{j=1}^B G_i(x_j) $$代表了这个 batch 里专家 $i$ 的平均激活概率。\n从而，这个形式和 GShard 的形式实际上是等价的：\n$$ \\mathcal{L}_{\\mathrm{load}} = \\sum_{i=1}^N f_i P_i = \\mathbb{E}[\\mathbf{f}\\cdot \\mathbf{P}] = N \\mathbb{E}[\\mathrm{Load}(X)\\cdot\\mathrm{Importance}(X)] $$Analysis 目前关于 load balancing loss 可以实现负载均衡的数学分析是比较少的。Switch Transformer 中进行了如下分析：\n[!quote] The auxiliary loss of Equation 4 encourages uniform routing since it is minimized under a uniform distribution.\n但实际上，这个分析是错误的，Switch transformer 中定义的 load balancing loss 并不是在均匀分布时取最小值，而是取最大值。考虑如下情形，\n$$ \\mathbf{f}=[1,\\dots,0],\\ \\mathbf{P}=[0,\\dots,1] $$此时，我们有\n$$ \\mathbf{f}\\cdot \\mathbf{P} = 0\\leq [1/N,\\dots,1/N]\\cdot [1/N,\\dots,1/N] = 1/N. $$我找了很多资料，最后发现苏剑林老师在 MoE环游记：2、不患寡而患不均 这篇 blog 中进行了分析。我这里基于苏剑林老师的 blog 进行了一下总结，推荐大家看苏剑林老师的原文。\n分析的核心思想就是，尽管 load balancing loss 目标函数没有意义，但是我们通过 straight-through estimator (STE)，就可以得到一个有意义的目标函数，其梯度与 load balancing loss 目标函数的梯度是一致的，这样 load balancing loss 目标函数就完成了其负载均衡的目标。\n注意到 load balancing loss 的最终目标是让每个专家处理的 token 个数尽可能一致，也就是\n$$ \\mathcal{L}_{\\mathrm{load}} = \\sum_{j=1}^B \\left(f_i-\\frac1B\\right)^2 $$但是 $f_i=\\mathrm{Load}_i(X)$ 是一个离散变量不可导，因此，基于 STE, 我们可以将其改写为\n$$ \\mathcal{L}_{\\mathrm{load}} = \\sum_{j=1}^B \\left(P_i+\\mathrm{sg}[f_i-P_i]-\\frac1B\\right)^2 $$其中 $\\mathrm{sg}[\\cdot]$ 满足：\n$$ \\mathrm{sg}[\\cdot]=\\begin{cases} x, \u0026\\text{forward pass}\\\\ 0, \u0026\\text{backward pass} \\end{cases} $$此时我们目标函数的梯度就变成了\n$$ \\begin{aligned} \\nabla_\\theta\\mathcal{L}_{\\mathrm{load}} \u0026= \\sum_{j=1}^B 2\\nabla_\\theta\\left(P_i+\\mathrm{sg}[f_i-P_i]-\\frac1B\\right)\\left(P_i+\\mathrm{sg}[f_i-P_i]-\\frac1B\\right)\\\\ \u0026= 2\\sum_{j=1}^B\\nabla_\\theta P_i\\left(f_i-\\frac1B\\right)\\\\ \u0026= 2\\sum_{j=1}^B\\nabla_\\theta P_i\\left(f_i-\\frac1B\\right)\\\\ \u0026= 2\\nabla_\\theta \\left(\\sum_{j=1}^Bf_iP_i\\right) \\end{aligned} $$这样就对应上了我们之前的目标函数。总之，目标函数尽管本身没有意义，但是其对应的梯度却可以让模型实现负载均衡的目标。\nExtension Loss Free Load Balancing Strategy Loss-Free load balancing 是 DeepSeek 提出来的一个无需 load balancing loss 实现负载均衡的方法，其核心思想是 load balancing loss 会影响语言建模性能，因此在 $\\mathrm{TopK}$ 操作时，作者加入一些扰动，从而让负载高的专家被选择的概率降低，让负载低的专家被选择的概率提高。\n具体细节见 Loss-Free Balancing.\nGlobal-batch Load Balancing Loss Strategy Global-batch Load Balancing 是 Qwen 提出来的一个提高负载均衡的方法。其核心思想是，现有的 token choice 只是在 batch 层面达到负载均衡，这种局部均衡的性质可能会对模型的全局均衡性产生影响。因此，作者就预先在 global batch 层面对专家进行分配，从而提高专家在不同 domain 上的特化程度。\n具体细节见 Global-batch load balancing\nExperiments olmoe 对比了加入 load balancing loss 之后模型的表现变化情况，结果如下图所示\n可以看到，加入 load balancing loss 之后，模型的表现均超过了不加时的表现。\n作者进一步分析了不同专家在加/不加 load balancing loss 时的激活情况，结果如下图所示\n结果显示，load balancing loss 确实可以让不同专家的激活概率大致相当。\nConclusion 在本文中，我们介绍了 load balancing loss 的演化，定义，分析以及相关的实验结果。\nReference GShard Bengio et.at 2015 Noam et.al 2017 Loss free balancing Global-batch load balancing loss MoE环游记：2、不患寡而患不均 ","date":"2025-12-11T16:10:08+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/load-balancing-tutorial/","title":"Load Balancing tutorial"},{"content":"Qwen 在 25 年 2 月提出了 global batching load balancing loss strategy, 其在 global level 上考虑每个专家的负载均衡，从而提高模型的表现\nIntroduction 实现 MoE 模型负载均衡的原因有两点：\nEffectiveness: 通过负载均衡才能更高效地利用各个专家 Efficiency: 一般需要使用 expert parallel 来部署 MoE 模型，不均衡的负载会大幅度降低前向过程 已有的框架如 DeepSpeed, Megablocks 和 Megatron-Core 都是在 micro-batch level 上计算负载均衡损失的，但是，一个 micro-batch 通常只包含少数序列，因此 load balancing loss 就要求各个专家在每个序列上均匀分布。\n针对这个问题，作者在本文中提出的解决方法是在 global-batch 层面考虑负载均衡，\nMethod MoE 模型定义如下\n$$ y = \\sum_{i\\in N_E,g_i\\in\\mathrm{TopK}(g)}g_i(x)E_i(x) $$其中 $E_i$ 是对应的专家, $g_i$ 是对应的权重\nLoad balancing loss 定义如下\n$$ \\mathrm{LBL} = N_E\\sum_{i=1}^{N_E}f_iP_i $$一般来说，MoE 模型训练时会使用 expert parallel 策略，此时 load balancing loss 修改为\n$$ \\mathrm{LBL}_{\\mathrm{micro}}= \\frac{1}{N_P}\\sum_{j=1}^{N_P}\\left(N_E\\sum_{i=1}^{N_E}f_i^hP_i^j\\right) $$其中 $N_P$ 是 parallel groups 的个数。在这种情况下，模型需要再每个 parallel group 中实现负载均衡。但是某一个 mircro-batch 里可能只包含某一个 domain 的序列，因此各个专家被强制要求在每一个 domain 中也均衡分布。作者认为，这种方式会限制专家的能力，进而影响模型的表现。\n作者的解决方法在于，获取每个 parallel group 的 $f_i$ 然后求 global-batch 的 $\\bar{f}_i$.\n$$ \\mathrm{LBL}_{\\mathrm{global}}=N_E\\sum_{i=1}^{N_E}\\bar{f}_i\\bar{P}_i=N_E\\sum_{i=1}^{N_E}\\bar{f}_i\\left(\\frac{1}{N_P}\\sum_{j=1}^{N_P}P_j\\right) = \\frac{1}{N_P}\\sum_{j=1}^{N_P}\\left(N_E\\sum_{i=1}^{N_E}\\bar{f}_iP_i^j\\right) $$实际中，由于计算节点数量有限，micro-batch size 之和可能会小于 global-batch size, 因此我们会使用 gradient accumulation. 在这种情况下， 作者使用了一个 buffer 来保存多个 micro-batch 的专家选择次数，最终算法实现过程如下所示\nExperiments 作者使用了不同大小的模型来进行实验，作者用 Balance BSZ 来表现计算 expert selection frequency 时的 token 数，结果如下图所示\n可以看到，global LBL 可以提高模型的表现，并且，当 Balance BSZ 增加时，模型的表现也会提升。作者发现对于 Loss-Free Balancing 来说，使用 global batch 的效果也更好。\n作者还发现，global LBL 会提高专家的特化程度，结果如下图所示\n可以看到，使用默认的 LBL loss, 各个专家的特化程度很低，而使用本文的 global LBL 之后，专家的特化程度有了大幅度的提高。\n作者还进一步发现，随着 Balance BSZ 的提高，模型表现也持续提升，结果如下图所示\n作者认为，synchronization 和 buffer 机制相比于 micro-batch 来说可以带来大幅度提升。\nAnalysis 首先，为了分析 global batch LBL 优于 Micro batch 的关键原因，作者先同步所有 group 的矩阵 $G$, 然后再从全局 token 中随机选一批，用这批 token 计算专家选择频率。通过这个方法，我们可以保证 token 数量与 micro-batch 一致，但是分布于 Global-batch 一致。实验结果发现，这种方法的表现优于 micro-batch LBL, 说明 global batch LBL 的优势在于 Token 分布更全局，而不是 token 数量更多。\n作者还分析了 global batch LBL 和 micro batch LBL 两种模式，作者认为前者是后者的一个宽松版约束，作者发现从后者切换为前者之后，模型的表现可以得到进一步提升，但是其表现仍然不如一开始就使用 global batch LBL 更好。作者分析原因认为，这是因为 expert 收敛速度比较快。\n作者进一步通过降低 micro batch LBL 权重来探究是否可以达到同样的表现，结果发现适度江都权重确实可以提高模型的表现，但是降低太多会损害模型的表现，即此时出现了负载不均衡的现象\n作者对比了 global batch LBL 和 micro batch LBL 效率发现，前者比后者慢 $2\\%$ 左右，这个差距几乎可以忽略不计\n作者进一步分析了不同 balancing 方式的专家特化程度，结果如下图所示\n实验结果发现：\n使用 global batch LBL 之后，topK sum 明显变大，这说明 routing 和 language modeling 任务对齐地更好 global batch LBL 的专家在不同任务上的特化程度更明显 micro batch LBL 的 topK sum 比较小 Loss-free balancing 的 topK sum 介于 micro batch LBL 和 global batch LBL 之间 Conclusion 作者在本文中提出了 global LBL 来在全局为负载均衡提供指导，结果发现通过在更大的范围进行负载均衡的计算，我们可以有效提高专家的特化程度以及提高模型在下游任务上的表现\nReferences Arxiv ","date":"2025-12-11T16:09:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-global-batch-load-balancing/","title":"Notes on Global-batch load balancing"},{"content":"Introduction 传统的偏好优化主要基于 RLHF, 其过程为：SFT, reward modeling, RLHF. 其中 reward model 的训练至关重要，对最终模型的表现有非常大的影响。但是 RLHF 的问题在于其训练复杂且经常不稳定。\n为了解决这个问题，作者提出了 Direct Preference Optimization (DPO), DPO 通过构建 reward function 和最优策略之间的关系，进而通过训练 policy model 来同时完成 reward model 的训练。这样，我们就避免了 reward model 的训练。结果发现，DPO 的表现超过了之前的偏好优化方法。\nBackground 作者首先回顾了 RLHF, RLHF 的 pipeline 如下\n其包含了三个步骤：\nSFT RLHF 首先基于 base model 通过 SFT 得到一个初始模型 $\\pi^{\\mathrm{SFT}}$.\nReward modeling 接下来，我们给定输入 $x$, 对 $\\pi^{\\mathrm{SFT}}$ 采样得到 $(y_1,y_2)\\sim \\pi^{\\mathrm{SFT}}(y\\mid x)$. 输出 $y_1, y_2$ 然后由人类进行打分得到偏好关系 $y_wr^*(x, y_l)$. reward modeling 通常基于 Bradley-Terry (BT) 模型得到，BT model 定义人类真实偏好分布 $p^*$ 如下：\n$$ p^*(y_1\u003ey_2\\mid x)= \\frac{\\exp(r^*(x, y_1))}{\\exp(r^*(x, y_1))+\\exp(r^*(x, y_2))} $$假设我们从分布 $p^*$ 中采集到一个数据集 $\\mathcal{D}=\\{(x^{(i)},y_w^{(i)},y_l^{(i)})\\}_{i=1}^N$, 我们可以通过 MLE 来估计得到一个 reward model $r_{\\phi}(x,y)$, 通过将这个问题转换为一个二分类问题，我们得到对应的 negative log-likelihood loss 如下：\n$$ \\mathcal{L}_R(r_\\phi, D) = -\\mathbb{E}_{(x,y_w,y_l)\\sim \\mathcal{D}}\\left[\\log \\sigma\\left(r_\\phi(x,y_w)-r_\\phi(x,y_l)\\right)\\right] $$其中 $\\sigma$ 是 logistic function. 在 LLM 中，$r_\\phi(x,y)$ 通常由 $\\pi^{\\mathrm{SFT}}$ 初始化得到，然后我们在 $\\pi^{\\mathrm{SFT}}$ 最后一层加入一个 linear layer 来得到对应的 reward 的预测值。一般地，为了降低 reward function 的 variance, 之前的工作会进行 normalization, 即 $\\mathbb{E}_{(x,y)}\\sim\\mathcal{D}[r_{\\phi}(x,y)]=0$ for all $x$.\nRL fine-tuning 这个阶段，我们基于学习到的 reward function $r_\\phi(x,y)$ 来为 LLM 的训练提供奖励，作者使用了和 RLHF 一样的目标函数：\n$$ \\max_{\\pi_\\theta}\\quad \\mathbb{E}_{x\\sim \\mathcal{D}, y\\sim \\pi_\\theta(y\\mid x)}\\left[r_\\phi(x,y)\\right] - \\beta\\mathbb{D}_{\\mathrm{KL}}\\left[\\pi_\\theta(y\\mid x)\\Vert \\pi_{\\mathrm{ref}}(y\\mid x)\\right] $$这里 $\\beta\u003e0$ 是超参数，用于控制 $\\pi_\\theta$ 相对于 $\\pi_{\\mathrm{ref}}$ 的偏离程度， $\\pi_{\\mathrm{ref}}$ 一般就是 $\\pi^{\\mathrm{SFT}}$. 实际上， $\\pi_\\theta$ 也由 $\\pi^{\\mathrm{SFT}}$ 初始化.\nDPO DPO 的主要目标是构建一个更简单的 policy optimization 方法。与 RLHF 不同，DPO 跳过了 reward modeling 这一阶段，而是直接使用偏好数据来优化大语言模型\n为了实现这个目标，作者第一步就是构建 reward model 和 policy model 之间的关系。注意到\n$$ \\begin{aligned} \u0026\\max_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}, y\\sim \\pi_\\theta(y\\mid x)}\\left[r_\\phi(x,y)\\right] - \\beta\\mathbb{D}_{\\mathrm{KL}}\\left[\\pi_\\theta(y\\mid x)\\Vert \\pi_{\\mathrm{ref}}(y\\mid x)\\right]\\\\ \u0026= \\max_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\mathbb{E}_{y\\sim \\pi_\\theta(y\\mid x)}\\left[r_\\phi(x,y) - \\beta\\log\\frac{\\pi_\\theta(y\\mid x)}{\\pi_{\\mathrm{ref}}(y\\mid x)}\\right]\\\\ \u0026= \\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\mathbb{E}_{y\\sim \\pi_\\theta(y\\mid x)}\\left[\\log\\frac{\\pi_\\theta(y\\mid x)}{\\pi_{\\mathrm{ref}}(y\\mid x)}-\\frac{1}{\\beta}r_\\phi(x,y)\\right]\\\\ \u0026= \\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\mathbb{E}_{y\\sim \\pi_\\theta(y\\mid x)}\\left[\\log\\frac{\\pi_\\theta(y\\mid x)}{\\frac{1}{Z(x)}\\pi_{\\mathrm{ref}}(y\\mid x)\\exp\\left(\\frac1\\beta r_\\phi(x,y)\\right)}-\\log Z(x)\\right]\\\\ \\end{aligned} $$其中 $Z(x)$ 是 partition function, 定义如下\n$$ Z(x) = \\sum_{y} \\pi_{\\mathrm{ref}}(y\\mid x)\\exp\\left(\\frac1\\beta r_\\phi(x,y)\\right) $$注意到 partition function 只是 $x$ 和 $\\pi_{\\mathrm{ref}}$ 的函数，而不依赖于 $\\pi_\\theta$, 因此我们定义\n$$ \\pi^*(y\\mid x) = \\frac{1}{Z(x)}\\pi_{\\mathrm{ref}}(y\\mid x)\\exp\\left(\\frac1\\beta r_\\phi(x,y)\\right) $$易知 $\\pi^*$ 满足 $\\pi^*(y\\mid x)\\geq0, \\forall y$, 以及 $\\sum_y \\pi^*(y\\mid x)=1$. 因为 $Z(x)$ 不依赖于 $y$, 因此我们可以进一步简化上面的目标函数如下\n$$ \\begin{aligned} \u0026\\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\left[\\mathbb{E}_{y\\sim \\pi_\\theta(y\\mid x)}\\left[\\log\\frac{\\pi_\\theta(y\\mid x)}{\\pi^*(y\\mid x)}\\right]-\\log Z(x)\\right] \\\\ \u0026= \\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\left[\\mathbb{D}_{\\mathrm{KL}}\\left[\\pi_\\theta(y\\mid x)\\Vert \\pi^*(y\\mid x)\\right]-\\log Z(x)\\right]\\\\ \u0026= \\boxed{\\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\left[\\mathbb{D}_{\\mathrm{KL}}\\left[\\pi_\\theta(y\\mid x)\\Vert \\pi^*(y\\mid x)\\right]\\right]} \\end{aligned} $$而上述目标函数的最小值为 0，当且仅当 $\\pi_\\theta=\\pi^*$, 此时我们的最优 policy 为\n$$ \\pi_\\theta^*(y\\mid x) = \\pi^*(y\\mid x) = \\frac{1}{Z(x)}\\pi_{\\mathrm{ref}}(y\\mid x)\\exp\\left(\\frac1\\beta r_\\phi(x,y)\\right),\\ \\forall x\\in\\mathcal{D} $$直接求解 $\\pi_\\theta^*$ 非常困难，因为这涉及到 $Z(x)$ 的计算，这个时候作者就提出了一个关键改变，即我们从上述的 $\\pi_\\theta^*$ 反向推到出 $r(x,y)$:\n$$ r(x,y) = \\beta\\log\\frac{\\pi_\\theta(y\\mid x)}{\\pi_{\\mathrm{ref}}(y\\mid x)}+\\beta\\log Z(x) $$我们可以基于这个公式来推到出最优的 reward function $r^*$ 以及对应的最优策略 $\\pi^*$.\n此时，我们的表达式里仍然含有 $Z(x)$, 但是当我们使用 Bradley-Terry 模型之后，我们就可以得到真实的人类偏好分布，计算过程如下所示\n$$ \\begin{aligned} p^*(y_w\u003ey_l\\mid x)\u0026= \\sigma\\left(r_\\phi(x,y_w)-r_\\phi(x,y_l)\\right)\\\\ \u0026= \\boxed{\\sigma\\left(\\beta\\log\\frac{\\pi_\\theta(y_w\\mid x)}{\\pi_{\\mathrm{ref}}(y_w\\mid x)}-\\beta\\log\\frac{\\pi_\\theta(y_l\\mid x)}{\\pi_{\\mathrm{ref}}(y_l\\mid x)}\\right)} \\end{aligned} $$这样，基于 MLE 的目标函数就是\n$$ \\mathcal{L}_{\\mathrm{DPO}}(\\pi_{\\theta};\\pi_{\\mathrm{ref}})= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\sigma\\left(\\beta\\log\\frac{\\pi_\\theta(y_w\\mid x)}{\\pi_{\\mathrm{ref}}(y_w\\mid x)}-\\beta\\log\\frac{\\pi_\\theta(y_l\\mid x)}{\\pi_{\\mathrm{ref}}(y_l\\mid x)}\\right)\\right] $$通过这种方式，我们就避免了 reward model 的训练。\n接下来，作者分析了一下 DPO 目标函数的梯度，\n$$ \\begin{aligned} \\nabla_\\theta \\mathcal{L}_{\\mathrm{DPO}}(\\pi_{\\theta};\\pi_{\\mathrm{ref}})\u0026=-\\nabla_\\theta \\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\log\\sigma(u)\\right]\\\\ \u0026= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\frac{\\nabla_\\theta \\sigma(u)}{\\sigma(u)}\\nabla_\\theta u\\right]\\\\ \u0026= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\frac{\\sigma(u)(1-\\sigma(u))}{\\sigma(u)}\\nabla_\\theta u\\right]\\\\ \u0026= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[(1-\\sigma(u))\\nabla_\\theta u\\right]\\\\ \u0026= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\sigma(-u)\\nabla_\\theta u\\right]\\\\ \u0026= -\\mathbb{E}_{(x,y_w,y_l)\\in\\mathcal{D}}\\left[\\beta\\sigma\\left[\\beta\\log\\frac{\\pi_\\theta(y_l\\mid x)}{\\pi_{\\mathrm{ref}}(y_l\\mid x)}-\\beta\\log\\frac{\\pi_\\theta(y_w\\mid x)}{\\pi_{\\mathrm{ref}}(y_w\\mid x)}\\right]\\left[\\nabla_\\theta\\log\\pi(y_w\\mid x) - \\nabla_\\theta \\log\\pi(y_l\\mid x)\\right]\\right]\\\\ \\end{aligned} $$从梯度来看，当我们的 reward 估计错误时，即 $\\sigma(-u)\u003e0$ 时， DPO 会提高 $y_w$ 的生成可能性以及降低 $y_l$ 的生成可能性，从而提高模型对于偏好输出的可能性。\n最终，DPO 的 pipeline 如下：\n收集偏好数据 $y_1,y_2\\sim\\pi_{\\mathrm{ref}}(\\cdot\\mid x)$, 然后通过人类标注得到偏好数据集 $\\mathcal{D}=\\{(x^{(i)},y_w^{(i)},y_l^{(i)})\\}_{i=1}^N$ 基于 $\\mathcal{L}_{\\mathrm{DPO}}(\\pi_{\\theta};\\pi_{\\mathrm{ref}})$ 优化大语言模型参数 一般来说，我们将 $\\pi_{\\mathrm{ref}}$ 初始化为 $\\pi^{\\mathrm{SFT}}$, 但是如果我们没有 $\\pi^{\\mathrm{SFT}}$ 时，我们可以通过最大似然估计来得到 $\\pi_{\\mathrm{ref}}$, 即\n$$ \\pi_{\\mathrm{ref}} = \\arg\\max_{\\pi}\\mathbb{E}_{(x,y_w,y_l)\\sim\\mathcal{D}}\\log \\pi(y_w\\mid x) $$Theoretical Analysis of DPO Definition 两个 reward function $r(x,y)$ 和 $r'(x,y)$ 等价当且仅当存在函数 $f$ 满足 $r(x,y)-r'(x,y)=f(x)$。\n上述定义给出了一个等价关系，将 reward function 分割成了不同的等价类。接下来，作者给出了两个引理：\n第一个引理说明同一个等价类里的 reward function 对应的偏好分布一致\nLemma 1 在 Plack-Luce 框架下，如 Bradley-Terry model, 同一个等价类里的 reward function 得到的偏好分布是一致的\n第二个引理说明了最优策略对应的 reward function 都在一个等价类里\nLemma 2 在有限制的情况下，同一个等价类里的 reward function 得到的最优策略是一样的\nLemma 2 说明我们只要从最优的等价类里任意找到一个 reward function, 则最终的效果是一样的。\nTheorem 1 假设我们有一个 reference model $\\pi_{\\mathrm{ref}}(\\cdot\\mid x)\u003e0$ for all prompt-answer pairs $(x,y)$, 则对于某个模型 $\\pi(y\\mid x)$, 则与 $\\pi_{\\mathrm{ref}}$ 对应的等价类里的 reward function 都可以表示为如下形式 $r(x,y) = \\beta\\frac{\\pi(y\\mid x)}{\\pi_{\\mathrm{ref}}(y\\mid x)},\\ \\beta\u003e0$.\n我们也可以通过 Theorem 1 来显示推导出 DPO 选择的 reward function:\n$$ \\sum_{y}\\underbrace{\\pi_{\\mathrm{ref}}(y\\mid x)\\exp\\left(\\frac1\\beta r_\\phi(x,y)\\right)}_{\\pi(y\\mid x)}=1 $$回顾前面 $\\pi^*(y\\mid x)$ 的定义，我们知道 $\\pi(y\\mid x)$ 实际上是针对 $r(x,y)$ 推导出来的最优策略的 partition function.\n注意到我们的初始目标函数可以改写为如下形式\n$$ \\min_{\\pi_\\theta}\\ \\mathbb{E}_{x\\sim \\mathcal{D}}\\left[\\underbrace{r_{\\phi}(x,y)-\\beta\\log Z(x)}_{f(r_\\phi, \\pi_{\\mathrm{ref}},\\beta)}-\\beta\\underbrace{ \\frac{\\pi_\\theta(y\\mid x)}{\\pi_{\\mathrm{ref}}(y\\mid x)}}_{\\text{KL}}\\right] $$此时，我们可以将 $f(r_\\phi, \\pi_{\\mathrm{ref}},\\beta)$ 里的 normalization term 视作为 $\\pi_{\\mathrm{ref}}$ 的 soft value function, 这个 soft value function 不影响最终的结果，但是没有的话会导致训练时的 variance 很高。 DPO 通过 re-parameterization, 得到的奖励函数不需要 baseline, 因而解决了训练不稳定的问题。\nExperiments 作者首先对比了不同偏好优化算法的表现，如下图所示\n从实验结果可以看到，DPO 对于不同的 KL Values 和不同的采样温度其表现都非常好，并且更加 robust\nConclusion 在本文中，作者提出了 DPO，一个针对 LLM 偏好优化的训练范式，DPO 构建了最优的 policy 与 reward function 之间的关系，从而避免了训练 reward model, 让模型可以直接从偏好数据集中进行学习和训练。\nReferences openreview ","date":"2025-12-09T10:43:11+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-dpo/","title":"Notes on DPO"},{"content":"DeepSeek 在 24 年 11 月发布了 DeepSeek-V3, 一个仅花费 2.8M H800 hours 的大语言模型，且在各个 benchmark 上达到了 SOTA 表现\nIntroduction 作者在本文中提出了 DeepSeek-V3, 一个 671B-A37B 的 MoE 大语言模型。\n在训练目标和架构上，作者做了如下改进：\nefficiency inference: 采用了 DeepSeek-V2 提出的 MLA cost-effective training: 采用了 DeepSeekMoE 提出了 MoE 架构 auxiliary-loss-free strategy: 采用了 Loss-Free Balancing 提出的 loss balancing 策略 multi-token prediction: 采用了 MTP 的训练目标来提升模型的表现 在训练上，作者做了如下改进：\n使用了 FP8 混合精度进行训练并验证了其在大规模模型上的有效性 作者构建了 DualPipe 算法用于高效的 pipeline parallelism 构建了 cross-node all-to-all communication kernel 来高效使用 InfiniBand 以及 NVLink bandwidth 优化了 memory footprint, 来避免使用 tensor parallelism 预训练阶段，DeepSeek-V3 使用了14.8T token. 在 mid-training 阶段，作者将模型的上下文长度由 8K 扩展到 32K, 再扩展到 128K.\n后训练阶段，作者使用了 SFT 和 RL 两个阶段来提高模型的表现，作者还对 DeepSeek-R1 进行蒸馏来提高模型的 reasoning 能力\nArchitecture Basic Architecture DeepSeek-V3 的架构与 DeepSeek-V2 的架构一致，如下图所示\nMLA 的介绍见 MLA, MoE 架构的介绍见 DeepSeekMoE. DeepSeek-V3 在 DeepSeekMoE 的基础上做了两点改变：\n受 Loss-Free Balancing 启发，作者使用了 sigmoid fu nction 来计算 affinity score 对于 selected affinity score 应用了 normalization 在 DeepSeekMoE 的基础上，作者使用了 Loss-Free Balancing. 其表达式如下\n$$ g_{i,t} = \\begin{cases} s_{i,t}, \u0026 s_{i,t}+b_i\\in\\mathrm{Topk}(\\{s_{i,j}+b_j\\mid 1\\leq j\\leq N\\}, K)\\\\ 0, \u0026\\text{otherwise} \\end{cases} $$其中 $b_i$ 仅影响 routing, 训练时，如果对应的 expert 负载不均衡，则对 $b_i$ 进行更新，更新方式为增加/减少 $\\gamma$, 这里 $\\gamma$ 是一个超参数\n为了提高 routing 在 sequence 层面的负载均衡，作者还是用了一个 complementary sequence-wise balance loss:\n$$ \\begin{aligned} \\mathcal{L}_{\\mathrm{Bal}} \u0026= \\alpha\\sum_{i=1}^{N_r} f_iP_i\\\\ f_i \u0026= \\frac{}{}\\sum_{t=1}^T\\mathbb{1}(s_{i,t}\\in\\mathrm{TopK}(\\{s_{j,t}\\mid 1\\leq j \\leq N_r\\}, K_r))\\\\ s_{i,y}' \u0026= \\frac{s_{i,t}}{\\sum_{j=1}^{N_r}s_{j,t}}\\\\ P_i \u0026= \\frac1T\\sum_{t=1}^T s_{i,t}' \\end{aligned} $$其中 $\\alpha$ 是 balance factor, 为超参数，$T$ 是 sequence length, 加入这个损失后，每个 sequence 上的负载会变得更加均衡\n与 DeepSeekMoE 一样，作者也在 node 层面实现负载均衡，具体做法就是，根据 node 中所包含 expert 的 affinity score 之和来选取最高的 $M$ 个 nodes, 这样就可以进一步提高 computation 和 communication 之间的 overlap.\n由于 DeepSeek-V3 负载均衡比较好，因此作者使用了 no token-dropping strategy.\nMTP 受 MTP 启发，DeepSeek-V3 也构建了 MTP 模块，作者认为 MTP 模块有两个优势：\nMTP objective 提供了更多的学习信号，进而提高了数据使用效率 MTP 可以让模型更好预测未来的 token 与 MTP 不同，DeepSeek-V3 【todo】, DeepSeek-V3 所使用的 MTP 模块架构图如下所示\nMTP 模块使用了 $D$ sequential modules 来预测未来的 $D$ 个 token. 其中，第 $k$ 个 MTP 模块包含一个共享的 embedding layer $\\mathrm{Emb}(\\cdot)$ , 一个共享的 output head $\\mathrm{OutHead}(\\cdot)$, 一个 transformer block $\\mathrm{TRM}_k(\\cdot)$ 和一个 projection matrix $M_k\\in\\mathbb{R}^{d\\times 2d}$.\n对于第 $i$ 个 token 以及第 $k$ 个 MTP 模块，作者首先将第 $k-1$ 个 MTP 模块的第 $i$ 个 token 的 hidden states $h_i^{k-1}\\in\\mathbb{R}^d$ 和第 $i+k$ 个 token 的 embedding $\\mathrm{Emb}(t_{i+1})\\in\\mathbb{R}^d$ 联合在一起\n$$ h_i'^{k}=M_k[\\mathrm{RMSNorm(h_{i}^{k-1});\\mathrm{RMSNorm(\\mathrm{Emb}(t_{i+k}))}}] $$其中 $[\\cdot;\\cdot]$ 代表 concatenation 操作。当 $k=1$ 时，$h_{i}^{k-1}$ 就代表了 main model 的输出。每个 MTP 模块的 embedding layer 和 main model 的 embedding layer 是共享的，接下来， $h_i'^k$ 作为第 $k$ 个 MTP 模块 transformer block 的输入，得到\n$$ h_{i}^k = \\mathrm{TRM}_k(h_{i}'^l) $$最后，共享的 output head 输出对应的概率分布：\n$$ P_{i+k+1}^k = \\mathrm{OutHead}(h_i^k) $$这里的 $\\mathrm{OutHead}(\\cdot)$ 的权重与 main model 也是共享的。作者这里提到，所使用的思想与 EAGLE 是类似的，但是不同的地方在于，EAGLE 主要是用于 speculative decoding, 而 MTP 主要用于提升训练。\nMTP 的训练目标为未来 $T-k-1$ 个 token 的 cross-entropy loss:\n$$ \\mathcal{L}_{\\mathrm{MTP}}^k = \\mathrm{CrossEntropy}(P_{2+k:T+1}^k,t_{2+k:T+1}) = -\\frac1T\\sum_{t=k+2}^{T+1}\\log P_i^k[t_i], $$其中 $T$ 是 sequence length, $t_i$ 是第 $i$ 个位置对应的 ground truth token, $P_i^k[t_i]$ 代表低 $k$ 个 MTP 模块给出的 $t_i$ 的预测概率。最后，作者对所有的 MTP loss 进行求和，得到\n$$ \\mathcal{L}_{\\mathrm{MTP}} = \\frac{\\lambda}{D}\\sum_{k=1}^D \\mathcal{L}_{\\mathrm{MTP}}^k. $$Infra DeepSeek-V3 训练使用了 2048 张 H800, 每个 node 包含 8 张 GPU, node 内部使用 NVLink 和 NVSwitch 进行连接，node 之间使用 InfiniBand 进行连接\n与之前的 DeepSeek 系列相同，DeepSeek-V3 也是用了 HAI-LLM 框架来支持训练。训练时，DeepSeek-V3 使用了 16-way PP, 64-way EP (spanning 8 nodes, GShard) 以及 ZeRO-1 DP.\n作者主要进行了三点优化：\n构建了 DualPipe 用于高效 pipeline parallelism 构建了 cross-node all-to-all communication kernels 来高效利用 IB 以及 NVLink bandwidth 优化了训练时的 memory footprint, 使得训练时不再依赖 TP Training Framework DualPipe DeepSeek-V3 中，由于 cross-node EP, computation-to-communication ratio 近似为 1:1, 为了解决这个问题，作者提出了 DualPipe. DualPipe 的核心思想是将 forward 和 backward 过程中的 computation 以及 communication 进行重叠。与 ZeroBubble 类似，作者将每个 chunk 分为四个部分：attention, all-to-all dispatch, MLP 以及 all-to-all combine. 对于 attention 和 MLP, 作者还进一步将 backward 拆分为 针对权重和输入的 backward. 其示意图如下所示，这里橙色部分代表 forward, 绿色代表了针对输入的 backward, 蓝色代表了针对权重的 backward\n示意图里包含两个 block, 我们记为 block1, block2, 前向计算过程为\n1 dispatch(F, block1) -\u0026gt; MLP(F, block1) -\u0026gt; combine(F, block1) -\u0026gt; attention(F, block2) 其中，dispatch(F, block1) 与 MLP 的反向传播 MLP(B, block1) 计算重叠, MLP(F, block1) 与 MLP 反向传播的 dispatch dispatch(B, block1) 通信重叠，combine(F, block1) 与 attention 反向传播的 attention(B, block2) 重叠，attention(F, block2) 与反向传播的 combine combine(block2) 重叠。下面是一个具体的例子\n作者进一步对比了 DualPipe, 1F1B 和 ZeroBubble, 结果如下表所示\nMethod Bubble Parameter Activation 1F1B $(PP-1)(F+B)$ $1\\times$ $PP$ ZB1P $(PP-1)(F+B-2W)$ $1\\times$ $PP$ DualPipe (Ours) $(\\frac{PP}{2}-1)(F\\\u0026B+B-3W)$ $2\\times$ $PP+1$ 这里 $F$ 是 forward chunk 的执行时间，$B$ 是 backward chunk 执行的时间，$W$ 是一个 chunk \u0026ldquo;backward for weights\u0026rdquo; 的执行时间，$F\\\u0026B$ 是一个 chunk 前向反向传播重叠的时间。可以看到，DualPipe 只使用了额外的 $1/PP$ 倍的 peak activation memory, 就大幅度降低了 bubble 时间\nCross-node All-to-all Communication 作者针对 DualPipe 构建了 cross-node all-to-all communication kernels 来提高通信效率。\n作者提到，跨节点通信使用的是 IB, 节点内部通信使用的是 NVLink， 通信方式如下图所示。\n对于 H800 来说,NVLink 的带宽为 160GB/s, IB 的带宽为 50GB/s, 因此 NVLInk 的通信效率是 IB 的 3.2 倍，即节点内部通信效率高于节点之间的通信效率。为了提高通信效率，作者限制每个 token 只能被分发到至多 4 个节点上（DeepSeek-V2 里提到的 device limited routing）。其具体通信方式为想传输到目标 node index 相同的 rank 上，然后再通过节点内部通信传输到目标的 rank 上。通过这种方式，我们可以让节点间通信与节点内部通信进行重叠，从而每个 token 可以从一个节点上选取 3.2 个专家，这样 DeepSeek-V3 可以在不损失效率的情况下最高选取 $13$ 个专家。\n作者进一步采用了 warp specialization 技巧来将 20 个 SM 划分为 10 个通信 channel. 作者分别针对 dispatching 和 combing 阶段使用了不同数量的 warps.\nMemory saving 作者使用了如下技巧来减少内存访问：\nRecomputation of RMSNorm and MLA Up-Projection. 在 backward 过程中重新计算 RMSNorm operation 以及 MLA up projection 来避免存储器对应的输出 Exponential Moving Average in CPU. 作者将 EMA 参数保存在 CPU 中，然后进行异步更新来进一步减少内存访问 Shared Embedding and Output Head for Multi-Token Prediction. 作者将 embedding layer 和 output head 放在一个 PP rank 上，这样就可以提高 MTP 的内存访问效率 FP8 Training 作者提出了一个基于 FP8 的混合精度训练框架。作者提出了两个改进方案：\n分组量化，将 tensor 按照 tile 分组或者按照 block 分组来分组量化，这样就避免了全局量化的精度损失 高精度累加，作者在乘法计算时，使用了 FP8 格式，然后在累加阶段，使用了更高精度的格式 为了进一步减少内存和通信开销，对于 activation 的 cache 以及 dispatch, 作者使用了 FP8 数据格式，然后对于优化器状态，作者使用了 BF16 数据格式。下面是对上面改进的具体说明。\nMixed Precision Training 作者在本文中提出了使用 FP8 混合精度进行预训练，作者参考了 low precision training 构建 FP8 训练框架，即计算量高的使用 FP8 精度，计算量低的使用原本的数据精度, 框架如下图所示\n其中各个模块使用的精度如下表所示\nPrecision Modules FP8 Linear (Fprop, Dgrad, Wgrad) higher precision - embedding\n- output head\n- Moe gating\n- normalization\n- attention operator\n- master weights\n- weight gradients\n- optimizer states Enhancing Low-precision Training Accuracy 作者介绍了几个策略用于提高 FP8 混合精度训练的表现：\nfine-grained quantization: 对于激活值，作者将 tensor 分割为 $1\\times 128$ 大小的 groups, 然后每个 group 内部进行 quantization; 对于权重，作者将其分割为 $128\\times 128$ 大小的 groups, 然后进行 quantization, 这样可以降低 quantization error, 如上图左图所示 increasing accumulation precision: 低精度训练会带来 underflow 的问题，而 FP8 GEMM 累加仅能保留约 14bit 精度，远低于 FP32 的 32bit 精度。为了解决这个问题，作者采用了 Tensor Core 不分累加 +CUDA core 高精度聚合的协同策略。即在 Tensor Core 上执行 MMA 指令时，先按照 14bit 精度对 $N_C$ 个元素进行累加，当达到 $N_C$ 时，作者将结果复制到 CUDA core 的 FP32 寄存器中，在 FP32 精度下完成累加。过程如上图右图所示 Mantissa over exponents. 与之前的工作不同，作者使用了 E4M3 的数据格式来达到更高的精度 Online Quantization. 为了降低 quantization 的误差，作者实时计算了 activation block 以及 weight block 的最大绝对值来导出 scaling factor 并量化为 FP8 精度 Low Precision Storage and Communication 作者通过压缩 cached activation 和 optimizer states 来进一步减少内存访问以及通信访问次数\nLow-precision optimizer states: 对于 optimizer states, 作者使用了 BF16 数据格式来保存 AdamW 的优化器的一阶和二阶动量，但是对于 master weight 和 gradients 作者仍然使用了 FP32 来保证训练的数值稳定性 Low-Precision Activation: 对于 linear operator, 作者将其 cache activation 使用 FP8 格式进行存储，对于 attention 的输出，作者使用了 E5M6 数据格式来存储 activations, 这些 activations 的 tile size 为 $1\\times 128$, 作者还是用了 2 的幂次作为 scale factor 来减少 quantizationerror; 对于 SwiGLU, 作者将其输入也保存为 FP8 的数据格式来降低内存消耗 Low-Precision Communication: 在 MoE up-projection 之前，作者将 attention 的输出量化为 FP8 数据格式再执行 dispatch 操作，这样可以降低通信开销。在反向传播的时候同理，先进行 FP8 量化，然后再进行反向 dispatch. 对于 combine 阶段，为了避免精度损失，作者还是使用了 BF16 数据格式 作者对比了 FP8 和 BF16 精度训练，结果如下图所示\n实验结果显示，FP8 混合精度训练的损失降低不足 $0.25\\%$.\nInference and Deployment 作者在 H800 集群上部署 DeepSeek-V3, 作者分别针对 prefilling 和 decoding 两个阶段进行了优化\nPrefilling prefilling 阶段在 4 节点 32 GPU 上进行，并行策略如下\n模块 并行策略 说明 Dense MLP 1-wat TP 减少 TP 通信 Attention 4-way TP, SP, 8-way DP SP 用于长文本处理 MoE 32-way EP 每个 GPU 包含 8 个专家 为了实现负载均衡，作者提出了redundant experts的策略，具体就是将负载比较高的专家进行复制。模型会周期性进行统计，然后计算出负载比较高的专家，接下来每个 GPU 除了原来的 8 个专家之外，还会 host 一个额外的 redundant expert.\n为了提高 throughput, 作者同时处理两个 micro-batch, 来重叠 attention, MoE 和 dispatch, combine 通信，如下图所示\n作者还探索了 dynamic redundancy 策略，即每个 GPU host 更多的专家（比如 16 个）然后 inference 的每一步仅激活其中的 9 个，在进行 all-to-all operation 时，作者首先进行 global optimal routing 来计算最合适的专家。作者认为 prefilling 计算量很大，因此 routing 的计算量可以忽略不计\nDecoding 在 decoding 阶段，作者在 40 个节点 320 GPU 上进行部署，并行策略如下\n模块 并行策略 说明 Attention 4-way TP, SP, 80-way DP SP 用于长文本处理 MoE 320-way EP 每个 GPU 包含 8 个专家 这个阶段，每个 GPU 只 host 一个专家，64 个 GPU 负责 host redundant expert 以及 shared expert. 作者通过在 IB 上直接进行 P2P 的传输来降低通信的开销，作者还是用了 IBGDA 来进一步最小化 latency 以及提高通信效率\n在这个阶段，attention 的计算占据了大部分时间，因此，作者采取了如下的 overlap 策略\n由于 decoding 阶段一般 batch size 比较小，因此整个 decoding 的瓶颈在于 memory access, 因此为了避免影响 attention 的计算效率，作者仅安排一小部分 SM 用于 dispatch-MoE-combine.\nSuggestions on Hardware Design Communication Hardware 尽管作者将计算与通信重叠来提高训练效率，但是已有的通信仍然依赖于 SM, 这样限制了计算的效率。作者希望工能够构建专门针对通信的设备，另一方面，作者希望统一 IB 和 NVLink 来降低实现难度\nComputation Hardware 提升 FP8 GEMM 的累加精度，当前精度只有 14bit, 作者希望能够在 GPU 设计上进行改进，将这个精度提升到 34-bit 支持 tile 以及 block 层面的 quantization, 尽管前文作者已经设计出了改进的算法，但是 Tensor Core 以及 CUDA 之前平凡的数据移动降低了计算效率，作者希望未来能够支持细粒度的 quantization online quantization, 作者希望将 FP8 cast 以及 TMA 结合在一起，从而 quantization 可以在传输的时候完成计算，减少了内存读写。作者还建议使用 warp-level cast instruction Transposed GEMM operations. 本文中，矩阵被切分为不同的 tiles, 在计算的时候需要先加载这些 tiles 然后进行 dequantization, transpose 等操作，作者希望未来能够直接支持 transposed reads of matrices Pre-training Pre-training Data 相比于 DeepSeek-V2, DeepSeek-V3 提升了数学和代码数据的比例，以及增加了多语种数据。最终训练数据一共包括 14.8T\n作者还使用了 DeepSeekCoder-V2 里应用的 Fill in the middle 策略来让模型基于上下文越策中间的文本，对应的数据格式如下\n1 \u0026lt;|fim_begin|\u0026gt;f_pre\u0026lt;|fim_hole|\u0026gt;f_suf\u0026lt;|fim_hole|\u0026gt;f_middle\u0026lt;|fim_end|\u0026gt; 这个结构与 sequence packing 结合在一起。\nTokenizer 基于 BBPE, 大小为 128K tokens. 在训练时，作者将随机一部分 combine token 进行切分来减少 token boundary bias 问题\nHyper-parameters 模型参数如下表所示，最终 DeepSeek-V3 拥有 671B 总参数，激活参数为 37B\nvariable notation value layers $\\ell$ 61 dense layers - 3 hidden dimension $d$ 7168 num of attention heads $n_h$ 128 head dimension $d_h$ 128 KV compression dimension $d_c$ 512 query compression dimension $d_c'$ 1536 decouple query and key dimension $d_h^R$ 64 routed expert $N_r$ 256 shared expert $N_s$ 1 MoE hidden dimension $d_{MoE}$ 2048 activated experts $K$ 8 limited node routing $M$ 4 MTP depth $D$ 1 训练时，作者使用了 learning rate scheduling, batch size scheduling 等方法。对于 PP, routed expert 会均匀分布为 8node 对应的 64 个 GPU 上，预训练时模型的上下文长度为 4k\nLong Context Extension 作者使用了 YARN 来扩展模型的上下文长度，模型上下文长度经过两个额外的训练阶段从 4K 扩展到 32K 再扩展到 128K, 均训练了 1000 步。\nYARN 配置与 DeepSeek-V2 基本一致，在额外的 decoupled query and key 上作者没有应用这一点。参数配置如下表\nparameter $s$ $\\alpha$ $\\beta$ $\\sqrt{t}$ value 40 1 32 $0.1\\ln s+1$ 第一阶段的上下文长度为 32K, batch size 为 1920, 第二个阶段的上下文长度为 128K, batch size 为 480.\nPerformance DeepSeek-V3 base 的表现下图所示，作者对比了 DeepSeek-V2, Qwen2.5, LLaMA 3.1\nDiscussion 作者首先验证了 MTP 的有效性，结果如下图所示\n可以看到，在模型架构相同，训练数据相同的情况下，1-MTP 的效果超过了 baseline 的表现，说明了 MTP 策略的有效性\n接下来，作者还验证了 Loss-Free Balancing 的有效性，结果如下图所示\n可以看到，loss-free Balancing 的效果比 loss balancing 的效果更好\n接下来，作者对比了 loss-free balancing 和 sequence-wise auxiliary loss, 即 batch-wise v.s. sequence-wise. 作者认为，前者的约束更灵活，因为其不要求 in-domain balance. 作者在测试集上进行可视化，结果如下图所示\n实验结果发现，loss-free 策略对应的 expert specialization 更强。作者进一步设计了一个 batch-wise auxiliary loss 来实现 batch-wise load balance, 结果发现，这种策略也能达到和 loss-free balancing 一样的效果，这说明了 batch-wise load balancing 效果更好\n最后，作者提了两点 loss-free 策略的问题：\n在特定的 sequence 或者小 batch 里出现负载不均衡。作者通通过增大 batch size 来解决这个问题 在推理阶段因为 domain-shift 导致的负载不均衡。作者实现了一个基于 redundant expert 的推理框架来解决这个问题 Post-training SFT post-training 包含 1.5M 样本，数据包括 reasoning 数据以及 non-reasoning 数据，前者由 DeepSeek-R1 合成，后者由 DeepSeek-V2.5 合成\nSFT 时，作者训练了两个 epoch, 使用了 sequence packing 技巧\nRL Reward model 包含 rule-based reward model 和 model-based reward model.\nRL 训练使用的算法为 GRPO\nPost-training Performance Post-training Discussion 作者首先探究了 Distillation 对模型表现的影响，作者使用 DeepSeek-R1 来蒸馏 DeepSeek-V2.5, 结果如下图所示\nModel LiveCodeBench-CoT MATH-500 Pass@1 Length Pass@1 Length DeepSeek-V2.5 Baseline 31.1 718 74.6 769 DeepSeek-V2.5 +R1 Distill 37.4 783 83.2 1510 可以看到，distillation 可以显著提高模型的表现。但是其问题在于也会让模型输出的长度增加。作者认为知识蒸馏是 post-training optimization 的一个重要方向。\n接下来，作者讨论了 self-rewarding, 具体做法就是使用 DeepSeek-V3 来对评估结果进行投票，结果发现最终的效果很好，作者认为 LLM 可以很好地将非结构化信息转换为 rewards.\n最后，作者讨论了 MTP. MTP 可以于 speculative decoding 结合，进一步提高 decoding 的速度。通过实验作者发现，second token prediction 的接受率在 $85\\%\\sim 90\\%$ 之间，说明了其有效性。\nConclusion 在本文中，作者提出了 DeepSeek-V3, 一个 671B-A37B 的 MoE 大语言模型，训练 token 数为 14.8T. 作者使用了 loss-free-balancing strategy 来实现负载均衡。训练时作者使用了 FP8 混合精度。评估发现 DeepSeek-V3 的达到了 SOTA 表现。\n作者认为，DeepSeek-V3 model size 太大，不适合部署。第二，部署策略需要进一步改进。\n最后，作者认为未来工作有以下几点：\n改进模型架构，进一步提高训练以及推理效率，扩展模型的上下文 提升训练数据的数量和质量 提高模型的 reasoning 能力 更详尽的评估 References Arxiv ","date":"2025-12-08T11:14:45+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseek-v3/","title":"Notes on DeepSeek-V3"},{"content":"介绍 DeepMind 在 6 月 17 号发布了 Gemini2.x 系列的技术报告，包括\nGemini 2.5 Pro Gemini 2.5 Flash Gemini 2.0 Flash (earlier) Gemini 2.0 Flash-Lite (earlier) 技术报告简单说了一些技术细节，主要还是模型的评估\n注：Gemini 2.0 Flash 和 Gemini 2.0 Flash-Lite 将要被 Gemini 2.5 Flash 和 Gemini 2.5 Flash-Lite 取缔，见最新的 blog\nGemini2.x 系列亮点：\n领先的 coding 和 reasoning 能力 超过 1M 的上下文，可以处理超过 3 个小时的 video 集成 long context, multimodal 和 reasoning 三种能力的 agentic workflow 能力 模型能力对比\nModel Gemini 1.5 Flash Gemini 1.5 Pro Gemini 2.0 Flash-Lite Gemini 2.0 Flash Gemini 2.5 Flash Gemini 2.5 Pro Input Modalities Text, Image, Video, Audio Text, Image, Video, Audio Text, Image, Video, Audio Text, Image, Video, Audio Text, Image, Video, Audio Text, Image, Video, Audio Input length 1M 2M 1M 1M 1M 1M Output modalities Text Text Text Text, Image* Text, Audio* Text, Audio Output length 8K 8K 8K 8K 64K 64K Thinking No No No Yes Dynamic Dynamic Supports tool use? No No No Yes Yes Yes Knowledge cutoff November 2023 November 2023 June 2024 June 2024 January 2025 January 2025 模型场景使用对比\nModel Gemini 2.5 Flash-Lite Gemini 2.5 Flash Gemini 2.5 Thinking No Yes Yes 使用场景 大规模调用 日常使用 coding 或者 reasoning 人物 速度 非常快 快 一半 表现 一半 强 非常强 输入价格 0.1 0.3 1.25 输出价格 0.4 2.5 10 模型表现\n模型吞吐量对比\n架构，数据与训练 架构 Gemini2.5 是一个基于 MoE 的 transformer 架构，支持 text, vision, audio 模态\nFlash 系列使用的是知识蒸馏的方法训练得到的，训练时使用了 $k$-sparse 的策略，也就是只保留教师模型输出概率最高的 $k$ 的词以及对应的概率。作者认为知识蒸馏可以有效提高小模型的能力。\nInfra Gemini 系列在 TPUv5p 的架构上进行训练。作者主要提了两点：\nSlice-Granularity Elasticity：可以在部分 TPU 出现故障时快速切换并继续训练 Split-Phase SDC detection：通过轻量级重放和校验机制，在几分钟内就能识别出有问题的硬件设备 Post-training post-training 包含 SFT，reward model 以及 RL 的训练。\n在 RL 阶段，奖励来自 verifiable rewards 和 model-based generative rewards\n能力提升 技术报告提到了几个方面能力的提升\ncode pre-training 阶段，加入了大量的代码数据，作者还评估了代码数据的质量 post-training 阶段，作者基于 reasoning 能力构建了一系列的工程任务，来提高模型解决问题的能力\nFactuality 通过 search 和 tool use，reason about output 以及 issue follow-up queries 来验证 factual accuracy\nMultilinguality 预训练时使用了 400 多种语言的语料进行训练\nAudio 训练模型完成 audio generation 任务，生成的时候使用了causal audio representation，训练数据覆盖了 200 多种语言\nVideo 通过降低每帧视频对应的 visual token 个数（258-\u0026gt; 66），来让模型可以处理 3 个小时的视频\nEvaluation 对比了 Claude_4, o3, DeepSeek-R1 和 Grok-1\nConclusion 结论里作者主要提到了两点\n模型能力的提升已经超过了 benchmark 的构建速度和成本 未来如何设计经济的，覆盖广的，能动态调整难度的 benchmark 是一个关键问题 技术报告中作者还提到了 Gemini Plays Pokemon 的 case study，作者提到了两点问题：\n作者发模型对视觉信息的依赖程度并不是很高 尽管模型上下文长度超过了 1M，但是对于这种复杂的 long horizon 问题，当输入超过了 100K token 之后，模型倾向于重复过去的行为，而不是生成新的计划 因此，未来如何解决 multi-turn, long-horizon 的 agentic task 也是一个值得探究的方向。 Reference Gemini 2.5 Flash/Flash Lite Gemini 2.5 Technical Report ","date":"2025-12-06T18:14:15+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gemini2.5/","title":"Notes on Gemini2.5"},{"content":"Introduction 作者在本文中主要想完成三个目标：\n基于已有框架尝试训练一个 MoE 大语言模型 详细分析 MoE 模型的 routing 机制 为后续 MoE 模型开发提供经验 基于这三个目标，作者在本文中发布了 4 个 MoE 模型。预训练时，作者使用了 $52.25\\%$ 的代码数据来提高模型的表现。作者还是用了 UL2 作为训练目标。\n通过实验作者发现 MoE 机制存在以下三个性质：\nContext-independent Specialization: 即 MoE 模型倾向于基于 token semantics 进行聚类，而不是 context early routing learning: routing 的分布情况在训练早期就已经确定了 drop towards the end: 使用 token dropping 策略之后，越往后的的 token 被丢弃掉的概率也就越大 Method Model 作者介绍了模型使用的数据集如下\n模型的 tokenizer 基于 umT5, 大小为 256K, umT5 支持多语种，并且还有 fallback 机制来处理 OOV 的 token\n模型配置如下\nModel OpenMoE-Base/16E OpenMoE-8B/32E OpenMoE-34B/32E total params 650M 8.7B 34B activated params 142M 2.1B 6B total experts 16 32 32 activated experts 2 2 2 shared experts 1 1 1 Layout every 4 every 6 every 4 hidden_dim 768 2048 3072 ffn_hidden_dim 3072 8192 12288 num_heads 12 24 24 head_dim 64 128 128 num_layers 12 24 32 Training 作者使用了 Load Balancing loss 和 ST-MoE 提出的 Router Z-loss, 最终训练的目标函数与 ST-MoE 一致\n作者在训练时，还是用了 UL2, UL2 结合了 mask language modeling 和 casual language modeling 两种训练方式。\n作者首先对各个实验配置进行了验证，其中关键接论文模型很容易在 code data 上达到比较高的准确率以及比较低的 loss, 作者认为这是因为 code data 中存在大量的特殊符号。\n训练过程中，作者发现模型在训练一定步数之后，容易出现过饱和现象，因此作者将训练目标函数由 UL2 降为 CasualLM, 并未作者还将代码数据的比例降低到 $15\\%$.\nExperiments Analysis OpenMoE 作者发现，大部分专家对于不同的 domain 没有出现 specialization 情况，对于 math domain, specialization 现象比较明显，作者认为这是因为 math domain 包含更多的 special tokens 对于不同的语言，有部分专家出现 specialization 现象 部分专家对于 position ID 有 specialization 现象，并且连续的 token 更偏好相同的专家 作者还发现，部分专家对于 Token ID 有 specialization 现象，作者将这种现象称为 Context-independent Specialization. 专家还会对语义相似的 token 进行聚类，并且这种聚类在训练早期就已经发生，作者认为其原因在于重新分配 token 会增加最终的 loss 对于 token dropping, 作者发现月考后的 token, 其被 drop 的概率比例也越高。并且对于指令跟随数据，更多的 token 都会被丢掉，因此作者认为指令跟随数据是 MoE 模型的一种 OOD 数据 Analysis on other MoE Models 作者还分析了 Mixtral MoE 和 DeepSeekMoE 两个模型的专家路由情况。这两个模型采用了 dropless token routing 的机制。\n首先，作者分析了这两个模型对于 TokenID 的敏感性，结果发现，DeepSeek-MoE 的 specialization 现象比较明显，而 Mixtral MoE 由于使用了 upcycling, 其 specialization 现象不太明显，作者认为这是因为 upcycling 导致每个专家的权重都差不多。最终，作者认为对于 training from stratch 的 MoE model 这个 specialization 现象更明显。\nConclusion 作者在本文中提出了一个全开源的 MoE 大模型系列，包括 0.6B, 8.7B 和 34B 三个 size, 作者还对 MoE 中的 routing 进行了详细的分析，结果发现【【】】，这些发现有助于后续的 MoE 模型架构的研究。\nReferences Arxiv ","date":"2025-12-06T18:08:11+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-olmoe/","title":"Notes on olmoe"},{"content":"Introduction 作者在本文中提出了 Qwen3-VL 系列多模态大模型，包括 4 个 dense 模型和两个 MoE 模型，模型的上下文长度为 256K, 通过数据和训练上的优化，作者保持了模型的纯文本能力。最终 Qwen3-VL 包括 non-thinking 和 thinking variants.\n在架构上，Qwen3-VL 进行了三点改进：\nInterleaved MRoPE: 作者解决了 Qwen2.5-VL 提出的 MRoPE 在长视频理解场景下的频谱不平衡问题 DeepStack: 作者使用了 DeepStack 来提取 ViT 不同 layer 的视觉特征 Explicit Video timestamps: 作者使用了绝对时间来标记 frame 来提供更直接的时间信息 在数据上，作者使用了 image caption, OCR, grounding, spatial reasoning, code, long documents 以及 temporally grounded video 等数据，作者 还是用了 GUI-agent interaction 数据来提高模型的 action 能力\n在训练上，Qwen3-VL 包含两个大的阶段：pre-training 和 post-traing, pre-training 包含 4 个小阶段，post-training 包含 3 个阶段。\nArchitecture Qwen3-VL 的架构如下所示\n其中，\nLLM: LLM 使用了 Qwen3 系列大语言模型，包括 2B, 4B, 8B, 32B 四个 dense model 以及 30B-A3B, 235B-A22B 两个 moe 模型 Vision Encoder: encoder 基于 [[SigLip-2]] 初始化，然后使用了 dynamic input resolutions 进行 continue training, 作者使用了 CoMP 提出的 2D-RoPE 以及 interpolate absolute position embedding, 最终包括 SigLip2-SO-400M 和 SigLip-Large (300M) 两个 size, 后者用于 2B 和 4B 两个 size Patch Merger: 一个 2 层的 MLP, 将四个 visual token 压缩为 1 个 Interleaved MRoPE 这部分介绍见 [[MRoPE-Interleave]]\nDeepStack 受 Deepstack 启发，作者从 vision encoder 的中间层（具体来说是第 8， 16， 24 层）提取对应的视觉特征，然后经过 MLP 与 LLM 对应 layer 的视觉 token 直接进行相加。\nVideo Timestamp 作者发现，Qwen2.5-VL 中使用的 MRoPE 存在如下问题：\n将 temporal position 与绝对时间绑定之后，对于长视频会产生非常大且稀疏的 temporal position ids 需要使用不同的 FPS 进行采样来提高模型的泛化性 为了解决这个问题，作者使用了一个 textual token-based time encoding strategy, 其中每个 video temporal patch 对应的 timestamp 表示为 \u0026lt;3.0 seconds\u0026gt;, 这样视频会被处理为以下格式\n1 \u0026lt;0.0 seconds\u0026gt; \u0026lt;video token\u0026gt; \u0026lt;video token\u0026gt; ... \u0026lt;4.0 seconds\u0026gt; \u0026lt;video token\u0026gt; \u0026lt;video token\u0026gt; 在训练时，作者还使用了 seconds 以及 HMS 两种格式来提高模型对于不同格式的泛化能力。作者认为，虽然这种表示会提高上下文长度，但是也能够提高模型 video grounding 或者 dense captioning 等时序信息敏感任务的表现\nPre-training Training Recipe 预训练阶段包含 4 个阶段，如下图所示\nStage 0: 这一阶段的目的是对齐视觉特征和文本特征，只训练 Patch merger, 训练使用了 67B token, 覆盖 image-caption, knowledge, OCR 数据，上下文长度为 8192 Stage 1: 这一阶段所有参数都参加训练，训练使用了 1Ttoken, 作者在训练是加入了纯文本数据，最终数据包含 interleaved image-text, visual grounding, VQA, STEM, video 数据，上下文长度为 8192 Stage 2: 这一阶段的目的是扩展模型的上下文长度到 32K, 训练使用了 1T token, 数据包括长视频以及 agent-oriented instruction-following 数据 Stage 3: 这一阶段的目的是将模型的上下文长度进一步扩展到 262K, 训练使用了 100B token. 数据包括长视频以及长文本 Data Image Caption Data: 作者使用了 Qwen2.5-VL 32B 来进行 re-captioning, 然后进行了 de-duplication 以及 clustering 来提高数据的质量和多样性 Interleaved Text-Image Data: 作者对文档进行分裂，然后使用微调的 Qwen2.5-VL 7B 来进行解析，对于长文本，作者将连续页面拼接在一起。作者使用了对齐以及页数来保证数据的质量 Knowledge Data: 作者构建了多个类别的数据，然后对这些数据进行 refine OCR: 作者构造了 30M 的数据以及 1M 的多语种数据 Document Parsing Data: 作者从 CC 上收集了 3M PDF 以及处理了自有的 4M 数据，最终数据集里包含合成数据和真实数据；对于长文档理解数据，作者通过将 single-page 数据 merge 在一起得到，然后作者构造了 long document VQA 数据 Grounding and counting Data: grounding 数据包括 box-based 和 Point-based 两种形式，均从开源数据集收集得到，前者包括 RefCOCO, Object365, 后者包括 PixMo; 对于 Counting, 作者基于 grounding 数据构造了 direct counting, box-based counting 以及 point-based counting 三种形式 Spatial Understanding: 数据包括 spatial understanding 和 3D grounding 两类数据，前者的数据使用了相对位置关系来提高 spatial reasoning 的 robustness; 后者使用了 Omni3D 来统一数据格式 Code: 包括 Qwen3, Qwen3-Coder 的纯文本 coding 数据，以及多模态 coding 数据，覆盖了将 UI 截图转换为 HTML/CSS 以及从图片生成 SVG 等任务 Video: 包括 Dense Caption Synthesis 以及 Spatial-Temporal Video Grounding 两个任务。作者还对不同来源不同长度的数据进行了平衡 STEM: 作者构造了一个合成数据 pipeline, 合成了 1M point-grounding samples, 2M perception-oriented VQA 数据，最终数据集包含 6M 标注图表数据，覆盖了 STEM 相关学科；对于多模态推理数据，作者收集了 60M 的 K12 以及本科生级别的练习题，作者还合成了 12M 的多模态推理数据。除了多模态推理数据，作者还加入了纯文本推理数据 Agent: 这部分数据包括 GUI, function calling 以及 Search 三部分， GUI 数据通过数据合成得到，Function calling 数据通过强模型生成轨迹得到，search 数据通过收集执行搜索轨迹得到 Post-training Post-training 包含三个阶段：\nSFT: 提高模型的指令跟随能力，SFT 又分为了两个小阶段，上下文长度分别为 32K 和 256K, 对于 instruct 和 reasoning 版本，作者设计了不同的数据格式，后者包含 CoT reasoning trace Strong-to-Weak Distillation: 提高小模型的能力，这里应该是和 Qwen3 一样，将大模型的能力蒸馏到小模型里 RL: 提高模型的 reasoning 能力以及人类偏好对齐。这里包含了 Reasoning RL 以及 General RL 两个阶段，覆盖了 math, OCR, grounding, instruction following 等 domain 整体的训练 pipeline 我猜测应该是这样：\nCode-start Data Code-start Data 分为 SFT 数据和 Long CoT SFT 数据，前者用于训练 instruct 版模型，后者用于训练 reasoning 版模型\nData tasks samples training filtering SFT spatial reasoning\nimage-grounded reasoning\nspatio-temporal grounding\nlong document understanding 1.2M (1/3 are text-only) - stage 1: 32K\n- stage 2: 256K - query - rule-based\n- model-based Long CoT SFT VQA, OCR, 2D/3D grounding, video analysis, STEM, agent text:multimodal = 1:1 - difficulty\n- multi-modal\n- response quality Strong-to-Weak Distillation 蒸馏过程包括两个阶段：\noff-policy Distillation: 使用教师模型的输出进行训练提高模型基本的 reasoning 能力 On-policy Distillation: 使用教师模型输出的 logit 作为蒸馏信号提高模型的 reasoning 能力 RL Reasoning RL 作者收集了 30K 的 RL 数据，然后对通过率超过 90% 的数据进行过滤 (16 responses per query), 对于 reward, 作者构建了一个 unified reward framework 来提供奖励\n训练时，作者使用了 SAPO 算法进行训练\nGeneral RL 作者采用了一个 multi-task RL 的范式来提高模型在不同任务上的表现，reward 主要包含两个方面：\ninstruction following: 评估模型遵循用户指令的能力，包括内容，格式，长度等 preference alignment: 对于开放式问题，评估模型帮助性，事实准确性等方面的表现 基于这两个方面 reward 有两个部分组成：\nrule-based reward: 基于规则的 reward, 比如格式要求等 model-based reward: 使用 Qwen2.5-VL 72B 和 Qwen3 作为 judge model 来提供奖励 为了解决模型的重复性实处，中英文混杂等问题，作者构造了一个数据集来故意触发模型这些问题然后加以改正。\nThinking with Images 作者还够在了数据提高模型的 \u0026ldquo;thinking with images\u0026rdquo; 的能力，训练包含两个阶段：\nStage 1: 作者构造了 10K Grounding 数据，然后对 Qwen2.5-VL 32B 进行 SFT 来模仿 agent 的行为: think -\u0026gt; act -\u0026gt; analyze feedback -\u0026gt; answer, 然后作者使用 multi-turn, tool-integrated RL 来进一步提高模型的 reasoning 能力 Stage 2: 作者从 Qwen2.5-VL 32B 蒸馏得到 120K multi-turn agentic interactions 数据集， 然后作者使用了相似的 cold-start SFT 以及 tool-integrated RL pipeline 来训练 Qwen3-VL 这里 RL 训练的 reward 包含以下几部分：\nanswer accuracy reward multi-turn reasoning reward tool-calling reward Experiments Performance Qwen3-VL 235B-A22B 的表现如下图所示\nAblation Study 作者对比了以下 Qwen3-ViT 和 SigLIP-2 的表现，结果如下图所示\n实验结果显示，使用 1.7B 的 Qwen3 和 1.5T tokens 进行训练之后，Qwen3-ViT 的表现超过了 SigLIP2 的表现，验证了 Qwen3-ViT 的有效性\n作者对比了 Deepseek 和 baseline 的表现，结果如下图所示\n可以看到，相比于 baseline, DeepStack 的表现更好，说明了 DeepStack 可以提供更丰富的视觉信息。\n作者还评估了以下 Qwen3-VL 在视频版大海捞针任务上的表现，实验结果发现，对于 30 分钟的视频，Qwen3-VL 的准确率为 $100\\%$, 通过 YARN 上下文扩展策略，模型在 2 个小时视频上的准确率为 $99.5\\%$.\nConclusion 作者在本文中提出了 Qwen3-VL 系列多模态大模型，在架构上，作者使用了 interleaved-MRoPE, DeepStack 等改进策略，在数据上，作者扩展了训练数据的多样性，在训练上，作者分别训练了 instruct 版本和 reasoning 版本。最终评估发现，Qwen3-VL 达到了 SOTA 表现。\n作者认为，未来的工作在于\n基于 Qwen3-VL 构建具身智能 agent 提高模型的可交互感知，tool-augmented reasoning 以及 real-time multimodal control 能力 提高模型与人类学习，合作的能力 统一理解与生成多模态大模型 References blog Arxiv ","date":"2025-12-05T10:12:01+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen3-vl/","title":"Notes on Qwen3 VL"},{"content":"Introduction 相关工作包括 GRPO 以及 GSPO\nSAPO 的关键思想有两点：\ntokne-level soft trust region 可以保证 sequence-level coherence 非对称的 temperature 可以针对 postive token 和 negative token 进行不同的优化 Method 作者首先给出了 SAPO 的目标函数如下：\n$$ \\mathcal{J}_{\\mathrm{SAPO}}(\\theta) = \\mathbb{E}_{(x, y)\\sim\\mathcal{D},\\{y_i\\}_{i=1}^G\\sim \\pi_{\\theta_{old}}(\\cdot\\mid x)}\\left[ \\frac{1}{G}\\sum_{i=1}^G\\frac{1}{|y_i|}\\sum_{t=1}^{|y_i|}f_{i,t}(r_{i,t}(\\theta))\\hat{A}_{i,t} \\right], $$其中，\n$$ f_{i,t}(x) = \\sigma(\\tau_{i,t}(x-1))\\cdot \\frac{4}{\\tau_{i,t}}, \\tau_{i,t} = \\begin{cases} \\tau_{pos}, \u0026 \\text{if }\\hat{A}_{i,t}\u003e0\\\\ \\tau_{neg}, \u0026\\text{otherwise} \\end{cases} $$这里\n$$ \\hat{A}_{i,t} = \\hat{A}_{i} = \\frac{r(x,y_i) - \\mathrm{mean}(\\{r(x,y_i)\\}_{i=1}^G)}{\\mathrm{std}(\\{r(x,y_i)\\}_{i=1}^G)},\\quad r_{i,t}(\\theta) = \\frac{\\pi_{\\theta}(y_{i,t}\\mid x, y_{i,\u003c t})}{\\pi_{\\theta_{old}}(y_{i,t}\\mid x, y_{i,\u003c t})} $$$\\tau_{pos}$ 和 $\\tau_{neg}$ 分别是 positive token 以及 negative token 对应的温度, $\\sigma(x)=1/(1+e^{-x})$ 是 sigmoid function.\n对 $\\mathcal{J}_{\\mathrm{SAPO}}(\\theta)$ 求导得到\n$$ \\nabla_\\theta \\mathcal{J}_{\\mathrm{SAPO}}(\\theta) = \\mathbb{E}_{(x, y)\\sim\\mathcal{D},\\{y_i\\}_{i=1}^G\\sim \\pi_{\\theta_{old}}(\\cdot\\mid q)}\\left[ \\frac{1}{G}\\sum_{i=1}^G\\frac{1}{|y_i|}\\sum_{t=1}^{|y_i|}w_{i,t}(\\theta)r_{i,t}(\\theta)\\nabla_\\theta \\log \\pi_\\theta(y_{i,t}\\mid q, y_{i, ","date":"2025-12-05T10:09:06+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-sapo/","title":"Notes on SAPO"},{"content":"Introduction 已有的 MLLM 将视觉 token 作为一个 1d sequence, 输入给 LLM. 在本文中，作者将 visual token 注入到 LLM 的不同 layer 中来提高视觉信息的利用率\nMethod 首先，对于输入的图片 $I$, 我们将其分为高精度图片版本 $I_{high}$ 和低精度图片版本 $I_{low}$, $I_{low}$ 通过 vision encoder 和 MLP 得到对应的视觉 token $X_v$ 作为 LLM 的输入，然后在 LLM transformer block 的第 $i$ 层，其对应的视觉 token $X_{i,v}$ 会与 stack feature $X_{v}^i$ 相加，这里 $X_v^i$ 是对高精度图片输入的一个采样，即\n$$ X_v^i = \\mathrm{Sampling2D}(\\mathrm{MLP}(\\mathrm{ViT}(I_{high}))) $$算法伪代码如下所示\n1 2 3 4 5 6 7 8 9 10 11 12 # H0: Input embeddings for LLM (Original inputs args for traditional LMM); # vis_pos: the location of visual tokens; # X, Xstack: Original visual tokens, Extra high-resolution visual token list; # lstart, n: Index of starting layer, and layer interval for stacking. def forward(H0, Xstack, lstart, n, vis_pos): H = H0 for (idx, TransformerLayer) in enumerate(self.layers): # DeepStack: if idx \u0026gt;= lstart \u0026amp; (idx − lstart) % n == 0: H[vis_pos] += Xstack[(idx − lstart)//n] # Original Transformer: H = TransformerLayer(H) Experiments Ablation Study 作者进一步验证了不同实验配置，结果发现在 early layer 进行 deepstack 效果最好，越往后效果越差\n作者还在 ViT 上应用了 DeepStack 策略，结果发现 ViT 的效果也有所提升\n作者还发现，模型表现提升是因为加入了 high-reoslution image token 信息\nConclusion 作者在本文中提出了 DeepStack, 一个提高 MLLM 中视觉信息利用率的方法，作者验证了这个方法的有效性。\nReference paper ","date":"2025-12-04T17:32:41+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepstack/","title":"Notes on DeepStack"},{"content":"Introduction Transformer 架构在 NLP 领域已经成为了事实上的标准。而在 CV 领域，目前还是 CNN 架构占据了主导。\n在本文中，作者就探究是否可以应用 Transformer 架构来完成图像识别的任务。\n作者发现，当在小规模数据集上进行训练时，Transformer 架构的表现是不如 CNN 架构的，作者认为原因是 Transformer 架构缺乏如 translation equivariance 和 locality 等 inductive biases. 但是当数据集规模上去之后，作者发现这种 inductive bias 可以通过大规模训练抵消掉。其中，最好的模型在 ImageNet 上达到了 $88.55\\%$ 的准确率。\nMethod ViT 的架构如下图所示\n为了能够处理图片，对于输入的图片 $x\\in\\mathbb{R}^{H\\times W\\times C}$, 其中 $(H,W)$ 是图片大小，$C$ 是 channel 的个数。作者将其展开为一个 2D patch 序列, $x_p\\in\\mathbb{R}^{N\\times (P^2\\times C)}$, 其中 $(P, P)$ 是 image patch 的大小，$N=HW/P^2$ 是 image patch 的的个数，也是 Transformer 输入 token 的个数，Transformer 的 hidden size 为 $D$, 作者使用了一个 linear layer 来将 image patch 转换为 Transformer 的输入。\n与 BERT 一致，作者使用了一个 [class] token 来作为输入 patch 序列的 embedding, 即 $z_0^0=x_{class}$, 其对应的输出 $z_L^0$ 作为该图片的表示。作者还使用了 1D 的 position encoding. 最终，ViT 表达式如下所示\n$$ \\begin{aligned} z_0 \u0026= [x_{class};x_p^1\\mathbf{E};x_p^2\\mathbf{E};\\cdots;x_p^N\\mathbf{E};]+\\mathbf{E}_{pos}, \u0026\\mathbf{E}\\in\\mathbb{R}^{(P^2\\cdot C)\\times D},\\mathbf{E}_{pos}\\in\\mathbb{R}^{(N+1)\\times D}\\\\ z_{\\ell}'\u0026=\\mathrm{MultiHeadAttention}(\\mathrm{LayerNorm}(z_{\\ell-1}))+z_{\\ell-1},\u0026\\ell=1,\\dots,L\\\\ z_{\\ell} \u0026= \\mathrm{MLP}(\\mathrm{LayerNorm}(z_{\\ell}'))+z_{\\ell}',\u0026\\ell=1,\\dots,L\\\\ y\u0026=\\mathrm{LayerNorm}(z_L^0) \\end{aligned} $$作者分析认为，相比于 CNN 架构，Transformer 架构缺乏 inductive bias, 这是因为每个 patch 对其周围的 2D 临近信息用的非常少，ViT 架构必须从零开始学习这些位置以及空间关系。\n作者还介绍了混合架构，即 patch embedding 由一个 linear layer 替换为一个 CNN 架构\n在 finetune 时，作者移除了 pre-trained 的 prediction head, 然后加入了一个 $D\\times K$ 的 feed forward layer, 其中 $K$ 是分类的类别数。\nExperiments 作者使用的数据集如下所示\nDataset classes images ImageNet 1K 1.3M ImageNet-21K 21K 14M JFT 18K 303M 作者在 ImageNet, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102 这些数据集上进行评测。\n模型配置如下表所示\nModel Layers $D$ $D_{FFN}$ # heads # params ViT-Base 12 768 3072 12 86M ViT-Large 24 1024 4096 16 307M ViT-Huge 32 1280 5120 16 632M 为了简便，作者在后续实验过程中，对模型名字进行了改写，比如 ViT-L/16 就代表了 ViT-Large model, 其对应的 patch size 为 16.\nViT 的实验结果如下图所示\n可以看到相比于 CNN 架构，ViT 所需要的训练时间更少，且效果更好。\n为了验证 ViT 模型具有更高的训练效率，作者在不同的数据集上进行的实验，结果如下图所示\n可以看到，当数据集比较小时，大部分 ViT 模型表现都比 ResNet 差，但是当数据集规模增加之后，ViT 的表现逐渐超过了 ResNet. 并且，随着数据集规模的增加，大模型的表现也比小模型好。\n作者进一步在四个不同大小的随机数据集上进行了验证，结果如下图所示\n实验结果显示，inductive bias 在小规模数据集上比较有效，但是当数据集规模上去之后，直接从数据集中学习相关的 pattern 更加高效和有效。\n由于不同的模型参数可能不太一致，为了更加公平地对比，作者使用算力来进行对比，结果如下图所示\n实验结果显示，ViT 的训练效率比 ResNet 更高，其次，hybrid model 在算力比较小的时候，其表现比 ViT 更好，但是当算力提升之后，两者表现差不多。最后，ViT 随着算力的提升，还有进一步提升的空间。\n作者还对 ViT 的 attention 进行了进一步的分析，结果有以下几点：\nrow-column structure 很明显，同一行同一列的 patches 有比较相似的 embedding ViT 学习到的 position embedding 对于比较近的 patches 其相似度也更高 对于 attention 来说，某些 heads 在 early layer 就已经开始关注全局信息，并且随着 layer 的提升，attention distance 也在提升，说明模型关注全局信息的能力逐渐增强 作者尝试了对 Transformer 进行 scaling up, 结果如下图所示\n实验结果显示，使用更多的 layer 可以有效提高模型的表现，但是当 layer 数大于 16 之后，其提升有限。另一方面，提高宽度对模型表现影响不大。作者发现，降低 patch size 也可以提高模型的表现。\n作者还对比了不同类型的 position encoding, 结果发现，使用/不使用 position encoding 对模型表现影响非常大，但是不同的 position encoding 总体来说表现差别不是很大。\nConclusion 作者在本文中提出了 ViT, 一个基于 Transformer 架构的图像识别模型，作者通过在大规模数据集上进行训练验证了 ViT 架构的有效性。\n作者发现目前还存在如下挑战：\n如何使用自监督的方式进行训练，比如类似 BERT 的 masked patch prediction 将 ViT 应用于其他的视觉任务，比如检测和分割 进一步 scaling ViT References openreview ","date":"2025-12-04T11:00:44+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-vit/","title":"Notes on ViT"},{"content":"Introduction 作者首先提到已有的 vision foundation model (VFM) 存在两个问题：\n不能处理动态分辨率图片输入，尽管我们可以使用 Bilinear interpolation 和 multi-resolution training 等方法，但是模型对于动态分辨率图片输入处理能力仍然不足 VFM 和 LLM 之间存在 representation gap 针对这两个问题，作者提出了 CoMP, 一个 continual pre-training pipeline, CoMP 包含两个模块：\nC-RoPE, 一个针对 VFM 的 continual RoPE, 用于帮助 VFM 处理动态分辨率图片输入 Alignment Loss, 用于对齐 VFM 和 LLM 的 representation Method CoMP 整体的架构图如下所示\n可以看到，CoMP 本质上就是一个多模态大模型，知识我们训练的目标为 VFM\nC-RoPE C-RopE 的核心思想是结合绝对位置编码以及相对位置编码来使得 pre-trained ViT 可以接受任意精度图片输入，对于输入图片 $X_V\\in\\mathbb{R}^{H\\times W}$, 首先经过 patchify 得到 $N=HW/P^2$ 个 patch, 这里 $P$ 是 patch size, 每个 patch 大小为 $x_p\\in\\mathbb{R}^{N\\times (P^2\\cdot C)}$, $C$ 是 channels. 然后 VFM 每一个 layer 的计算过程为\n$$ \\begin{aligned} z_0 \u0026= [x_p^1E;\\dots;x_p^NE] + \\mathrm{Int}(E_{pos})\\\\ q_i,k_i,v_i \u0026= \\mathrm{Proj}_q(z_i), \\mathrm{Proj}_k(z_i), \\mathrm{Proj}_v(z_i)\\\\ y_i \u0026= z_i + \\mathrm{Proj}_o(\\mathrm{Softmax}\\left((Rq_i)^T(Rk_i) / D_v\\right)v_i)\\\\ z_{i+1} \u0026= y_i + \\mathrm{FFN}(y_i) \\end{aligned} $$这里 $E\\in\\mathbb{E}^{P^2\\cdot C\\times D_v}$, $E_{pos}\\in\\mathbb{R}^{N\\times D_V}$ 分别是 patch embedding 和 learnable position embedding, $\\mathrm{Int}(\\cdot)$ 是 bilinear interpolation.\n相比于 ViT, C-RoPE 做出了两点改动：\n使用了 Interpolation 来支持动态分辨率图片输入 使用了 RoPE 来高效利用位置信息 Text-supervised Generative Pre-training 作者还是用了 LLM 的 cross-entropy loss 来进行对齐，text decoding loss 定义为\n$$ \\mathcal{L}_{dec} = -\\frac1T\\sum_{i=V+1}^{V+T}\\log P(X_i\\mid X_{","date":"2025-12-04T10:58:30+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-comp/","title":"Notes on CoMP"},{"content":"","date":"2025-12-02T18:21:54+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseek-r1/","title":"Notes on DeepSeek-R1"},{"content":"Introduction 作者首先提到如何提高模型的训练效率以及 inference efficiency 是两个尚未解决的问题。\n基于这两个问题，作者在本文中提出了 DeepSeek-V2，一个开源的 MoE 模型，DeepSeek-V2 的亮点在于训练和推理都非常高效。最终 DeepSeeK-V2 包含 236B 总参数，激活参数为 21B, 上下文长度为 128K. 作者还开源了 DeepSeek-V2-Lite, 一个 15.7B-A2.4B 的 MoE 模型，用于学术研究。\nDeepSeek-V2 主要改进点为：\n基于 DeepSeekMoE, 使用了 MoE 架构 使用了 MLA 压缩 KV cache, 大幅度提高推理效率 DeepSeek-V2 预训练使用了 8.1T tokens, 相比于 DeepSeek-LLM, 预训练数据主要增加了中文数据以及提高了数据的质量。\n接下来，作者收集了 1.5M 对话数据来进行 SFT, 最终作者基于 DeepSeek Math 提出的 GRPO 来进行对齐。\nArchitecture DeepSeek-V2 的模型架构如下\n模型基于 DeepSeekMoE 开发得到，相比于 DeepSeekMoE, DeepSeek-V2 主要是使用了 MLA\nMLA 这部分介绍见 MLA\nDeepSeekMoE Architecture 关于架构的介绍见 DeepSeekMoE\nDevice-Limited Routing 由于 DeepSeek-MoE 使用了细粒度的专家，因此专家会分布在更多的设备（GPU）上，计算时，基于 routing 的 expert 所在设备，会产生不同大小的通信开销。为了降低通信开销，作者构建了 device-limited routing mechanism. 具体的做法就是，在 Routing 之前，先基于 experts 的 affinity score 挑选 $M$ 个设备，然后基于这 $M$ 个设备的专家挑选 top-K 专家进行计算。\n作者通过实验发现，当 $M\\geq3$ 时，device-limited routing 可以和标准的 top-K routing 表现差不多。\nAuxiliary Loss for Load Balance 作者使用了三个 loss 来实现负载均衡。其中，expert level 和 device level 的 load balancing loss 与 DeepSeekMoE 相同。第三个 loss 是 communication balance loss, 这个 loss 的目的是让每个设备的通信开销保持平衡。损失函数的表达式如下所示\n$$ \\mathcal{L}_{communication} = \\alpha\\sum_{i=1}^D f_iP_i $$其中 $\\alpha$ 是超参数，$D$ 是 expert group 的个数。\n$$ f_i = \\frac{D}{MT}\\sum_{t=1}^T\\mathbb{1}(\\text{Token }t \\text{ is sent Device }i),\\quad P_i = \\sum_{j\\in\\mathcal{E}_i}P_j $$device limited routing 让每个 device 发送至多 $MT$ 个 hidden states 到其他设备上。而 communication balancing loss 则让每个设备最多从其他设备接收 $MT$ 个 hidden states.\nToken-Dropping Strategy 尽管前面已经增加了 load balance loss, 但毕竟不是硬约束。因此，作者就从硬件层面提出了 Token dropping 策略，来提高训练效率。核心思想就是，在训练时，主动丢弃部分 token, 强制让各个设备的计算量不会超过额度限制，进而减少资源浪费。\n具体做法就是，在训练之前，先将每个设备的 capacity factor 设置为 1 （定义见 Switch Transformer）. 然后按照 affinity score 来丢弃一些分数比较低的 token, 直到该设备的 token 数量刚到达到 capacity。为了避免过度学习导致模型表现较差，对于 $10\\%$ 的训练数据，作者不执行 token dropping 策略。\n最终，在 inference 时，可以根据需求来决定是否丢弃 token, 比如在 low latency 场景，我们可以丢弃低价值的 token, 在高精度场景，我们就可以保留所有的 token.由于在训练阶段已经才去过 token dropping 策略，因此在推理时不管是丢弃还是全部保留模型都能比较好的适应。\nPre-training Data 预训练数据与 DeepSeek-LLM 基本上差不多，作者针对中文数据，数据质量进行了改进。最终预训练数据包括 8.1T token, 其中中文数据比英文数据多 $12\\%$.\ntokenizer 与 DeepSeek-LLM 一致。\nModel Configuration 模型配置如下表所示\nModel DeepSeek-V2 DeepSeek-V2-Lite Date 2024-5 2024-5 # Total Parameters 236B 15.7B # Activated Parameters 21B 2.4B # tokens 8.1T 5.7T # Dense Layers 1 1 # MoE Layers 60 26 Hidden Dim 5120 2048 Dense Intermediate Dim 12288 10944 MoE Intermediate Dim 1536 1408 Attention MLA MLA # Attention Heads 128 16 # Key-Value Heads 128 16 # Routed Experts 160 64 # Experts Active Per Token 6 6 # Shared Experts 2 2 这里比较特殊的一点在于，模型在第一层使用了 MoE layer, 这个做法的原因在后面的 olmoe 里有提到，核心思想是 early layer 特别是第一层 layer 收敛比较慢。\nMLA 的配置如下 (DeepSeek-V2)\nfield value $d_c$ 512 $d_c'$ 1536 $n_h$ 16 $d_h$ 128 $d_h^R$ 64 Training Recipe 训练的配置也与 DeepSeek-LLM 差不多，对于 MoE，作者使用了 PP 将不同的 layers 分配在不同的 device 上，然后 MoE 的 experts 被分配在 8 个 device 上 ($D=8$), 对于 device-limited routing, 每个 token 发送到至多 3 个 device, 也就是 $M=3$.\nInfra 在 infra 上，DeepSeek-V2 也是用了 HAI-LLM 框架进行训练。这里面使用了 16-way zero-bubble PP, 8-way EP, ZeRO-1 DP.\n由于 DeepSeek-V2 的激活参数比较少，因此，作者没有使用 TP, 进而降低通信开销。作者还将 shared experts 的计算与 expert all-to-all 通信进行重叠来提高计算效率。作者还使用了 kernel fusion 和 flash attention 2 来加速训练。\nLong Context 在预训练阶段结束之后，作者使用了 YARN 来将模型的上下文从 4K 扩展到 128K. 超参数设置为\nParameter Value $s$ 40 $\\alpha$ 1 $\\beta$ 32 target context length 160K scaling factor $\\sqrt{t} = 0.0707\\ln s + 1$ 作者在 32K 的上下文下额外训练了 1000 步，然后在推理阶段通过 YaRN 将模型的上下文长度扩展到 128K.\nEvaluation 作者对比了 DeepSeek-LLM, Qwen1.5, Mixtral MoE 以及 LLaMA 3, 实验结果如下\nEfficiency 作者对比了以下 DeepSeek-MoE 和 DeepSeek-LLM 的训练效率，结果发现，对于 1T 的 token, DeepSeek-LLM 需要 300.6K GPU hours, 而 DeepSeek-V2 仅需要 127.8K GPU hours. 也就是说，DeepSeeK-V2 节省了 $42.5\\%$ 的训练成本\n在推理时，作者首先将模型的精度转换为 FP8，然后作者进一步对模型进行 KV cache quantization 来进一步压缩每个 token 的 KV cache 到 6bits. 最终，DeepSeek-V2 的 throughtput 为 50K tokens/s.\nPost-training post-training 分为 SFT 和 RL 两个阶段。\nSFT 在 SFT 阶段，作者构建了 1.5M 样本，包括 1.2M 有帮助性的样本和 0.3M 安全性相关的样本。模型训练了 2 个 epoch, 学习率为 $5e-6$.\nRL 作者使用了 GRPO 算法来进一步对齐模型的表现。\n作者通过实验发现，在 reasoning data, 如 code 和 math 相关数据上进行训练时，可以有效提高模型的表现。因此作者将 RL 的训练分为两个阶段，第一个阶段用于提高模型的 reasoning 能力，第二个阶段用于对齐人类偏好。\n在第一个阶段，作者首先训练了一个针对 code 和 Math 的 reward model $\\mathrm{RM}_{\\mathrm{reasoning}}$, 然后基于这个 reward model 来训练 policy model:\n$$ r_i=\\mathrm{RM}_{\\mathrm{reasoning}}(o_i) $$在第二阶段，作者使用了一个 Multi-reward 框架，包括一个 helpful reward model $\\mathrm{RM}_{\\mathrm{helpful}}$, 一个 safety reward model $\\mathrm{RM}_{\\mathrm{safety}}$ 和一个 rule-based reward model $\\mathrm{RM}_{\\mathrm{rule}}$, 最终的 reward 为\n$$ r_i = c_1\\mathrm{RM}_{\\mathrm{helpful}}+c_2\\mathrm{RM}_{\\mathrm{safety}}+c_3\\mathrm{RM}_{\\mathrm{rule}} $$训练时，reward model 由 SFT model 初始化得到，然后基于 point-wise 或者 pair-wise loss 进行训练。\nEvaluation chat 版本的模型评估结果如下所示\nDiscussion 作者讨论了三点发现：\nSFT data 数量。已有工作认为进需要 10K 左右的样本就可以进行 SFT，但是作者发现当数据量小于 10K 时，模型在 IFEval benchmark 上的表现大幅度下降。作者认为，这是由于数据过少导致模型很难掌握特定的技能。因此，作者认为足够的数据以及数据质量都很重要，特别是写作类任务和 open-ended QA 类任务。 alignment tax. 作者发现通过 human preference alignment, 模型在 open-ended generation benchmark 上的保险有了很大提升。与 RLHF 一样，作者也发现了 alignment 之后模型在一些 benchmark 上表现也会下降。作者通过改进解决了这个问题，作者认为如何在不损失模型表现的情况下实现对齐是一个值得探究的方向。 online RL. 作者发现 Online RL 比 offline RL 的表现更好。作者认为如何根据不同的任务来选取 offline RL 和 online RL 也是一个值得探究的问题。 Conclusion 在本文中，作者提出了 DeepSeek-V2, 一个基于 MoE 架构的大语言模型系列，模型的上下文为 128K. 作者基于 DeepSeek-MoE, 提出了 MLA 来提高模型的 inference 效率，并大幅度降低了训练的成本。\n作者介绍了几点未来工作：\n进一步 scaling up MoE 模型，降低模型的训练以及推理成本 进一步对齐模型和人类的价值观，然后最小化人类监督信号 扩展模型到多模态版本 Reference Arxiv HuggingFace ","date":"2025-12-02T18:21:54+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseek-v2/","title":"Notes on DeepSeek-V2"},{"content":"DeepSeek 在 2024 年 5 月提出了 multi-head latent attention (MLA), 用于提高 attention 的 Inference 效率\nIntroduction 传统的 multi head attention (MHA) 虽然效果好，但是在 inference 时，其 KV cache 会变成瓶颈，影响推理效率。为了解决这个问题，已有的工作如 MQA 和 GQA 通过共享权重来减少 KV cache 内存占用，但是结果发现模型的表现也会降低。\n为了解决这个问题，作者提出了 multi-head latent attention (MLA), 来压缩 KV cache.\nRelated Work MHA 令 $d$ 为 hidden size, $n_h$ 为 attention heads 的个数，$\\ell$ 为 transformer layer 的层数，$d_h$ 为每个 head 的 dimension, $h_t\\in\\mathbb{R}^d$ 为 attention layer 中第 $t$ 个 token 对应的 hidden states。对于标准的 MHA, 我们首先计算 Q, K, V 如下：\n$$ q_t=W^{Q}h_t,\\quad k_t=W^Kh_t,\\quad v_t = W^Vh_t $$其中，$W^Q,W^K,W^V\\in\\mathbb{R}^{d_hn_h\\times d}$ 分别为 query, key, value projection layer 的权重。接下来 MHA 的计算方式如下\n$$ \\begin{aligned} o_{t,i} \u0026= \\sum_{j=1}^t\\mathrm{softmax}_j\\left(\\frac{q_{t,i}^Tk_{j,i}}{\\sqrt{d_h}}\\right)v_{j,i},\\\\ u_t\u0026= W^O[o_{t,1};o_{t,2};\\dots,;o_{t,n_h}] \\end{aligned} $$其中 $q_t=[q_{t,1};q_{t,2};\\dots,;q_{t,n_h}]$, $k_t=[k_{t,1};k_{t,2};\\dots,;k_{t,n_h}]$, $v_t=[v_{t,1};v_{t,2};\\dots,;v_{t,n_h}]$. $W^O\\in\\mathbb{R}^{d\\times d_hn_h}$ 为 output projection 的权重。在 inference 阶段，每个 token 需要缓存其 key 以及 value 对应的值，从而每个 token 的 kv cache 占用为 $2n_hd_h\\ell$. 当序列长度过大时，KV cache 会影响整体的 inference efficiency.\nMQA \u0026amp; GQA MQA 通过在所有的 heads 中共享 key 和 value 来实现降低 kv cache 的作用，在 MQA 中，$W^K, W^V\\in\\mathbb{R}^{d_h\\times d}$, 在计算时，对应的 $k_t$ 和 $v_t$ 通过广播机制参与 attention 的计算。此时，KV cache 占用为 MHA 的 $1/n_h$, 即 $2d_h\\ell$.\n但是，MQA 的问题是表达能力太弱（表现差），因此后续 GQA 进行了改进，GQA 在 MQA 和 MHA 之间进行了权衡，即将 heads 分为若干个 group, 每个 group 中共享 key 和 value, 即 $W^K, W^V\\in\\mathbb{R}^{n_gd_h\\times d}$, 这里 $n_g$ 是 group 个数，在计算 attention 时，key 和 value 在 group 内部共享，此时，GQA 的 KV cache 占用是 MQA 的 $n_g$ 倍，即 $2n_gd_h\\ell$.\n这部分具体介绍见 MQA 和 GQA.\nMLA MLA 的架构图如下所示\nMLA 使用 low-rank joint compression 来压缩 key 以及 value 的 KV cache:\n$$ c_t^{KV} = W^{DKV}h_t,\\quad k_t^C = W^{UK}c_t^{KV}, v_t^C = W^{UV}c_t^{KV} $$这里 $c_t^{KV}\\in\\mathbb{R}^{d_c}$ 为 key 以及 value 压缩后的 latent vector. $d_c\u003c","date":"2025-12-02T18:21:54+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-mla/","title":"Notes on MLA"},{"content":"Introduction 已有的 MoE 模型往往都会使用 load balancing loss 来避免 Imbalanced routing, 但是加入这个额外的损失之后，模型训练的梯度也会受到影响。\nDeepSeek 基于这个问题提出了 Loss-Free Balancing, 该方法不引入额外的 loss item, 而是在 routing 的结果上加入一个 bias item, bias item 可以根据 expert load 来动态更新进而实现 load balancing.\nMethod MoE definition\n$$ \\begin{aligned} h_t \u0026= u_t + \\sum_{i=1}^Ng_{i,t}\\mathrm{FFN}_i(u_t),\\\\ g_{i,t} \u0026= \\begin{cases} s_{i,t}, \u0026 s_{i,t}\\in\\mathrm{Topk}(\\{s_{i,j}\\mid 1\\leq j\\leq N\\}, K)\\\\ 0, \u0026\\text{otherwise} \\end{cases}\\\\ s_{i,t} \u0026= G(u_t^Te_i), \\end{aligned} $$其中 $G$ 是 gating function, $e_i$ 是第 $i$ 个专家对应 gating function 的权重\nload balancing loss 有两个作用：\n避免 routing collapse, 即模型只选择固定的少数专家完成任务 减少通信开销 但是引入 load balancing loss 会对 LLM 的训练产生影响，为了避免对模型性能造成影响，我们需要小心设置 load balancing loss 的权重，实验结果如下图所示\n为了解决这个问题，作者提出了 Loss-Free Balancing, 具体做法就是在每个专家的 gating score 上加入一个 bias term, 然后再决定对应的专家：\n$$ g_{i,t} = \\begin{cases} s_{i,t}, \u0026 s_{i,t}+b_i\\in\\mathrm{Topk}(\\{s_{i,j}+b_j\\mid 1\\leq j\\leq N\\}, K)\\\\ 0, \u0026\\text{otherwise} \\end{cases} $$注意这里的 bias item 仅影响 top-K 操作，其对最终的输出没有影响。\n为了实现负载均衡，作者根据上一个 batch 的 expert load 情况来调整 bias item, 如果某一个专家的 load 太大，则对应的 bias item 会变小。 其算法实现过程如下：\n作者对比不同的负载均衡算法如下表所示\nLoad Balancing Methods Balanced Expert Load Interference Gradients Future Token Leakage Loss-Controlled (strong auxiliary loss) balanced strong no leakage Loss-Controlled (weak auxiliary loss) imbalanced weak no leakage Expert Choice balanced none with leakage Loss-Free (Ours) balanced none no leakage Experiments 作者使用了 DeepSeekMoE 来进行实验，作者使用了 sigmoid function 作为 gating function, 因为作者发现 sigmoid function 效果比 softmax 效果更好。\n作者提出了 maximal violation (MaxVio) 来量化一个 MoE layer 的负载均衡程度\n$$ \\mathrm{MaxVio} = \\frac{\\max_i\\mathrm{Load}_i-\\overline{\\mathrm{Load}_i}}{\\overline{\\mathrm{Load}_i}} $$其中 $\\mathrm{Load}_i$ 代表了分配给第 $i$ 个专家的 token 个数，$\\overline{\\mathrm{Load}_i}$ 代表了理想情况下的负载均衡。\n实验结果如下图所示\nModel Size Load Balancing Methods Validation Perplexity MaxVioglobal 1B Loss-Controlled 9.56 0.72 1B Loss-Free 9.50 0.04 3B Loss-Controlled 7.97 0.52 3B Loss-Free 7.92 0.04 实验结果显示，本文提出的 Loss-Free Balancing 效果更好，且负载均衡更高\n作者还展示了训练过程的负载情况如下\n可以看到，Loss-Free 的负载一直比 load balancing loss 效果更好\nAblation Study 作者首先探究了 bias term 更新速率对模型表现的影响，结果如下图所示\n结果显示，使用过大的 update rate 会影响最终的负载均衡，而较小的 update rate 收敛速率比较慢。因此作者在本文中使用了 $u=0.001$ 这个设置\n作者还对比了 baseline model 使用 softmax 个 sigmoid gating 两种方式，结果如下图所示\n实验结果显示，Sigmoid function 对于超参数更加 robust, 且表现也更好一些。\n作者还尝试了不同的 bias 更新方式，结果显示尽管不同的更新方式有可能会提高 load balance, 但是最终模型表现提升不大。\nConclusion 作者在本文中提出了 Loss-Free Balancing 策略，一个针对 MoE 负载均衡而不需要额外损失项的方法\nReferences Arxiv ","date":"2025-11-21T15:38:52+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-loss-free-balancing/","title":"Notes on Loss-free Balancing"},{"content":"Introduction 作者在本文中提出了 Mixtral 8x7B, 一个 MoE 模型，模型上下文为 32K. 作者还对模型进行 finetune 得到了 Mixtral 8x7B-Instruct, finetuning 包含 SFT 和 DPO 两个阶段。\nMethod 模型架构与 Mistral-7B 基本相同，参数如下表所示\nParameter Value dim 4096 n_layers 32 head_dim 128 hidden_dim 14336 n_heads 32 n_kv_heads 8 window_size 4096 context_len 32768 vocab_size 32000 num_experts 8 top_k_experts 2 MoE 的架构与 GShard 基本一致\nResults 作者探究了专家的 specialization, 结果有三点发现：\n不同专家对于不同 domain 的数据并没有出现 specialization 在 math domain 上，专家的分布有一个明显的区别。 连续的 token 往往会被分配到同一个专家上 Conclusion 作者在本文提出了 Mistral 8x7B, 一个 MoE 大语言模型\nReferences Arxiv ","date":"2025-11-01T15:32:30+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/mixstral-8x7b/","title":"Mixstral 8x7B"},{"content":"Introduction 作者主要强调 Mistral 7B 的表现超过了 LLaMA 2 7B 和 LLaMA 34B 的表现。\nMistral 7B 主要使用了 GQA 以及 SWA 两个方法来加速推理和减少内存占用，进而提高 batch size 和 throughput.\nMethod Mistral 7B 的模型配置如下表所示\nParameter Value dim 4096 n_layers 32 head_dim 128 hidden_dim 14336 n_heads 32 n_kv_heads 8 window_size 4096 context_len 8192 vocab_size 32000 SWA 的示意图如下所示\n对于第 $k$ 层，模型可以感知到 $W\\times k$ 的 tokens, 进而可以在提高训练效率的同时保持模型的表现。作者发现，在 flash attention 中使用 SWA 之后，模型的效率提升了 2 倍左右。\n并且使用了 SWA 之后，我们的 kv cache 也就随之固定了，因此我们可以使用一个 rolling buff cache, 其大小为 $W$, 对于第 $i$ 个 token, 我们将其保存在 cache 中的第 $i \\% M$ 个位置。\n作者还进一步将 sequence 分割为多个 chunk, 每个 chunk 的大小都是 window size $M$, 这样在计算 attention 的时候，对于当前的 chunk, 我们使用 self-attention, 对于 cache 中的 attention, 我们使用 SWA, 然后对于 past token, 由于这部分不在 sliding window 内因此不参与计算\nConclusion 作者在本文中提出了 Mistral 7B, 一个基于 GQA 和 SWA 的大语言模型。\nReferences Arxiv ","date":"2025-11-01T15:28:19+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/mixstral-7b/","title":"Mixstral 7B"},{"content":"Introduction 已有的模型如 DeepSeek-MoE, Mixtral MoE, Qwen1.5 等 MoE 模型基本只开源权重。也有一些开源的模型，比如 OpenMoE 等，但是开源信息不全。基于这个目的，作者提出了 olmoe 模型系列，包括 olmoe-7B-A1B 和 olmoe-7B-A1B-instruct 两个版本。\nMethod Pretraining 模型的架构如下图所示，MoE 架构与 dense 架构不同的地方在于 decoder layer 中的 FFN 被替换为了 MoE layer.\n模型的配置如下表所示\n训练的目标函数为\n$$ \\mathcal{L} = \\mathcal{L}_{CE} +\\alpha\\mathcal{L}_{LB} +\\beta\\mathcal{L}_{RZ} $$其中 $\\alpha,\\beta$ 为系数， $\\mathcal{L}_{CE}$, $\\mathcal{L}_{LB}$ 以及 $\\mathcal{L}_{RZ}$ 分别代表 cross-entropy loss, load balancing loss 以及 routing Z loss.\n预训练数据包括 DCLM 和 Dolma1.7 两个数据集的混合，作者将预训练数据集称为olmoe-mix. 数据集的配比如下\nPost-training 在 post-training 时，作者将训练分为 instruction tuning 和 preference tuning 两个阶段，在 instruction dataset 中，作者加入了更多的代码和数学数据来提高对应的能力。数据集如下表所示\nAblation Study MoE Settings MoE vs. Dense 作者对比了 MoE 模型和 dense 模型的训练效率，为了方便对比，作者使用 olmo-7B 和 olmo-1B 作为 baseline, 最终 olmoe 的总参数为 6.9B, 激活参数为 1.3B. 实验结果如下图所示\n实验结果发现，MoE 模型所需要的 token 或者 FLOPs 是 dense 模型的 $1/3$, 但是由于 MoE 模型需要额外的内存开销，因此从训练时间上来看，MoE 模型训练时间仅比 dense 模型快 $2$ 倍左右。\nExpert Granularity DeepSeekMoE 提出使用细粒度的专家来提供更多的组合可能性。作者探究了不同的粒度对模型表现的影响，结果如下图所示\n结果显示，当专家粒度从 8E-1A 扩展到 32E-4A 时，模型在 HellaSwag 上的表现提升了 $10\\%$, 但是进一步扩展到 64E-8A 时，模型的表现提升不到 $2\\%$, 这说明了无限制提升粒度对模型的提升越来越有限。在本文中，作者使用了 64 个专家。\nShared Experts DeepSeekMoE 提出使用共享专家来学习 common knowledge, 作者对这种方法进行了实验，结果如下图所示。\n可以看到，加入一个 shared expert 之后，模型的表现没有变化，作者认为减少 routed expert 之后，模型的组合可能性降低为原来的 $10\\%$ 左右。因此作者认为没有必要使用共享专家，因此作者在 olmoe 中没有采用共享专家这个方法。\nExpert Choice vs. Token Choice 作者探究了 routing 的策略，一个是 expert choice (EC), 另一种是 token choice (TC), 分别代表了每个 expert 选取固定的 token 数和每个 token 选取固定的 expert 数这两种情况。实验结果如下图所示\n可以看到，token choice 的表现明显更好。EC 虽然可以实现负载均衡。但是因为自回归模型在生成时是无法提前确定生成的 token 数的，因此 EC 很可能导致算力资源浪费或者是 token dropping. 在本文中，作者采用了 TC 这种策略。\nSparse Upcycling 作者还对比了从零开始训练 MoE 与基于 dense model upcycling 的方式训练 MoE，sparse upcycling 的相关工作有 MiniCPM, Qwen2 以及 Mixtral MoE.结果如下图所示\n结果发现，upcycling 确实可以提高训练效率，但是这种方法的缺陷在于：\nupcycling 受 dense model 的超参数限制 upcycling 的训练不是很稳定 因此在本文中作者没有采取 upcycling 的做法。\nLoad Balancing Loss 作者还探究 Load Balancing loss 对模型表现的影响，结果如下图所示\n可以看到，加入 load balancing loss 之后，模型的表现均超过了不加时的表现。\n作者进一步分析了不同专家在加/不加 load balancing loss 时的激活情况，结果如下图所示\n结果显示，load balancing loss 确实可以让不同专家的激活概率大致相当。\nRouter Z-loss ST-MoE 提出了 Router z-loss 来提高 MOE 训练的稳定性和表现。其表达式如下所示\n$$ \\mathcal{L}_{RZ}(x) = \\frac{1}{B}\\sum_{i=1}^B\\left(\\log\\sum_{j=1}^{N_E}\\exp(x_i^{(i)})\\right)^2 $$实验结果如下图所示\n可以看到，加入 router Z-loss 之后，模型训练的稳定性有所提升。因此在本文中作者使用了这个 loss.\nGeneral Pre-training Settings Initialization 作者探究了不同初始化策略对模型训练的影响，结果发现使用 truncate normal initialization 的训练稳定性更高\nQK-Norm 作者探究了 QK-norm 对模型训练的影响，结果发现 QK-norm 可以提高模型训练的稳定性，结果如下图所示\nAdamW Epsilon 作者发现，在 AdamW 优化器中，使用更小的 eps 可以提高模型的表现，因此作者将 eps 设置为 $1e-8$.\nAdaptation Settings 在 post-training 阶段，作者在三个方面进行了实验：\n是否加入 load balancing loss: 结论是不加，因为负载均衡在 pre-training 阶段已经实现了 是否使用 annealing: 结论是使用，因为效果更好 使用 DPO 还是 KTO, 结论是两种方法结果差不多 实验结果如下表所示\nLoad Balancing Precision Switch Transformer 中提出使用 float32 精度来进行 routing 的计算，作者通过实验发现，这一方法并不能提高模型训练的稳定性，因此作者没有采用这一策略。\nMoE Analysis Router Saturation 作者探究了训练过程中激活的专家和训练结束后激活的专家的匹配程度，结果如下图所示\n结果发现，训练 $1\\%$ 的数据之后，就有 $40\\%$ 的 routing 和训练完毕的 routing 一致，当训练 $40\\%$ 的数据之后，这个比例提升到了 $80\\%$.\n作者还发现，later layers 比 early layers 饱和更快，early layer, 特别是 layer 0, 饱和的非常慢。作者认为，这是 DeepSeekMoE 放弃在第一层使用 MoE layer 的原因，因为 load balancing loss 收敛更慢。\nExpert Co-activation 作者分析了 expert 之间的互相依赖程度，作者通过可视化发现，不同的 expert 之间 co-activation 的比例比较小，说明 expert redundancy 比较低\nDomain Specialization 作者还探究了不同 expert 对于不同 domain 的 specialization 程度，作者发现对于 specialized domain 的数据，expert 会出现一定程度的 specialization, 但是对于通用 domain 的数据，expert 的 specialization 程度比较低。这个结论与 Mixtral MoE 的结论不同，作者认为这个原因是 Mixtral MoE 使用了 upcycling 的方式，这会限制模型的表现。因此，作者进一步强调 MoE 从零开始训练是一个更好的训练方式。\nVocabulary Specialization 作者还探究了 vocabulary 中不同 token index 与激活专家之间的关系，结果发现 later layers 的 specialization 程度更高，这与 saturation 的趋势一致\nConclusion 作者在本文中提出了 olmoe, 一个全开源的 moe 大模型系列，作者详细介绍了针对 MoE 架构和通用架构的设计，为后来的模型架构设计提供了基础。\nReferences arxiv ","date":"2025-11-01T15:23:58+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-olmoe/","title":"Notes on olmoe"},{"content":"Introduction 作者认为，训练大规模的模型存在如下问题：\n缺乏有效的 Model parallelism 算法 随着设备数的增加，训练时间与 model size 呈现超线性增长的关系 (tensorflow) 在 nodes 数增多时，构建所需要的时间也大幅度增长 在多个设备上 partition model 比较困难 作者在本文中构建了一个基于 sparse MoE 架构的 600B 模型。为了解决这些问题，作者作出了如下贡献：\n作者提出了基于 MoE 架构的模型，来减少计算和通信开销 作者提出了 Gshard, 来自动化实现并行 作者提出了 SPMD (single program multiple data) 来减少计算表示和编译的难度 Method 架构如下图所示\n其中激活专家个数为 2 个，MoE layer 和 FFN layer 交替出现。\n作者对 GATE 函数进行了如下优化：\nexpert capacity. 每个 expert 都有一个处理 token 的上限值，超过该上限值之后 GATE 直接输出 0 进入下一层，假设 batch size 为 $N$, 专家个数为 $E$, 则该阈值定义为 $N/E$ group dispatching. 将 token 拆分为 $G$ 个 group, 每个 group 并行处理，每个 group 里每个专家处理的 token 个数上限为 $N/(GE)$. Auxiliary loss. 作者使用了 Load Balancing loss 来实现负载均衡 Random routing. 作者选取了 top-2 的专家，当第二个专家的权重太小是，作者直接忽略第二个专家，简化为选取 top-1 的专家 算法运行如下所示\nParallel Implementation 第一步是将算法转化为线性代数的方式，算法的代码如下所示\n1 2 3 4 5 6 7 gates = softmax(einsum(\u0026#34;GSM, ME-\u0026gt;GSE\u0026#34;, inputs, wg)) combine_weights, dispatch_mask = Top2Gating(gates) dispatched_expert_inputs = einsum(\u0026#34;GSEC, GSM-\u0026gt;EGCM\u0026#34;, dispatch_mask, reshaped_inputs) h = einsum(\u0026#34;EGCM, EMH-\u0026gt;EGCH\u0026#34;, dispatched_expert_inputs, wi) h = relu(h) expert_outputs = einsum(\u0026#34;EGCH, EHM-\u0026gt;GECM\u0026#34;, h, wo) outputs = einsum(\u0026#34;GSEC, GECM-\u0026gt;GSM\u0026#34;, combine_weights, expert_outputs) 第二步是通过 API 来实现并行执行\n1 2 3 4 5 6 7 8 9 10 # Partition inputs along group (G) dim. inputs = split(inputs, 0, D) # Replicate the gating weights wg = replicate(wg) gates = softmax(einsum(\u0026#34;GSM, ME-\u0026gt;GSE\u0026#34;, inputs, wg)) combine_weights, dispatch_mask = Top2Gating(gating_logits) dispatched_expert_inputs = einsum(\u0026#34;GSEC, GSM-\u0026gt;EGCM\u0026#34;, dispatch_mask, reshaped_inputs) # Partition dispatched inputs along expert (E) dim. dispatched_expert_inputs = split(dispatched_expert_inputs, 0, D) h = einsum(\u0026#34;EGCM, EMH-\u0026gt;EGCH\u0026#34;, dispatched_expert_inputs, wi) 第三部就是基于 compiler infra 来基于 sharding annotation 自动化分割一个 computation graph.\nExperiments 作者在机器翻译的任务上训练了若干模型，结果发现：\n层数更多的模型表现更好 提高 expert capacity 有效提高模型的表现 使用更多的 expert 可以在 high-resourced 任务上提高表现 dense 模型相比于 MoE 模型拥有更强的迁移能力 从训练效率上来看\n层数更多的模型的 sample efficiency 也更高 600B 的模型也可以在 4 天之内训练完毕 从内存使用效率上来看\n层数相同时，weight memory 以及 activation memory 不随专家个数增加而增加 专家个数比较少（128）时，模型可以达到 roofline performance 的 $70\\%$, 专家个数比较多（2048）时，模型依然可以达到 roofline performance 的 $48\\%$. Conclusion 作者探究了如何提高大规模模型的训练效率，作者提出了基于 MoE 架构的模型，然后作者还设计了 GShard, 一个自动化分割大规模模型的深度学习模块。实现结果发现，增加模型 size 可以提高模型的表现。作者还发现，SPMD 是一个更好的提高计算效率的方法。\nReference arxiv ","date":"2025-10-29T11:22:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/gshard/","title":"GShard"},{"content":"Introduction 已有的工作如 GShard 和 Switch Transformer 均验证了 MoE 模型的有效性，但是他们的问题在于训练不稳定。\n因此，在本文中，作者就提出了 ST-MoE 来解决这个问题。作者的主要贡献如下：\n探究了如何平衡模型的表现与训练稳定性 提出了 router Z-loss 来解决训练的不稳定性 探究了如何设定 MoE 模型 fine tuning 时的超参数 对 MoE 的性质进行了分析 Training Stability 作者发现，直接训练 MoE 模型会面临训练不稳定性的问题，实验结果如下图所示\n可以看到不同的实验，一次训练崩溃，一次训练正常。\n作者接下来探究了如何提高训练的稳定性，作者有三点发现：\n大多数方法都可以提高训练稳定性，但是也会是的模型表现更差 router Z-loss 可以在不损失模型表现的情况下提高训练的稳定性 multiplicative components 如 RMSNorm 等可以提高模型表现，但是训练稳定性更差 作者构建了一个 baseline 的 MoE 模型，总专家个数为 $N=32$, 激活专家个数为 $K=2$, Transformer block 按照 4 个 block 为一组，每组里包含一个 MoE layer. 为了避免初始化带来的误差，作者使用 6 个随机数种子进行训练。\nMultiplicative Components multiplicative components 指的是 GEGLU 和 RMSNorm 这些操作，其实验结果如下所示\nMethod Fraction Stable Quality Baseline 4/6 $-1.755 \\pm0.02$ Remove GEGLU 3/3 $-1.849 \\pm0.02$ Remove RMS Norm. Scale Param 3/3 $-2.020 \\pm0.06$ 结果发现，移除掉 multiplicative components 之后，模型训练变得稳定，但是效果也变差了。\nAdding Noise 接下来作者尝试了 Switch Transformer 中的 jitter noise 和 dropoout 方法，实验结果如下图所示\nMethod Fraction Stable Quality Baseline 4/6 $-1.755 ± 0.02$ Input jitter ($10^{-2}$) 3/3 $-1.777 \\pm 0.03$ Dropout (0.1) 3/3 $-1.822 \\pm0.11$ 实验结果显示，加入噪声对模型的表现存在负面影响。\nConstraining Activations 作者接下来分析了以下 Switch Transformer 中 router 存在的问题，作者发现尽管在 router 中使用 float32 可以提高训练稳定性，但是这还不够，因此作者使用了如下的 router Z-loss:\n$$ \\mathcal{L}_z(x) = \\frac1B\\sum_{i=1}^B\\left(\\log\\sum_{j=1}^N e^{x_j^{(i)}}\\right)^2 $$其中 $B, N$ 分别是 batch size 和专家个数，$x_j^{(i)}$ 代表了第 $j$ 个专家对 $i$ 个 token 的激活 logits. 通过增加这个 loss, 我们可以让 routing layer 的 logits 都在一个合理的范围内。\nMethod Fraction Stable Quality ($\\uparrow$) Baseline $4/6$ $-1.755 \\pm 0.02$ Update clipping (clip = 0.1) $3/3$ $-4.206 \\pm 0.17$ Router Z-Loss $3/3$ $-1.741 \\pm 0.02$ 可以看到，加入 router Z-loss 之后，模型的稳定性和表现都有了提升。\n因此，在本文中，作者使用的损失函数为\n$$ \\mathcal{L} = \\mathcal{L}_{CE} + \\alpha \\mathcal{L}_B + \\beta \\mathcal{L}_z $$其中 $\\mathcal{L}_{CE}, \\mathcal{L}_B$ 分别是 cross-entropy loss 和 Load Balancing loss, $\\alpha,\\beta$ 是对应 loss 的权重。\nNumerical Precision 接下来作者探究了数值精度对训练效率和稳定性的影响，作者发现使用低精度进行训练的优势有：\n通信开销更小 计算消耗更小 内存需求更小 但是低精度进行训练的问题是存在严重的 roundoff error, 其示意图如下所示\nFine-tuning 作者在本节探究了如何 fine-tune 一个 MoE 模型。\n作者首先通过实验给出了 MoE 模型容易过拟合的结论，实验结果如下图所示\n可以看到，在两个人呢误伤，MoE 模型相比于 dense 模型都更容易过拟合。\n接下来作者探究了超参数特别是 batch size 和 learning rate 对模型表现的影响，结果如下图所示\n可以看到，与 dense 模型相反，MoE 模型需要更大的 batch size 以及更小的 learning rate.\nDesign Sparse Models 首先作者分析了专家个数 $N$ 的设置，作者认为当专家个数超过 $256$ 之后，带来的收益就微乎其微了。并且，当专家个数增加之后，虽然算力没有变化，但是通信开销以及计算优化会变得更加困难。作者还认为，在 TPU 上，每个 core 应该进放置一个 expert\n作者还对 capacity factor 进行了实验，实验结果显示，提升 capacity factor 可以提高模型的表现。但是同时也会带来计算开销，从而让模型训练更慢，实验结果如下图所示\nModel Train CF Step Time (s) ($\\downarrow$) ST-MoE-L 1.25 $2.397$ ST-MoE-L 2.0 $2.447 (+7\\%)$ ST-MoE-32B 1.25 $4.244$ ST-MoE-32B 2.0 $4.819 (+14\\%)$ 最终结论如下：\n使用 top-2 routing, 即激活 $K=2$ 个专家，capacity factor 设置为 $1.25$, 每个 core 上最多放置一个专家 在评估时可以动态调整 capacity factor Tracing Tokens through the Model 在这一节里，作者分析了 MoE 模型专家的性质。\n首先，作者发现，每一层里，至少有一个专家对于 sentinel token 比较敏感。并且，encoder 里的专家 specialization 更强，比如某些专家会负责 punctuation, verbs, proper names 等\n接下来，作者发现，专家的 specialization 在 decoder 里几乎没有，作者认为这是因为 target token distribution 不同导致的，即：\ndecoding 是只有一小部分 token 被一个专家处理 decoding 过程中大部分 token 都是 sentinel token 因此，每个 group 相比于 encoder 来说通常只 cover 一小部分的 semantic space\nencoder 和 decoder 专家的 specialization 实验结果如下\nLayer 1 Layer 2 Layer 3 Layer 4 Layer 5 Layer 6 Uniform (32-experts) Encoder 2.2 1.8 1.6 1.7 1.7 1.2 3.5 Decoder 3.4 3.4 3.4 3.4 3.4 3.4 3.5 可以看到，decoder 的专家分布和均匀分布差不多，这说明 decoder 里的专家 specialization 更弱。\n作者还探究了不同专家对于不同语言的敏感程度，结果发现，专家对于不同语言并没有出现明显的 specialization 现象，作者分析原因认为，输入一般只含有一种语言，因此各个专家均需要去处理对应语言的 token, 这就导致了专家对于不同语种的数据处理基本上没有太大差别。\nConclusion 在本文中，作者提出了 ST-MoE 系列 MoE 大语言模型，作者通过实验优化了 MoE 模型训练的不稳定性问题。\nReferences Arxiv ","date":"2025-10-29T11:19:37+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/st-moe/","title":"ST-MoE"},{"content":"Introduction Kaplan scaling law 和 Chinchilla scaling law 已经给出了针对 dense transformer model 的 scaling law, 即在给定算力下如何找到最优的 model size 和 dataset size. 已有的工作如 T5 主要是扩展 dense model 来达到更高的表现, 但是他们的问题是对算力要求比较高。\n为了解决这个问题，作者尝试在不增加算力的情况下提升模型的参数量，为了达到这个目的，作者构建了基于 MoE 架构的模型 Switch Transformer.\n作者的主要贡献如下：\n提出了基于 MoE 架构的 Switch Transformer model 探究了针对 MoE 架构的 scaling law 将 MoE model 的能力蒸馏到 small dense model 里去 若干提升训练效率和稳定性的技巧 Architecture Switch Transformer 的架构如下图所示\nMoE MoE 的定义见 MoE tutorial, 我们假设有 $N$ 个专家，其中激活 $K$ 个专家。\n之前的工作认为我们只有在激活 $\u003e2$ 个专家时，模型才能够比较好的训练，但是在本文中，作者决定只使用 1 个 expert, 也就是 $K=1$. 作者将激活一个专家的 layer 称为 Switch layer.\n作者认为 Switch Layer 有三个优势：\nrouter computation 现在只需要将每个 token route 到 1 个 expert 每个专家的 capacity 更小，负载更加均衡 routing 的实现更简单，且通信开销也降低了 Efficient Sparse Routing 作者首先定义了expert capacity, 也就是每个 expert 处理的 token 数量，其定义如下\n$$ \\text{expert capacity} = \\left(\\frac{\\text{tokens per batch}}{\\text{number of experts}}\\right) * \\text{capacity factor} $$其示意图如下所示\n提升 capacity factor 可以减少 token overflow 的概率，但同时也会导致计算和内存的浪费。作者通过实验发现应该尽可能降低 dropped token 比例。\n为了平衡每个 expert 处理 token 的个数，作者设计了 load balancing loss, 见 Load Balancing loss 来要求每个 expert 处理的 token 数基本一致。\nParallelism 作者在本节介绍了针对 MoE 模型的并行策略，如下图所示\n这里，我们给定 notation 如下\nTerm Description $B$ Number of tokens in the batch. $N$ Number of total cores. $n$ Number of ways for data-parallelism sharding. $m$ Number of ways for model-parallelism sharding. $E$ Number of experts in Switch layers. $C$ Expert capacity, the batch size of each expert. 作者将一个 Logical mesh 分为两个维度，一个是 data-parallel sharding, 用 $n$ 表示，另一个是 model-parallel sharding, 用 $m$ 表示，从而总的 cores 数为 $N=n\\times m$.\nData Parallelism 对于 data prallelism 来说，每个 core 上的模型是一样的，而一个 batch 的数据被切分为了 $B/n$ 份，此时 $n=N$, $m=1$. 这种并行策略的好处是仅在 forward 和 backward 完成之后才会进行一次通信。\nModel Parallelism 对于 model parallelism, 每个 core 上的数据都是全部数据（通过拷贝得到），模型被 shard 到所有 core 上，此时 $n=1,m=N$, 这种情况下，每次 forward 或者 backward，都需要进行 $N$ 次通信（假设我们使用 pipeline parallelism）， 这种并行策略的好处是每个 core 上的计算压力小了，但是缺点是通信开销多了。\nModel and Data Parallelism 第三种总策略是同时进行 model parallelism 和 data parallelism, 这种情况下，每个 core 上保存 $B/n$ 的数据以及 shard 为 $m$ 份的 sharding weight.\nExpert and Data Parallelism 第四种策略是同时进行 expert parallelism 和 data parallelism, 作者在这里假设 $n=N$ 且 $E=n=N$, 也就是每个 core 上保留 $B/n$ 的数据, 然后还有一个对应的专家。\n首先，对于输入大小为 $[B,d]$ 的 tensor, 我们会进行拆分得到大小为 $[n, B/n, d]$ 的 tensor, 代表 $n$ 个设备上分别存储 $[B/n, d]$ 的数据，首先我们计算路由权重，得到\n$$ [n,B/n, d]\\times [d, E] \\to [n, B/n, E] $$权重的每个值 $[i,j,k]$ 代表了第 $i$ 个 core 上,第 $j$ 个 token 分配一个第 $k$ 个专家的概率，我们对其进行 softmax 与 top-K 操作之后，就得到了对应的专家（注意 Switch Transformer 中的 $K=1$, 我们通过 one-hot 编码将其输出转化为一个二进制矩阵，大小为 $[n,B/n, E, C]$, 这里的每个值 $[i,j,k,l]$ 代表了第 $i$ 个 core 上,第 $j$ 个 token 分配给第 $k$ 个专家第 $l$ 个 token, 接下来我们再基于输入 [n,B/n,d] 计算 core 到 expert 的数据，其大小为 [n,E,C,d], 计算方式为\n$$ \\mathrm{einsum}([n,B/n, d], [n,B/n,E,C], \\mathrm{dim}=[B/n]) $$里面的元素 [i,j,k,:] 代表了第 $i$ 个设备路由到第 $j$ 个专家的第 $k$ ($k","date":"2025-10-28T09:38:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/switch-transformer/","title":"Switch Transformer"},{"content":"Introduction 本文中关注的研究问题为：\n给定一个 FLOPs budget, 如何平衡 model size 和 dataset size 之间的关系？\n即，我们希望求解如下优化问题：\n$$ N_{opt}(C), D_{opt}(C) =\\arg\\min_{N,D,\\ \\mathrm{s.t.}\\ FLOPs(N,D)=C} L(N,D) $$作者通过训练 400 多个模型，构建了对应的 scaling law.\n已有工作如 Kaplan scaling law 已经发现模型参数和大语言模型表现之间的关系，一个结论就是计算最优并不代表达到最优的 loss. 在本文中，作者也有相同结论，但是作者认为大模型应该使用比 Kaplan scaling law 推荐的更多的 training token. 基于这个发现，作者训练了 Chinchilla, 一个 70B 的 LLM, Chinchilla 相比 Gopher 表现有了大幅度的提升。\nScaling Law Fix Model Size and Very Dataset Size 这个方法中，作者通过改变训练步数，来研究 FLOPs 与模型表现之间的关系，结果如下图所示\n通过对实验结果进行拟合，作者发现存在关系 $N_{opt}(C)\\propto C^a$ 以及 $D_{opt}(C)\\propto C^b$, 拟合的结果为 $a=b=0.5$.\nIsoFLOPS Profiles 这个方法中，作者使用了不同的模型大小以及算力来构建最优模型参数量与算力之间的关系。作者给定 9 个算力配置，然后选取不同参数量的模型，训练的 token 数由算力和模型参数量决定，实验结果如下图所示\n结果显示，不同大小的模型的表现 (loss) 随算力上升先下降后上升。因此给定算力，存在一个最优的 model size. 作者基于拟合出来的曲线得到了 Gopher 使用的算力配置下的最优 model size 和 training tokens. 同样的，作者得到 $a=0.49,b=0.51$.\nFitting a Parametric Loss Function 这个方法中，作者对 $L(N,D)$ 进行建模，作者使用了如下的公式\n$$ L(N,D) = E + \\frac{A}{N^{\\alpha}} + \\frac{B}{D^{\\beta}} $$第一项代表了建模的误差，第二项代表了数据集充分大损失与模型参数之间的关系，第三项代表了当模型充分训练时，损失与数据集大小之间的关系。\n为了求解 $(A,B,E,\\alpha,\\beta)$, 作者基于训练收集到的数据 $L(N_i,D_i)$, 通过 L-BFGS 算法来最小化 Huber loss 进行求解，结果得到 $(A,B,E,\\alpha,\\beta)=( 406.4, 410.7, 1.69, 0.34, 0.28)$.\n将结果带入带上面的表达式中，然后求出梯度为 0 的点，就得到\n$$ N_{opt}(C) = G\\left(\\frac C6\\right)^a, D_{opt}(C) = G^{-1}\\left(\\frac C6\\right)^b, \\text{ where }G=\\left(\\frac{\\alpha A}{\\beta B}\\right)^{1/(\\alpha+\\beta)}, a=\\frac{\\beta}{\\alpha+\\beta}, b=\\frac{\\alpha}{\\alpha+\\beta} $$带入数值之后就得到 $a=0.46$, $b=0.54$. 作者对结果可视化如下图所示，左图是拟合曲线的 Contour plot, 右图对左图的一个切片\nOptimal Model Scaling 作者将三种方法的结果以及 Kaplan scaling law 的结果总结放在下表中，作者假设 $N_{opt}(C)\\propto C^a$ 以及 $D_{opt}(C)\\propto C^b$\nApproach $a$ $b$ Kaplan 0.73 0.26 Approach 1 0.50 0.50 Approach 2 0.49 0.51 Approach 3 0.46 0.54 结果表明，三种方法的结论差不多：model size 和 dataset size 增长 debility 差不多。\n作者因此给出来的不同模型大小所需要的算力以及 token, 结果如下表所示\nParameters Approach 1 Approach 2 Approach 3 FLOPs Tokens FLOPs Tokens FLOPs Tokens 400 M 1.92e+19 8.0 B 1.84e+19 7.7 B 1.84e+19 7.7 B 1 B 1.21e+20 20.2 B 1.20e+20 20.0 B 1.20e+20 20.0 B 10 B 1.23e+22 205.1 B 1.32e+22 219.5 B 1.32e+22 219.5 B 67 B 5.76e+23 1.5 T 6.88e+23 1.7 T 6.88e+23 1.7 T 175 B 3.85e+24 3.7 T 4.54e+24 4.3 T 4.54e+24 4.3 T 280 B 9.90e+24 5.9 T 1.18e+25 7.1 T 1.18e+25 7.1 T 520 B 3.43e+25 11.0 T 4.19e+25 13.4 T 4.19e+25 13.4 T 1 T 1.27e+26 21.2 T 1.59e+26 26.5 T 1.59e+26 26.5 T 10 T 1.30e+28 216.2 T 1.75e+28 292.0 T 1.75e+28 292.0 T 作者基于发现的 scaling law, 对已有模型进行了探究，发现现有的大模型都存在 under-training 的现象，结果如下图所示\n实验结果显示，现有的大模型的 size 应该更小（或者需要更大的算力）。作者最终的结论就是，现有的比较小的模型，需要更多的算力才能达到更好的表现。\nChinchilla 基于上一节的发现，作者提出了 Chinchilla, 一个 70B 的模型，训练使用了 1.4T token. 训练的数据集为 MassiveText 的扩展版本，训练使用的优化器为 AdamW, tokenizer 为 SentencePiece.\n模型配置如下表所示\nModel Layers Number Heads Key/Value Size dmodel Max LR Batch Size Gopher 280B 80 128 128 16384 $4\\times 10^{-5}$ $3M\\to6M$ Chinchilla 70B 80 64 128 8192 $1\\times 10^{-5}$ $1.5M\\to3M$ Ablation Study learning rate schedule 作者还通过 ablation study 发现，cosine learning rate cycle length 应该和训练步数差不多，当 cycle length 太长时，模型表现会下降。\nOptimizer 作者对比了 Adam 和 AdamW 的表现，结果发现，AdamW 的表现优于 Adam.\nHigh Precision 训练时，作者使用了高精度也就是 float32 来保存梯度的状态，结果显示，不管是 Adam 还是 AdamW, 使用高精度都可以提高模型的表现\nComparison with Kaplan 作者还对比了 Chinchilla 和 Kaplan 的预测结果，如下图所示\n结果显示，基于 Chinchilla 预测得到的模型训练效果比 Kaplan 的更好。\nCurvature of the FLOPs-frontier 作者发现，FLOP-minimal loss frontier 存在 curvature, 也就是小模型和大模型预测出来的曲线是不一样的，作者将结果展示在下图中\n结果显示，从小模型拟合出来的结果比大模型拥有更高的算力使用效率，作者认为这是未来的一个研究方向。\nConclusion 本文中作者重新探究了针对 LLM 的 scaling law, 作者发现已有的大模型都存在 under-training 的现象，也就是说，模型需要更多的训练 token, 具体来讲，model size scaling 和 dataset scaling 应该处于同一水平。作者基于这个结论，提出了 Chinchilla, 一个 70B 的 LLM, 其表现超过了 280B 的 LLM.\nReferences arxiv ","date":"2025-10-22T14:39:23+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/chinchilla-scaling-law/","title":"Chinchilla Scaling Law"},{"content":"Introduction 作者首先就总结了本文的发现，主要有以下几点：\n损失与模型的 scale , 即除开 embedding 的模型参数，数据集大小以及算力强相关，与 model shape 比如 depth 或者 width 关系不大 scaling law 是非常光滑的，意味着 scaling law 是一个可预测的模型 overfitting 普遍存在，当参数和数据集大小同时增加时，模型的表现会增加，但是当其中一个量固定时，提升就比较小。并且当我们将模型参数提升 8 倍时，我们只需要将数据集的大小提升 5 倍就可以避免过拟合 训练的损失函数曲线与 model size 无关，因此我们可以预测给定大小模型的表现 模型在测试集和训练集上的表现高度相关，因此我们可以基于训练集的损失来预测模型的表现 大模型比小模型拥有更高的 sample efficiency, 即更小的训练步数就可以达到相同的表现 convergence 不能说明一切，我们可以通过 early-stopping 来提高算力使用效率，避免模型花费过多的算力在较小的提升上 最优的 batch size 与 loss 呈一个 power law 的关系 Notation Symbol Notation $L$ cross entropy loss $N$ non-embedding parametters $B$ batch size $S$ number of training steps $C\\approx 6NBS$ estimate of total training compute $D$ dataset size in tokens $B_{crit}$ critical batch size $C_{\\min}$ estimate of the minimum compute to reach a given value of loss $S_{\\min}$ estimate of the minimum steps to reach a given value of loss $\\alpha_X$ power-law exponents for the scaling law of loss $d$ hidden size of the model 为了简便，后续未经特殊说明，我们说模型参数均指的是 non-embedding 的参数\nScaling Law Overview 当数据集 $D$ 足够大时，损失与模型参数大小 $N$ 之间的关系为 $$ L(N) = \\left(\\frac{N_C}{N}\\right)^{\\alpha_N}, \\alpha_N\\sim 0.076, N_C\\sim 8.8\\times 10^{13} $$ 给定模型参数大小 $N$ , 损失与数据集大小 $D$ 之间的关系为 $$ L(D) = \\left(\\frac{D_C}{D}\\right)^{\\alpha_D}, \\alpha_D\\sim 0.095, D_C\\sim 5.4\\times 10^{13} $$ 给定足够大的数据集 $D$ 和最优模型大小 $N$ 时，损失与算力 $C$ 之间的关系为 $$ L(C_{\\min}) = \\left(\\frac{C_C^{\\min}}{C_{\\min}}\\right)^{\\alpha_D}, \\alpha_C^{\\min}\\sim 0.050, C_C^{\\min}\\sim 3.1\\times 10^{8} $$作者在不同大小的数据集，算力，模型大小下进行了测试，结果发现 scaling 与模型的 shape, transformer 的超参数之间的关系比较小。$\\alpha_N,\\alpha_D,\\alpha_C^{\\min}$ 等决定了当我们 scale up 数据及大小，模型大小和算力时损失的变化情况。比如当我们将模型参数提升至 2 倍时，模型的损失会降低至原来的 $0.95$.\n基于发现 1 和 2, 作者发现当我们将模型的 size 提升至原来的 2 倍时，模型的数据集大小应该提升至原来的 $1.67$ 倍，具体关系为 $D\\sim N^{0.74}$.\n作者使用了一个统一的公式来描述损失与数据及大小和模型参数大小之间的关系\n$$ L(N,D) = \\left[\\left(\\frac{N_c}{N}\\right)^{\\frac{\\alpha_N}{\\alpha_D}}+ \\frac{D_c}{D}\\right]^{\\alpha_D} $$ 当数据集充分大时，损失与模型参数大小以及更新步数 $S$ 的关系如下 $$ L(N, S) = \\left(\\frac{N_C}{N}\\right)^{\\alpha_N} +\\left(\\frac{S_C}{S_{\\min}(S)}\\right)^{\\alpha_S} $$这里 $S_C\\approx 2.1\\times 10^3$, $\\alpha_S\\approx 0.76$, $S_{\\min}(S)$ 是估计出来的最小优化步数\n最优的 batch size 与损失函数之间的关系如下 $$ B_{crit}(L) = \\frac{B_*}{L^{1/\\alpha_B}}, B_*\\sim 2*10^8 \\text{ tokens}, \\alpha_B\\sim 0.21 $$ 给定算力 $C$ 且无其他限制时，模型参数，数据及大小，batch size 和更新参数与算力之间的关系如下 $$ N\\propto C^{\\alpha_C^{\\min}/\\alpha_N}, B\\propto C^{\\alpha_C^{\\min}/\\alpha_B}, S\\propto C^{\\alpha_C^{\\min}/\\alpha_S}, D=BS $$其中\n$$ \\alpha_C^{\\min} = \\frac{1}{\\frac{1}{\\alpha_S}+\\frac{1}{\\alpha_B}+\\frac{1}{\\alpha_N}} $$实验的结果为 $N\\propto C_{\\min}^{0.73}$, $B\\propto C_{\\min}^{0.24}$ , $S\\propto C_{\\min}^{0.03}$. 也就是说，当我们提升算力时，提升模型的参数大小带来的收益是最高的。\nBackground 首先，transformer 的参数量通过计算可以得到\n$$ N\\approx 2dn(2d+d_{ff}) = 12nd^2 $$这里 $d$ 是 hidden size, $n$ 是 layer 个数，$d_{ff}$ 是 MLP 的 hidden size, 这里我们 假设 $d_{ff}=4d$. 计算时我们丢掉了 bias 以及 LayerNorm 的参数量。具体计算过程见 LLM parameter analysis\ntransformer 一次前向计算的 operations 数量大概为\n$$ C_{forward}\\approx 2N + 2nLd $$这里 $L$ 是输入的 token 长度。\n由于反向传播所需要的 FLOPs 是前向传播两倍，因此 transformer 的计算量为\n$$ C = C_{backward} + C_{forward} = 3C_{forward}\\approx 6N $$具体计算过程见 LLM FLOPs analysis。也就是说，对于参数量为 $N$ 的 transformer model, 每个 token 所需要的 FLOPs 为 $C\\approx 6N$\nEmpirical Results and Basic Power Laws Transformer Shape and Hyper-parameter Independence 作者基于 $N=12nd^2$, 在保持总参数量 $N$ 不变的情况下，分别调整 $n$, $d_{ff}$ 和 number of attention heads 的个数 （变化 $d$ 用于维持总参数量不变），结果如下图所示\n实验结果发现，损失对于 $d_{ff}/d$, $d/n$, $d/n_h$ 都比较 robust, 说明模型的损失对模型的 shape 依赖性比较低。\nNon-embedding Parameter Count 作者探究了以下 model size 对损失的影响，作者使用了不同的 $n$ 和 $d$, 然后训练得到的损失情况如下图所示\n作者发现，当包含 embedding parameter 时，损失不仅依赖于模型参数量，还依赖于 layer 层数 $n$, 但是当我们排除 embedding parameter 时，模型的损失便与 layer 层数 $n$ 关系不大。这个趋势可以用以下模型来表示\n$$ L(N) \\approx \\left(\\frac{N_c}{N}\\right)^{\\alpha_N} $$最终拟合的曲线如下图所示\nComparing to LSTMs and Universal Transformers 作者比较了 LSTM 和 Transformer 结构的损失，结果如下图所示\n可以看到，transformer 比 LSTM 拥有更强的学习能力， LSTM 架构对于 early context 表现比较好，但是随着 context 增加，LSTM 的表现逐渐弱于 transformer. 即 transformer 的长上下文能力强于 LSTM 架构。\nGeneralization Among Data Distributions 模型是在 WebText2 数据集上训练的，作者进一步在其他数据集上评估了以下模型的泛化性，结果如下图所示\n结果发现，模型在其他数据集上的泛化性很好。并且，模型的泛化性能仅与训练阶段的表现相关（validation loss），而与训练阶段（是否收敛）无关。\n作者还评估了 model depth 对模型泛化性的影响，结果如下图所示\n实验结果显示，model depth 对模型泛化性基本没有影响。\nPerformance with Data Size and Compute 作者探究了损失与 dataset size $D$ 之间的关系。作者固定一个模型，然后当 test loss 不再下降时停止训练，结果发现 test loss 与 dataset size $D$ 之间存在如下关系\n$$ L(D) \\approx \\left(\\frac{D_c}{D}\\right)^{\\alpha_D} $$拟合结果如下图所示\n接下来，基于前面计算的结果，我们有 $C\\approx 6ND=6NBS$, 这里 $B$ 是 batch size, $S$ 是训练步数。给定 $C$, 作者使用不同大小的模型进行训练，batch size $B$ 保持不懂，训练步数设置为 $S=C/6BS$,实验结果显示损失与算力 $C$ 之间满足如下关系\n$$ L(C) \\approx \\left(\\frac{C_c}{C}\\right)^{\\alpha_C} $$拟合结果如下图所示\n作者进一步探究了 sample efficiency 与 model size 之间的关系，实验结果如下图所示\n结果显示，随着 model size 增加，sample efficiency 也在增加\nCharting the Infinite Data Limit and Overfitting 作者在本节探讨了同时变化 $N$ 和 $D$ 对损失变化的影响。\nProposed Equation 作者基于三个原则进行建模：\n改变 vocabulary size 或者 tokenization 会 rescale loss 固定 $D$ 并且令 $N\\to\\infty$, 则最终损失应该接近 $L(D)$. 反之固定 $N$, 令 $D\\to\\infty$, 最终损失应该接近 $L(N)$ $L(N,D)$ 在 $D=\\infty$ 处应该是可解析的 基于以上三条原则，将模型选择为如下形式\n$$ L(N,D) = \\left[\\left(\\frac{N_c}{N}\\right)^{\\frac{\\alpha_N}{\\alpha_D}}+ \\frac{D_c}{D}\\right]^{\\alpha_D} $$作者基于不同配置进行训练，基于实验结果你和得到的参数如下\nParameter $\\alpha_N$ $\\alpha_D$ $N_c$ $D_c$ Value 0.076 0.103 $6.4\\times 10^{13}$ $1.8\\times 10^{13}$ 接下来，作者探究了模型的过拟合程度，作者定义如下 metric\n$$ \\delta L(N,D) := \\frac{L(N,D)}{L(N,\\infty)} - 1 $$带入 $L(N,D)$ 定义就得到\n$$ \\delta L(N,D) = \\left(1 + \\left(\\frac{N}{N_c}\\right)^{\\frac{\\alpha_N}{\\alpha_D}}\\frac{D_c}{D}\\right) - 1 $$通过测试不同的模型，作者发现 $\\delta L$ 的值在 $0.02$ 左右，将实验结果带入到上面的公式就得到\n$$ D \\geq (5\\times 10^3)N^{0.7379} $$也就是说对于参数量为 $N$ 的模型，需要 data size $D \\geq (5\\times 10^3)N^{0.7379}$ 才能避免过拟合。\nScaling Laws with Model Size and Training time 作者在本节构建了损失函数与 model size $N$ 以及训练时间的 scaling law\nAdjustment for Training at Critical Batch Size 已有结论说明，存在一个 critical batch size $B_{crit}$, 当 batch size 接近 $B_{crit}$ 时，增加 batch size 对计算效率影响比较小，但是当 batch size 大于 $B_{crit}$ 时，带来的提升比较小。另一方面，batch size 会影响梯度的噪声程度。因此，训练步数 $S$ 和处理的样本数 $E=BS$ 应该满足：\n$$ \\left(\\frac{S}{S_{\\min}}-1\\right)\\left(\\frac{E}{E_{\\min}}-1\\right) = 1 $$这里 $S_{\\min}$ 是达到损失 $L$ 所需要的最小训练步数，而 $E_{\\min}$ 是最小的训练样本数量。\n作者的实验结果如下\n作者将 critical batch size 定义为\n$$ B_{crit}(L) := \\frac{E_{\\min}}{S_{\\min}} $$使用 critical batch size 进行训练可以在计算效率和算力之间达到一个平衡。\n作者基于上面的实验结果探究了 critical batch size 和 model performance 之间的关系，实验结果如下图所示\n可以看到，critical batch size 与 model size 的关系不大，仅与损失 $L$ 有关。作者通过以下模型拟合 critical batch size:\n$$ B_{crit}(L) \\approx \\frac{B_*}{L^{1/\\alpha_B}} $$这里 $B_*\\approx 2\\times 10^8$, $\\alpha_B\\approx 0.21$.\n给定一个 target loss $L$, 当 batch size $B\u003e\u003e B_{crit}$ 时，作者定义最小训练步数为\n$$ S_{\\min}(S) := \\frac{S}{1+B_{crit}(L)/B} $$给定 target loss $L$ 和 model size $N$, 当 batch size $B\u003c\u003c B_{crit}$ 时，作者定义最小算力为\n$$ C_{\\min}(C) := \\frac{C}{1+B_{crit}(L)/B} $$Performance with Model Size and Compute 作者使用如下公式来探究损失与 model size 和 computer 之间的关系\n$$ L(N, S_{\\min}) = \\left(\\frac{N_C}{N}\\right)^{\\alpha_N} +\\left(\\frac{S_C}{S_{\\min}(S)}\\right)^{\\alpha_S} $$拟合结果如下表所示\nParameter $\\alpha_N$ $\\alpha_S$ $N_c$ $S_c$ Value 0.077 0.76 $6.5\\times 10^{13}$ $2.1\\times 10^{3}$ 基于这个拟合结果，作者得到了下图的结果\n作者还使用了不同的可视化方式，如下图所示\n实验结果显示，上面的公式拟合的很好。\nLower Bound on Early Stopping step 作者还探究了以下 early step 与模型大小以及数据集之间的关系，作者通过分析得到如下结果\n$$ S_{stop}(N,D) \\gtrsim \\frac{S_c}{[L(N,D)-L(N,\\infty)]^{1/\\alpha_S}} $$其中 $L(N,\\infty)$ 是在充分大数据集上的收敛损失。作者对实验结果进行了拟合，结果如下图所示\nOptimal Allocation of the Compute Budget 作者在本节探究了最优算力与 model size $N$ 和训练数据 $2B_{crit}S_{\\min}$ 之间的关系\nOptimal Performance and Allocations 作者首先基于\n$$ C_{\\min}(C) := \\frac{C}{1+B_{crit}(L)/B} $$绘制了如下曲线图\n作者发现，相比于 loss 与算力 $C$ 之间的关系，使用 $C_{\\min}$ 进行拟合效果更好。\n接下来，作者基于 $L(C_{\\min})$ 进一步探究了给定算力如何决定最优的 model size $N(C_{\\min})$. 其实验结果如下图所示\n实验结果显示，model size 和算力之间有如下关系\n$$ N(C_{\\min}) \\propto (C_{\\min})^{0.73} $$作者进一步探究了对于非最优模型大小与算力之间的关系，作者先构建了如下的关系\n$$ \\frac{C(N, N_{\\mathrm{eff}})}{C(N_{\\mathrm{eff}}, N_{\\mathrm{eff}})} = \\frac{N}{N_{\\mathrm{eff}}} \\left[ 1 + \\frac{\\alpha_S}{\\alpha_N} \\left( 1 - \\left( \\frac{N_{\\mathrm{eff}}}{N} \\right)^{\\alpha_N} \\right) \\right]^{-\\!1 / \\alpha_S}. $$对应的示意图为\n实现结果发现，大小为最优模型的 $0.6\\sim 2.2$ 倍只需要额外 $20\\%$ 的算力。作者强调，这个实验结果对于超大模型不一定适用。\n作者进一步推导了 $S_{\\min}$ 和 $C_{\\min}$ 之间的关系，由于 $C_{\\min}=6NB_{crit}S$, 且我们前面已经有 $B\\propto L^{-4.8}$ 和 $L\\propto C_{\\min}^{-0.05}$, 因此 我们有\n$$ B_{crit}\\propto L^{-4.8} \\propto (C_{\\min})^{-0.05\\times (-4.8)}\\propto (C_{\\min})^{0.24} $$以及\n$$ S_{\\min} \\propto \\frac{C_{\\min}}{6B_{crit}N(C_{\\min})} \\propto (C_{\\min})^{0.03} $$拟合的结果如下图所示\n因此，基于上面的结果，当我们增加算力时，我们的主要精力应该放在增加模型大小和提高 batch size 上，而训练步数基本可以保持不变。\nAnother way of Derivation 作者还给出了另一种建模 $L(C_{\\min})$ 的方式，即从 $L(N,S_{\\min})$ 中进行推导，作者将 $B_{crit}$ 和 $S_{\\min}$ 的表达式带入到 $L(N,S_{\\min})$ 然后求解最小值就得到\n$$ L(C_{\\min})= \\left(\\frac{C_C^{\\min}}{C_{\\min}}\\right)^{\\alpha_C^{\\min}} $$其中\n$$ \\alpha_C^{\\min} = \\frac{1}{\\frac{1}{\\alpha_S}+\\frac{1}{\\alpha_B}+\\frac{1}{\\alpha_N}} \\approx 0.054 $$这和前面的结果基本是吻合的，进一步进行推导得到\n$$ N(C_{\\min})\\propto (C_{\\min})^{\\alpha_C^{\\min}/\\alpha_N}\\approx (C_{\\min})^{0.71} $$这个结果也和上面的差不多。\nContradiction and a Conjecture 作者发现，尽管拟合的 scaling law 曲线非常好，但是由于自然语言不可能达到 zero entropy, 因此该曲线最终一定会失效。作者基于更大的模型进行了实验，结果发现，模型在某一点开始就比预测的损失曲线下降的更慢。作者认为这是因为 transformer 模型已经达到了 maximal performance 导致的。\n通过前面的分析，我们发现 $L(C_{\\min})$ 比 $L(D)$ 下降的快，因此两者必然在某一点相交。\n在前面的章节中，我们基于以下关系来决定数据集大小\n$$ D\\propto N^{0.74}\\propto (C_{\\min})^{0.74*0.73}\\propto (C_{\\min})^{0.54} $$这里我们利用了 $N(C_{\\min})$ 的结果\n另一方面，我们有\n$$ D(C_{\\min}) = \\frac{2C_{\\min}}{6N(C_{\\min})}\\propto (C_{\\min})^{0.26} $$可以看到，基于训练最优导出的数据集大小相比于拟合出来的数据集大小，实际上存在过拟合。\n作者进一步分析出了 $L(D(C_{\\min}))$ 和 $L(C_{\\min})$ 这两条曲线的交点，结果得到\n$$ C^*\\approx 10^4 \\text{ PF-Days}, N^*\\approx 10^{12}\\text{ parameters}, D^*\\approx 10^12\\text{ tokens}, L^*\\approx 1.7\\text{1.7nats/token} $$作者认为出现这种原因有以下几种情况：\n$L^*$ 给出了自然语言的 entropy 的一个估计，因此当模型充分大之后，模型可能已经获取到了数据中的所有知识 $L(C_{\\min})$ 可以作为数据集噪声的一个量化表现，其衡量了数据集的质量 Learning Rate Schedule 附录中，作者还探究了 learning rate 与损失之间的关系，作者使用了不同 learning rate schedule 对模型损失的影响，结果如下图所示\n实验结果显示，只要 learning rate 下降的不会太快，模型的表现基本上差不太多。\n作者基于实验结果得到了学习率和模型参数之间的关系如下\n$$ \\text{lr}(N)\\approx 0.003239 - 0.0001395\\log N $$也就是说，小模型用比较大的学习率，大模型用较小的学习率。\nConclusion 作者在本文中训练了大量不同配置的大模型，然后构建了损失（损失）与模型参数，数据及大小以及算力之间的关系。实验结果发现，损失与架构和优化参数之间的关系比较小，主要由模型参数量决定，更大的模型拥有更高的采样效率。\n作者认为，本文的局限在于损失函数不一定能够反应模型在其他语言任务上的表现。\nReferences arxiv ","date":"2025-10-22T14:10:52+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/kaplan-scaling-law/","title":"Kaplan Scaling Law"},{"content":"本文中，我们介绍如何计算基于 transformer 架构的 LLM 的 FLOPs, 计算完成之后，我们可以推导出算力 $C$ 与模型参数量 $N$，数据集大小 $D$ 之间的关系，即 $C\\approx 6ND$.\nBackground FLOPs FLOPs，floating point operations，表示浮点数运算次数，一般计算 FLOPs 仅考虑加法和乘法。\n假设 $A\\in\\mathbb{R}^{m\\times p}$, $B\\in\\mathbb{R}^{p\\times n}$, 则 计算 $C=AB$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp$ 次加法运算，因此总的 FLOPs 为 $2mnp$.\nAssumption 假设 $A\\in\\mathbb{R}^{m\\times p}$, $B\\in\\mathbb{R}^{p\\times n}$, $C\\in\\mathbb{R}^{m\\times n}$ 则 计算 $D=AB+C$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp+mn$ 次加法运算，因此总的 FLOPs 为 $2mnp+mn\\approx 2mnp$.\n基于上述结论，我们计算 FLOPs 时，忽略 element-wise 的操作，即我们做如下假设：\n忽略 normalization 中的小常数项运算 忽略 residual connection 和 bias term 的加法 忽略注意力的 MASK 和 softmax，因为这两者都是 element-wise operation. 使用 look-up 计算 embedding layer 注，基于以上假设，我们的结果与 Chinchilla Scaling law 的结果稍有不同，但结论不变。\nNotation Math Variable Code Variable Description $n$ num_hidden_layers Transformer block 个数 $\\vert V\\vert$ vocab_size 词表大小 $d$ hidden_size token embedding 的维度 $d_{ff}$ intermediate_size MLP 的中间层的维度 $h$ num_attention_heads query head 的个数 $s$ seq_len length of token sequence 注：为了避免混淆，我们使用 $n$ 来表示 decode layer 的个数。\nComputation 我们计算训练阶段的总 FLOPs, 记为 $C$, Kaplan scaling law 用 PF-days 作为单位，$1\\text{ PF-Days}=10^{15}\\times 24\\times 3600\\ FLOPs$. 训练阶段包括前向阶段 (forward pass) 和反向传播阶段 (backward pass). 因此\n$$ C = FLOPs(\\text{forward}) + FLOPs(\\text{backward}) $$Forward decoder-only transformer 的模型架构包含三个模块：\n1 层 embedding layer $n$ 层 decoder layer 1 层 lm head layer 因此模型总的 FLOPs 为\n$$ FLOPs(\\text{forward}) = FLOPs(\\text{embedding}) + n*FLOPs(\\mathrm{decode\\_layer})+FLOPs(\\mathrm{lm\\_head}) $$Embedding \u0026amp; Lm Head 首先，对于 embedding layer, embedding layer 本质上是一个 look up table, 计算过程中不涉及浮点数运算，因此 $\\boxed{FLOPs(\\text{embedding})=0}$.\n接下来，对于 lm_head, 这是一个 linear layer, 其权重大小为 $W\\in\\mathbb{R}^{d\\times |V|}$, 输入为 $x\\in\\mathbb{R}^{s\\times d}$, 因此 $\\boxed{FLOPs(\\mathrm{lm\\_head})=2sd|V|}$.\n因此，我们有\n$$ FLOPs(\\text{forward}) = n*FLOPs(\\mathrm{decode\\_layer})+ 2sd|V| $$Decode Layer 对于 decode_layer, 其又包含了四个模块：\npre-normalization attention post-normalization FFN pre-normalization 和 post-normalization 一般是一样的，因此\n$$ \\begin{aligned} FLOPs(\\mathrm{decode\\_layer}) \u0026= FLOPs(\\mathrm{pre\\_normoalization}) + FLOPs(\\mathrm{Attention}) + FLOPs(\\mathrm{post\\_normoalization}) +FLOPs(\\mathrm{FFN})\\\\ \u0026= 2*FLOPs(\\mathrm{normoalization}) + FLOPs(\\mathrm{Attention})+FLOPs(\\mathrm{FFN}) \\end{aligned} $$Normalization 现在有两种常见的 normalization, 也就是 LayerNorm 和 RMSNorm\nLayerNorm 定义如下\n$$ \\mathrm{LayerNorm}(x) = \\frac{x-\\mathbb{E}[x]}{\\sqrt{\\mathrm{var}[x]+\\epsilon}}\\odot \\beta + \\gamma $$其中 $\\beta,\\gamma\\in\\mathbb{R}^d$ 是可学习的参数。\n对输入 $x\\in\\mathbb{R}^{s\\times d}$, 均值需要约 $sd$ 次 FLOPs，方差 $\\mathrm{var}[x]$ 只需要约 $3sd$ 次 FLOPs，这两者的计算都可以忽略。接下来就是 scaling 和 shift, 这两者都是 element-wise 操作，我们这里，因此总的 FLOPs 为\n$$ \\boxed{FLOPs(\\mathrm{normoalization}) = 4sd} $$RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单\n$$ \\mathrm{RMSNorm}(x) = \\frac{x}{\\sqrt{\\|x\\|_2^2+\\epsilon}}\\odot \\gamma $$其中 $\\gamma\\in\\mathbb{R}^d$ 是可学习的参数\n对于 RMSNorm，其分析方式与 LayerNorm 基本一致，因此总的 FLOPs 为\n$$ \\boxed{FLOPs(\\mathrm{normoalization}) = 4sd} $$总之，不管使用哪种 normalization，其 FLOPs 都是\n$$ \\boxed{FLOPs(\\mathrm{normoalization}) = 4sd} $$Attention Attention 定义如下\n$$ \\mathrm{Attention}(X) = \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d}}\\right)V\\in\\mathbb{R}^{m\\times d} $$其中 $X\\in\\mathbb{R}^{s\\times d}$, $W_Q,W_K,W_V\\in\\mathbb{R}^{d\\times d}$\n$$ Q = W_QX\\in\\mathbb{R}^{s\\times d},\\quad K =W_KX\\in\\mathbb{R}^{s\\times d},\\quad V = W_VX\\in\\mathbb{R}^{s\\times d} $$$Q,K,V$ 计算的 FLOPs 为 $6*sd^2$. $QK^T$ 的 FlOPs 为 $2s^2d$, $\\mathrm{softmax}(\\cdot)V$ 的 FLOPs 为 $2s^2d$, 最后对于 multi-head attention 还有一个 output projection layer, 其权重为 $W_O\\in\\mathbb{R}^{d\\times d}$, 因此 FLOPs 为 $2sd^2$. 故 attention 最终的 FLOPs 为\n$$ FLOPs(\\mathrm{Attention})=6sd^2+2s^2d+2s^2d+2sd^2=\\boxed{8sd^2+4s^2d} $$FFN 对于 FFN，有两种常见的形式，一种基于 ReLU 激活函数，其定义如下\n$$ y = \\max(xW_1+b_1, 0)W_2 + b_2 $$其中 $W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$. $b_1\\in\\mathbb{R}^{d_{ff}}$, $b_2\\in\\mathbb{R}^{d}$.\n对输入 $x\\in\\mathbb{R}^{s\\times d}$, 其 FLOPs 为\n$$ FLOPs(\\mathrm{FFN_{ReLU}}) = 2sdd_{ff} + 2sd_{ff}d =\\boxed{4sdd_{ff}} $$其中第一项和第二项分别为为 $xW_1$ 与 $\\max(xW_1+b_1, 0)W_2$ 的 FLOPs.\n另一种基于 SwiGLU 激活函数，其定义为\n$$ \\mathrm{SwiGLU}(x) = x\\odot \\sigma(x) $$其中 $\\sigma(\\cdot)$ 是 sigmoid 函数\nFFN 的定义为\n$$ y = W_2(W_3x\\odot \\mathrm{SwiGLU}(W_1x)) $$其中 $W_3,W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$.\n对输入 $x\\in\\mathbb{R}^{s\\times d}$, 其 FLOPs 为\n$$ FLOPs(\\mathrm{FFN_{SwiGLU}}) = 2sdd_{ff} + 2sdd_{ff} + 2sd_{ff}d = \\boxed{6sdd_{ff}} $$Summary 最终，decoder-only transformer 的 FLOPs 计算量为 （我们假设使用 multi-head attention, 基于 SwiGLU 的 MLP）\n$$ \\begin{aligned} FLOPs(\\text{forward}) \u0026= FLOPs(\\text{embedding}) + n*FLOPs(\\mathrm{decode\\_layer})+FLOPs(\\mathrm{lm\\_head})\\\\ \u0026= n*FLOPs(\\mathrm{decode\\_layer})+2sd|V|\\\\ \u0026= n*(2*FLOPs(\\mathrm{normoalization}) + FLOPs(\\mathrm{Attention})+FLOPs(\\mathrm{FFN}))+2sd|V|\\\\ \u0026= n*(8sd+8sd^2+4s^2d+6sdd_{ff})+2sd|V|\\\\ \u0026= nsd^2\\left(\\frac8d + 8+\\frac{4s}{d}+\\frac{6d_{ff}}{d}+\\frac{2|V|}{nd}\\right)\\\\ \u0026\\approx \\boxed{nsd^2\\left(8+\\frac{4s}{d}+\\frac{6d_{ff}}{d}+\\frac{2|V|}{nd}\\right)} \\end{aligned} $$这里我们丢弃了 normalization 项，因为 normalization 的 FLOPs 是一个低阶项\nBackward 首先，我们有如下结论：\n神经网络 backward 过程的计算量（FLOPs）为 forward 过程的两倍\n我们使用一个简单的例子来证明这个结论，考虑线性层 $h=Wx$, 其中 $W\\in\\mathbb{R}^{m\\times d}$, 对于输入 $x\\in\\mathbb{R}^{d\\times 1}$ 其 forward 过程的计算量为 $2md$.\n对于反向过程，我们需要分别计算损失 $L$ 对权重和输入的梯度，即\n$$ \\frac{\\partial L}{\\partial x} = W^T\\frac{\\partial L}{\\partial h}\\in\\mathbb{R}^{d\\times 1}, \\quad\\frac{\\partial L}{\\partial W} = \\frac{\\partial L}{\\partial h}\\otimes x^T\\in{m\\times d}, $$这里 $\\frac{\\partial L}{\\partial h}\\in\\mathbb{R}^{m\\times 1}$ 为损失对输出 $h$ 的梯度。因此反向传播的总计算量为\n$$ 2dm + 2md = 4md $$这里的两项分别是对 $x$ 和 $W$ 求梯度的 FLOPs.\nOverall 将前向传播和反向传播的计算量汇总我们就得到一次前向传播和一次反向传播过程中，对于长度为 $s$ 的 token, 其 FLOPs 为 (multi-head attention, SwiGLU-FFN)\n$$ \\begin{aligned} C \u0026= FLOPs(\\text{forward}) + FLOPs(\\text{backward})\\\\ \u0026= 3FLOPs(\\mathrm{forward}) \\\\ \u0026\\approx \\boxed{3nsd^2\\left(8+\\frac{4s}{d}+\\frac{6d_{ff}}{d}+\\frac{2V}{nd}\\right)} \\end{aligned} $$Extension GQA GQA 与 MHA 不同的地方在于通过共享 key 和 value 来降低 KV cache 的占用，我们假设 group number 为 $g$, 则 key 和 value 的 FLOPs 现在变成了\n$$ 2sdd_h\\frac{h}{g}+2sdd_h\\frac{h}{g}=4sdd_h\\frac{h}{g} $$因此 attention 部分总的 FLOPs 变成了\n$$ FLOPs(\\mathrm{Attention})=4sd^2+2s^2d+2s^2d+4sdd_h\\frac{h}{g}=4sd^2+4s^2d+4sdd_h\\frac{h}{g} $$当 $g=h$ 时，GQA 就变成了 MHA, 此时的 FLOPs 也一致。\nMoE MoE 是针对 Dense FFN 的一个改进，介绍见 MoE, 我们假设一共有 $e$ 个路由专家，其中激活 $k$ 个。\nGate layer 一般是一个 linear layer, 其权重矩阵大小为 $W_{G}\\in\\mathbb{R}^{d\\times e}$, 因此 $FLOPs(\\text{router})= 2sde$.\nExpert layer 和前面提到的 FFN 一致，我们每次挑选出 $k$ 个专家进行计算，因此 expert 部分 $FLOPs(\\text{expert})=6ksdd_{ff}$.\n从而对于 MoE 来说，FFN 部分的 FLOPs 为\n$$ FLOPs(\\text{MoE}) = FLOPs(\\text{router})+FLOPs(\\text{expert})= \\boxed{2sde+6ksdd_{ff}} $$Simplification 我们已经得到了 transformer 的 FLOPs 计算表达式，但是其表达式比较繁琐，因此，在研究 scaling law 时，一般会进行简化。\n首先，在 LLM parameter analysis 中，我们已经给出了 LLM 参数量 $N$ （基于 Qwen3）的计算结果\n$$ N=n*(4d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd) + d(2|V|+1) $$我们这里对其进行简化，一般来说 $d_{ff}=8/3d$, $h_{kv}=h$, $h_d=d/h$, 带入就得到\n$$ N = n(4d+8d^2+2d^2+2d^2) + d(2|V|+1)=n(12d^2+4d)+ d(2|V|+1) $$我们忽略关于 $d$ 的一阶项，并且我们假设 $|V| \u003c\u003c 12nd$, 则最终模型参数量可以近似为\n$$ \\boxed{N \\approx 12nd^2} $$接下来，我们基于上面的配置简化 FLOPs 表达式\n$$ \\begin{aligned} C \u0026= 3nsd^2\\left(8+\\frac{4s}{d}+\\frac{6d_{ff}}{d}+\\frac{2|V|}{nd}\\right) \\\\ \u0026= 3nsd^2\\left(24+\\frac{4s}{d}+\\frac{2|V|}{nd}\\right)\\\\ \u0026\\approx 72nsd^2 \\\\ \u0026= 6sN \\end{aligned} $$这里我们利用了前面的 $|V| \u003c\u003c 12nd$ 假设，为了简便我们还舍弃了 $4s/d$.\n注意到 $s$ 代表 token 序列长度，如果训练集的总 token 个数为 $D$, 则最终对于包含 $D$ tokens 的数据集和包含 $N$ 参数量的 LLM, 其训练总 FLOPs 可以近似估计为\n$$ \\boxed{C\\approx 6ND} $$Experiments Setting 接下来我们定量分析一些模型的 FLOPs. 我们基于 Chinchilla scaling law 给出的实验配置 (Table A9), 我们筛掉 kv_size * n_heads != d_model 的配置，$|V|=32,000$.\n各部分的 FLOPs 计算代码如下\n1 2 3 4 5 def compute_flops(n, V, d, d_ff, s): lm_head_flops = 2 * V * d * s attention_flops = 8 * s * d * d + 4 * s * s * d feed_forward_flops = 4 * s * d * d_ff return lm_head_flops, attention_flops, feed_forward_flops 首先我们看一下不同大小模型的 FLOPs 分布情况\n可以看到，当模型越来越大，FFN 层的算力占比越来越高，这也是为什么后来采用 MoE 架构的一个原因。\n接下来，我们看一下模型 FLOPs 分布情况随上下文长度变化的情况\n可以看到，随 context length 增加，attention 的算力占比逐渐上升，这符合 attention 是一个平方复杂度的算法的结论。并且，当上下文足够长之后，计算还出现了 overflow (图像右端的突然下降)。\nConclusion 在本文中，我们介绍了如何计算基于 decoder-only transformer LLM 的 FLOPs, 并推导除了 Kaplan scaling law 中使用的公式 $C=6ND$, 这为后面的 infra 和 scaling law 的学习提供了基础。\nReferences Chinchilla Scaling law pytorch embedding layer transformer flops ","date":"2025-10-15T16:33:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/llm-flops-computation/","title":"LLM FLOPs Computation"},{"content":"快手提出了 Keye-VL 1.5, 一个强调 reasoning, video understanding 的 8B 多模态大模型。作者提出了 slow-fast video encoding strategy 来提高模型的视频理解能力，作者通过在预训练和后训练提高了模型的长上下文能力和 reasoning 能力\nIntroduction 作者回顾了多模态大模型的进展，然后提到视频理解能力仍然是一个挑战。为了解决这个问题，作者提出了 Keye-VL-1.5, 一个 8B 的视频理解多模态大模型。\nKeye-VL-1.5 主要做了三点改进：\n在架构上，使用了 Slow-Fast Video Encoding 在预训练阶段，使用多个 stage 来提升模型的长上下文能力 在 post-training 阶段，提高模型的 reasoning 能力和 alignment 表现 Method Architecture Keye-VL 1.5 的架构与 Keye-VL 一致。\nKeye-VL 1.5 主要做出的改进点为针对视频的 encoding 方式\n作者回顾了已有 MLLM 处理视频的方式，比如 Qwen2.5-VL 使用 3D convolution 来 merge 相邻的两帧，Seed1.5-VL 采用了 Dynamic Frame-Resolution Sampling 技巧，来根据 budget 和处理的任务来动态调整采样率 (frame) 和每一帧的图片精度 (resolution)。\n但是这些方法很难进行泛化。因此，作者在本文中就提出了 SlowFast video encoding stratrgy:\nSlow Pathway: 空间信息丰富（high resolution），时间信息简略 (low number of frames) Fast Pathway: 时间信息丰富（high number of frames），空间信息简略 (low resolution) 为了区分 slow/fast frames, 作者提出了一个基于 patch similarity 的 metric:\n第一帧始终定义为 slow frame 接下来的每一帧，如果其和上一帧的相似度超过 $95\\%$, 则定义为 fast frame; 反之则定义为 slow frame. 得到 slow/fast frames 之后，作者将 fast-frame 的 token budget 限制为 slow frame token budget 的 $30\\%$ 来平衡时间信息以及空间信息。接下来，作者使用二分搜索来决定 slow frame 的 token budget. 为了区分 slow frame 和 fast frame 的 token, 作者使用了特殊的 token 来进行分离。\n最终的处理结果如下图所示\nPre-training 预训练的数据和 Keye-VL 基本一致，我们主要介绍改进的点\n对于 Image caption 数据，作者认为这批数据可能会损害模型的指令跟随和 reasoning 能力，因此作者对数据进行了增广，主要是调整了数据的格式：\nQA, 数据格式为 \u0026lt;image, caption, [eos], question, answer\u0026gt; reverse QA, 数据格式为 \u0026lt;image, question, answer, [eos], caption\u0026gt; instruction following: 随机给一批数据作为输入，然后让模型基于特定 image 输出 caption 作者还构建了一个 trap question 来提高模型的 robustness 以及 faithfulness.\nOCR 数据在 Keye-VL 的基础上加入了两点：\nStructured Document and Code Understanding: 基于 markdown 和 HTML 等数据来获取 code OCR 数据 Instruction Following OCR: 基于特定指令进行 OCR 对于 grounding 数据，作者进一步加入了 temporal grounding 数据，作者首先使用 TEMPURA\n来将短视频分割成若干个 video clips. 然后作者使用 SOTA MLLM 来过滤数据，最后作者基于 Gemini2.5 来生成对应的 QA.\n预训练和 Keye-VL 一样，包含 3 个 stage\n前两个 stage，作者将模型的上下文限制为 8K, 使用了 DP 和 Zero-2 来减少内存开销。在 stage 3, 作者将模型的上下文从 8K 扩展到 128K, 对应的 base frequency 从 1M 提升到 8M. 训练数据包括长视频，长文本和大规模图片。作者将优化策略调整为 Zero-1, CP 和 PP 来支持 long-context 的训练。训练时数据分布为 video:images:text=24:50:26.\nPost-training SFT 阶段使用了 7.5M 多模态 QA 样本进行训练。\nMPO 阶段的数据相比 Keye-VL 有所减少，包含：\n250K 开源样本 150K 纯文本数据 26K 人类标注数据 对于 reward model 的训练，作者使用了 SFT 和 RL 两个阶段。SFT 阶段的数据包括 R1-Reward 和 MMPR, 训练之后作者还是用比较短的 good response 来避免产生较长的回答\n在 SFT 和 MPO 阶段之后，作者使用 LongCoT code-start 初步激活模型的 reasoning 能力。\n作者构建了一个 5 部的自动化数据生成 pipeline, 如下图所示\n步骤如下：\nMulti-Source Data Collection and Enhancement：收集数据 Multi-Path Reasoning Generation with Confidence Quantification: 基于 confidence 来挑选数据 Comprehensive Two-Level Quality Assessment: 基于答案和过程的正确性来提高数据质量 Human-in-the-Loop Quality Enhancement: 对于中等质量的数据请人类进一步进行标注 Dynamic Quality Scoring and Data Utilization Strategy: 对数据进行打分，高质量数据进行上采样 接下来就是 General RL 过程。\n对于通用的 RLVR 训练，作者使用了 GSPO 算法来进行训练。\n在训练过程中，作者采取了 progressive hint sampling 方式，也就是提供不同程度的 hint 来提高模型的训练效率。作者将 hint 分为五个等级：\nLevel 1 (Concept / Observation) Level 2 (Strategy / Method) Level 3 (Tools / Formula) Level 4 (Steps / Calculation) Level 5 (Complete Solution) 来提供不同程度的辅助，作者使用 Keye-VL 1.5 来确定最小的 hint level, 然后基于这个 level 提供信息进行 RL 的训练\n为了进一步提高模型的表现，作者采用了一个和 Seed1.5-VL 一样的迭代式训练策略，即反复进行 SFT 和 RL 来降低训练成本，提高训练效率。\n最后，再通用 RL 阶段之后，作者加入了 alignment RL 的训练来进行对齐，reward 包括三个方面：\nrule-based reward generative reward model-based reward 任务主要包括三个方面：\ninstruction following format adherence preference alignment 数据介绍如下：\ninstruction following:25 类硬约束，20 类软约束，数据包括 17K 多模态数据和 23K 纯文本数据，奖励包括 rule-based reward 和 generative reward reasoning: 12K 数学和逻辑推理数据 RAG: 提高模型的搜索能力，作者使用 GSPO 算法进行训练 Experiments 模型表现如下\n接下来作者进行了消融实验。\n首先是不同训练阶段对模型表现的影响，如下图所示\n实验结果显示，提高 SFT 训练数据可以有效提高模型在数学推理，逻辑推理和 OCR 任务上的表现。MPO 可以进一步提高模型的表现，Long CoT cold start 可以有效提高模型的 reasoning 表现。\n作者还探究了 model merging 对模型表现的影响，作者首先基于 base model 和 OCR 数据训练得到 OCR expert, 然后进行 merge, 实验结果如下\n实验结果显示，model merging 可以有效提高模型在 special domain 上的表现，并且还可以维持模型的通用能力\n作者还发现：\nexpert model 训练时间过长会影响最终 merge model 的表现 expert mode 训练的学习率应该要设置比较小 接下来，作者探究了 alignment RL 对模型表现的影响，\n实验结果说明，alignment RL 可以在保持模型 reasoning 能力的同时提高模型的指令跟随能力\n作者还探究了 hint 对模型表现的影响，使用不同 level hint 进行训练对模型表现影响的结果如下图所示\n实验结果显示，使用 high level 的 hint 可以有效提高模型输出的正确率。\n最后作者探究了 rejection sampling 对模型表现的影响，实验结果如下图所示\n结果发现，通过 rejection sampling，模型的表现有了进一步的提升。因此作者采用了 ST-RL-(RFT-SFT)-(RFT-RL) 的训练方式来进行训练\nConclusion 作者在本文中提出了 Keye-VL1.5, 一个强调 video understanding 和 reasoning 的 MLLM. Keye-VL 1.5 使用了 SlowFast video encoding strategy 来提高模型的效率和视频理解能力。作者详细介绍了模型的数据，训练和评估。\nReferences arxiv ","date":"2025-09-11T11:33:31+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-keye-vl-1.5/","title":"Notes on Keye-VL 1.5"},{"content":"作者提出了一个针对 Adam 优化器的 weight decay 方法\nIntroduction 作者首先回顾了动态梯度算法如 AdaGrad, RMSProp, Adam 的进展。已有工作表明动态梯度算法的泛化性要比 SGD with momentum 要差。作者在本文中探究了在 SGD 和 Adam 中使用 L2 regularization 和 weight decay 对最终模型表现的影响。结果表明，模型泛化性较差的原因在于对于 Adam, L2 regularization 的效果要比 SGD 差。\n作者有如下发现：\nL2 regularization 和 weight decay 不等价。在 SGD 中，L2 regularization 是等价的，但是在 Adam 中这个结论不成立。具体来说，L2 regularization 对历史参数的惩罚要小于 weight decay L2 regularization 对 Adam 效果提升有效 weight decay 对于 SGD 和 AdamW 都很有效，在 SGD 中，weight decay 与 L2 regularization 等价 最优的 weight decay 取决于 batch, batch 越大，最优的 weight decay 越小 通过 learning rate scheduler 可以进一步提高 Adam 的表现 作者在本文中的主要贡献是通过解耦梯度更新中的 weight decay 来提高 Adam 的 regularization.\n作者的主要 motivation 是提升 Adam 表现，让其可以和 SGD with momentum 相比\nMethod Weight decay 的定义如下\n$$ \\theta_{t+1} = (1-\\lambda)\\theta_t - \\alpha \\nabla f_t(\\theta_t) \\tag{1} $$其中 $\\lambda$ 是 weight decay rate, $\\nabla f_t(\\theta_t)$ 是第 $t$ 个 batch 的梯度，$\\alpha$ 是学习率。\n首先，对于标准的 SGD 来说，weight decay 与 L2 regularization 等价\nProposition 1 对于标准的 SGD 来说，对损失函数 $f_t(\\theta_t)$ 执行 weight decay （公式 $(1)$）与对损失函数 $f_t(\\theta_t)+\\lambda'/2\\|\\theta_t\\|_2^2$ 执行梯度下降算法是等价的，这里 $\\lambda'=\\lambda/\\alpha$。\n证明比较简单，只需要写出损失函数的梯度下降更新公式即可。\n基于这个结论，大部分优化算法都将 L2 regularization 和 weight decay 看做是等价的。但实际上，这个结论对于 adaptive gradient 方法来说是不成立的。结论如下\nProposition 2 令 $O$ 为一个 optimizer, 其目标函数为 $f_t(\\theta)$, 当不考虑 weight decay 时，梯度更新过程为 $\\theta_{t+1}\\gets \\theta_t-\\alpha M_t\\nabla f_t(\\theta_t)$. 当考虑 weight decay 时，梯度更新过程为 $\\theta_{t+1}\\gets (1-\\lambda)\\theta_t-\\alpha M_t\\nabla f_t(\\theta_t)$. 如果 $M_t\\neq kI$, 则不存在 $\\lambda'$, 使得 $O$ 在优化目标函数 $f_t^{reg}(\\theta)=f_t(\\theta)+\\lambda'/2\\|\\theta\\|_2^2$ 时，不考虑 weight decay 的梯度更新与 $O$ 在优化目标函数 $f_t(\\theta)$ 时，考虑 weight decay 的梯度更新等价。\n证明比较简单，只需要写出两个目标函数对应的梯度更新公式即可。\n作者通过分析发现，在 adaptive gradient 方法中，对于 L2 regularization，梯度和 regularization 是打包在一起考虑的。而 weight decay 是分开考虑的。这就导致了对于梯度比较大的权重，L2 regularization 的学习率较小，从而 regularization 效应减弱。而 weight decay 中，这种效应则不存在。因此 weight decay 的 regularization 效应更强。\n作者通过这个分析，给出了一个 weight decay 与 L2 regularization 相等的条件\nProposition 3 令 $O$ 为一个 optimizer, 其目标函数为 $f_t(\\theta)$, 当不考虑 weight decay 时，梯度更新过程为 $\\theta_{t+1}\\gets \\theta_t-\\alpha M_t\\nabla f_t(\\theta_t)$. 当考虑 weight decay 时，梯度更新过程为 $\\theta_{t+1}\\gets (1-\\lambda)\\theta_t-\\alpha M_t\\nabla f_t(\\theta_t)$. 如果 $M_t= \\mathrm{diag}(s)^{-1}$ ($s_i\u003e0,\\forall i$), 则 $O$ 在优化目标函数\n$$ \u003e f_t^{reg}(\\theta)=f_t(\\theta)+\\frac{\\lambda'}{2\\alpha}\\|\\theta\\odot \\sqrt{s}\\|_2^2 \u003e $$时，不考虑 weight decay 的梯度更新与 $O$ 在优化目标函数 $f_t(\\theta)$ 时，考虑 weight decay 的梯度更新等价。\n上面的结论显示，对于比较大的 preconditioner $s_i$, 其在相比于 L2 regularization 被 regularized 的效应更强。\n为了解耦这两个参数，作者提出了 SGDW 算法，其 weight decay 和梯度更新同时进行，算法如下图所示\n在算法中，为了支持同时给 $\\alpha$ 和 $\\lambda$ 做 scheduling, 作者提出了一个 scaling factor $\\eta_t$, $\\eta_t$ 由用户定义的 scheduler SetScheduleMultiplier(t) 决定。此时，针对 SGD with momentum 的 weight decay 与 L2 regularization 是等价的\n同理，我们也可以对 Adam 算法实行同样的操作，算法如下图所示\nConclusion 作者在本文中分析了 adaptive gradient 方法中 L2 regularization 与 weight decay 的不一致性。基于分析，作者提出了 SGDW 和 AdamW 两个优化算法。\nReferences arxiv ","date":"2025-09-04T10:27:03+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-adamw/","title":"Notes on AdamW"},{"content":"作者提出了 Adam, 一个一阶的优化方法，Adam 更加高效，且具有 scaling invariant 的性质。\nIntroduction 作者首先回顾了一下已有优化器的进展，其中主要是 SGD. 在本文中，作者提出了 Adam, 一个针对高维参数空间的一阶优化器，Adam 基于 gradient 的一阶和二阶信息为不同的参数安排不同的学习率。Adam 的来源是 adaptive moment estimation. Adam 主要是结合了 AdaGrad 和 RMSProp 两个算法的优点。\nAdam 与 RMSProp 的区别在于：\nRMSProp 在 rescaled gradient 上进行 momentum 的计算然后更新，而 Adam 直接使用一阶和二阶矩来进行估计 RMSProp 没有 bias-correction 项 Adam 的主要优势为：\n参数更新的量级与 gradient 的 scaling 无关 步长被 stepsize 超参数限制 不要求目标函数 stationary 对于稀疏梯度 work 的比较好 优化器自带 annealing Algorithm Adam 的算法如下图所示\n我们优化的目标函数如下\n$$ \\min_{\\theta}\\quad f(\\theta) $$这里 $f$ 一般是一个神经网络。我们记 $f(\\theta)$ 在 $\\theta_t$ 处的梯度为 $g_t=\\nabla_{\\theta}f(\\theta_t)$.\n算法运行时，会更新梯度 $m_t$ 以及梯度二阶矩 $v_t$ 的 exponential moving average. 超参数 $\\beta_1,\\beta_2$ 负责控制 exponential decay rates. 这里 $m_t$ 和 $v_t$ 分别是一阶动量（均值）和二阶动量（未中心化的 variance）的估计。由于 $m_t$ 和 $v_t$ 的初始化都是 0, 因此他们会引入 bias, 作者在后续通过修正解决了这个问题。\n假设 $\\epsilon=0$, 如果除了当前时刻 $t$ 之外，之前所有时刻的梯度 $g_i=0,i","date":"2025-09-04T10:11:55+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-adam/","title":"Notes on Adam"},{"content":"作者系统性分析了已有的 attention 机制，然后作者提出了混合的 attention 机制，来提高模型在长上下文的表现以及维持模型在短上下文场景下的表现。\nIntroduction 作者首先强调了提升 LLM 上下文长度面临的问题：\n如何有效处理长上下文输入 如何训练长上下文 LLM 如何降低长上下文 LLM 在 inference 时的 latency 以及 memory usage 对于建模长上下文输入，我们可以从 attention 机制或者位置编码来入手。前者类似的工作有 Landmark Attention 和 Focused Transformer, 但是这些方法的问题在于训练不稳定。 QK norm 可以比较好解决 softmax 分布过于极端的问题，但是其问题在于训练时的数值不稳定性，并且可能会影响模型的长上下文能力。\n另一方面，对于位置编码，已有的工作如 APE, AliBi, RoPE 等都可以提供位置信息。但是，这些方法很有可能回影响模型最终的 attention score 分布。另外，NoPE 探究了移除 position encoding 的可能性。\n还有一些工作目的是降低 softmax attention 的时间复杂度和空间复杂度。比如 sliding window attention, sparse attention, attention sink 等都可以降低整体的时间/空间复杂度。但是这些方法最终的表现都有所下降。\nObservation 作者首先对比了以下不同方法对模型长上下文能力的影响。\n作者训练了一个 8B 的模型，然后分别对比了三种方法：\nRoPE: base frequency 设置为 10,000, SFT 阶段扩展到 2M QK-Norm: 在 RoPE 的基础上，对 query 和 key 先进行 normalization 再进行 RoPE NoPE: 移除 attention 中的位置编码信息 作者分别评估了三种方法的表现，实验结果如下表\nModel Val Loss MMLU HellaSwag CommonsenseQA ARC-E ARC-C Needles 65k RoPE 1.52 48.55 73.74 68.30 81.05 39.13 9.82 QK-Norm 1.53 48.21 73.68 68.23 80.54 38.98 7.93 NoPE 1.58 47.61 72.16 66.42 76.94 37.12 9.03 实验结果显示，对于通用任务，RoPE 和 QK-Norm 的表现差不多，而 NoPE 的表现较差。对于长上下文任务，QK-Norm 的表现最差。\n接下来，作者分析了三种方法的 attention 分布情况。作者将上下文分为四个 segments：\nbegin: 开始的 10 个 token needle: 与 needle 相关的 tokens context: 通用的上下文 token qc: question/completion token, 语文题答案相关的 token 作者将 needle 放置在 50% 深度的位置。评测的实验结果如下\nContext Length Model Variants begin needle context qc 8k RoPE 0.3863 0.0328 0.3809 0.2000 QK-Norm 0.0242 0.0173 0.8020 0.1565 NoPE 0.3058 0.0454 0.4501 0.1987 32k RoPE 0.3541 0.0201 0.4343 0.1915 QK-Norm 0.0064 0.0056 0.8517 0.1364 NoPE 0.2807 0.0325 0.4981 0.1886 128k RoPE 0.3463 0.0010 0.4751 0.1776 QK-Norm 0.0010 0.0004 0.8993 0.0994 NoPE 0.0846 0.0073 0.8156 0.0925 实验结果显示：\nNoPE 是最关注 needle token 信息的，RoPE 次之，QK-Norm 最差 QK-Norm 更关注上下文信息，对其他的信息关注度较少 作者发现 QK-Norm 对初始的 token 信息关注度较少，作者认为这是因为 normalization 会让模型更关注邻近的 token 信息。为了验证这个观点，作者在不同的上下文长度下，分别计算了不同方法的 attention 分布情况，为了避免噪声，作者将开始的 10 个 token 以及最后的 $3\\%$ token 排除在外，实验结果如下图所示\n实验结果说明，相比于 softmax attention, QK-Norm 的 attention score 分布更平滑，其对于 need token 的注意力不如 RoPE, 这也说明了为什么 QK-Norm 的长上下文能力比较差。并且，QK-Norm 的 recency bias 更严重。\n作者通过计算 RoPE 和 QK-Norm 的 attention distribution 的 entropy 来进一步说明这一点，结果如下表所示\nModel 8k 32k 128k RoPE 6.02 6.95 7.62 QK-Norm 10.71 12.46 14.14 结果显示 QK-Norm 的熵更高，这意味着 QK-Norm 的 attention score 分布更分散，也就证明了 Qk-Norm retrieval 能力比较差。\nMethod 考虑到 NoPE 和 RopE 各自的优点，作者提出了一个混合架构，来结合 NoPE 与 RoPE. 具体做法就是 NoPE layer 和 RoPE layer 交替进行。作者将这个模型架构记为 RNoPE.\nRNoPE 不同 layer 与不同 base frequency 产生的 attention score 分布如下\nModel NoPE Layers - begin NoPE Layers - needle NoPE Layers - context NoPE Layers - qc RoPE Layers - begin RoPE Layers - needle RoPE Layers - context RoPE Layers - qc needles-128k RoPE - - - - 0.3541 0.0201 0.4343 0.1915 7.395 RNoPE-10k 0.3275 0.0765 0.5672 0.0287 0.0049 0.0004 0.6805 0.3142 8.036 RNoPE-100k 0.3263 0.0778 0.5633 0.0327 0.0241 0.0005 0.6782 0.2972 7.461 RNoPE-2M 0.3250 0.0712 0.5735 0.0303 0.1111 0.0046 0.6233 0.2611 7.022 RNoPE-4M 0.3486 0.0369 0.5981 0.0165 0.0960 0.0039 0.6774 0.2227 6.203 RNoPE-10k-swa 0.3303 0.0742 0.5634 0.0321 - - - - 9.562 实验结果显示，在 RNoPE 架构中，\n提升 base frequency 带来的增益逐渐递减 NoPE layer 的 retrieval 能力比较强，表现在 attention sink 现象以及对于 needle token 的 spike. 但是其 recency bias 比较弱 RoPE 的 retrieval 能力比较弱，但是其 attention sink 现象比较小 当 base frequency 增加止呕，RoPE 的 recency bias 会降低，表现在 qc token 的权重降低 作者发现，RoPE 的 receptive field 会影响 NoPE layer 的 retrieval 能力。作者总结得到两个 insight\nNoPE layer 对于 retrieval 任务比较擅长，而 RoPE layer 对于处理 local Information 比较擅长 限制 RoPE 的 receptive field 可以保证 NoPE layer 的 retrieval 能力 基于这两个 insight, 作者构建了 RNoPE-SWA 架构，RNoPE-SWA 相比于 RNoPE, 将 full attention 变成了 sliding window attention, 来避免 RoPE layer 对下游的 NoPE layer 产生影响。\n最终，作者基于 Command R+ 构建了模型，作者去除了 QK-Norm, 对于 NoPE layer, 作者使用了 full attention, 对于 RoPE layer, 作者使用了 window size 为 4096 的 sliding window attention. 作者通过实验验证了 NoPE layer 和 RoPElayer 的比例，实验结果发现 $1:3$ 的比例是最优的。最终每 4 个 layer 为一组，前三组为 SWA, 最后一层为 full attention.\n最终，在通用任务上评测的结果如下\nModel MMLU HellaSwag ARC-E ARC-C SATEn SATMath GSM8K Winogrande MBPP Baseline 57.5 75.8 84.6 48.5 70.0 30.9 40.9 68.5 39.1 RNope-SWA 59.5 76.2 82.5 48.8 71.9 30.5 42.7 69.5 39.3 在 Ruler retrieval 上评测的结果如下\nModel 8k 16k 32k 64k 128k 256k Baseline 96.6 94.4 95.1 89.1 83.0 57.1 RNope-SWA 96.1 96.1 94.9 92.0 90.0 74.8 在 Ruler QA 上评测的结果如下\nModel 8k 16k 32k 64k 128k 256k Baseline 53.5 50.0 52.5 45.5 36.0 30.0 RNope-SWA 55.5 52.5 55.5 49.0 46.0 42.5 实验结果说明，模型在通用任务任务上与 baseline 模型表现差不多，且长上下文能力更强\nConclusion 作者提出了 RNope-SWA, 一个混合 NoPE, RoPE position embedding 和 sliding window attention 的 attention 机制。RNope-SWA 可以在保持模型表现的同时提高计算效率和降低 KV cache 占用。\n尽管本文和已有的工作如 YoCo, Jamba-1.5 和 MiniMax-01 均证明了混合 attention 架构的有效性。但是，目前对于其工作机制尚缺乏一个比较系统性的理解。作者认为这是一个需要探究的方向。另外，作者还认为如何更好的汇总上下文信息与位置信息也是一个可以探究的方向。\nReferences arxiv ","date":"2025-09-02T11:24:10+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-rnope-swa/","title":"Notes on RNoPE-SWA"},{"content":"上海 AI LAB 提出了 InternVL 3.5 系列多模态大模型，InternVL 3.5 主要强调了模型的 reasoning 能力以及 inference 效率。\nIntroduction 作者首先介绍了已有多模态大模型的进展，然后提到有两个未解决的问题：\n如何构建一个 stable, effective, scalable 的 RL 框架来训练 MLLM 如何降低 MLLM 在长上下文场景下的计算开销过高的问题 为了解决这两个问题，作者提出了 InternVL3.5 多模态大模型系列，相比于 InternVL3, 模型在 reasoning 能力和 efficiency 上均进行了提升。作者主要通过 Cascade RL 框架来提高模型的 reasoning 表现，以及使用 Visual Resolution Router (ViR) 和 Decoupled Vision Language Deployment (DvD) 来提高模型的推理效率。\n总结起来，InternVL3.5 模型的贡献如下：\n开源了 InternVL3.5 系列多模态大模型，模型拥有先进的 reasoning 能力以及推理效率 提出了 Cascade RL, ViR 以及 DvD 等模块来提高模型的表现/推理效率 系统性评估了模型的表现 Method Architecture InternVL3.5 的架构与 InternVL3 的架构基本一致，包括一个 ViT, 一个 MLP 和一个 LLM, 模型的架构如下图所示\n模型配置如下表所示\nInternVl3.5 延续了 InternVL 系列的做法，即 20B 以下的模型使用 InternViT-300M, 20B 以上的模型使用 InternViT-6B 作为 Visual encoder. 大语言模型基于 Qwen3 和 gpt-oss.\n在 InternVL3.5 的基础上，作者还构建了 InternVL3.5-Flash 模型，InternVL3.5-Flash 不同的地方在于，其在 MLP 之前加入了一个 Visual Resolution Router (ViR) 模块，来处理不同分辨率的图片。在 InternVL3.5 中，每个 image patch 由 1024 个视觉 token 来表示，然后通过 pixel shuffle 压缩至 256 个。 在 InternVL3.5-Flash 中，ViR 首先会判断每个 patch 的 semantic richness, 然后基于 semantic richness 来将对应的 token 路由到不同的 pixel shuffle 模块中，最高可以将视觉 token 压缩到 64 个，这样就可以大幅度提高模型的推理效率。\nPre-training Pre-training 使用 next-token-prediction loss 来更新模型权重，给定多模态 sequence $x=(x_1,\\dots,x_L)$, 损失函数定义为\n$$ \\mathcal{L}(\\theta) = -\\sum_{x_i\\text{ is text token}} \\log p_{\\theta}(x_i \\mid x_1,\\dots,x_{i-1}) $$与 InternVL2.5 一样，作者还对不同长度的 sequence 进行了加权，避免模型产生 bias, 加权后的 loss 为\n$$ \\mathcal{L}(\\theta) = -\\sum_{x_i\\text{ is text token}}\\frac{w_i}{\\sum_j w_j} \\log p_{\\theta}(x_i \\mid x_1,\\dots,x_{i-1}),\\quad w_i = \\frac{1}{N^{0.5}} $$其中 $N$ 是训练样本中需要计算损失的 token 个数。\n训练数据蛀牙包含两部分：\n多模态数据，这部分数据基于 InternVL3 纯文本数据，基于 InternLM 系列和开源的数据集 最终，预训练数据一共包含116M 样本，250B token. 纯文本数据和多模态数据的比例为 $1:2.5$. 模型上下文长度为 $32K$.\nPost-training Post-training 包含三个阶段：\nSFT： 使用高质量对话数据提高模型表现 Cascade RL: 提高模型的 reasoning 能力 Visual Consistency Learning: 训练 ViR 模块，提高模型处理动态分辨率图片的能力 训练 pipeline 如下图所示\nSFT 阶段的训练数据包括三个方面：\nInternVL3 的指令跟随数据 多模态 reasoning 数据 能力扩展数据，包括 GUI 交互，具身交互以及 SVG 理解与生成的数据 Cascade RL 阶段结合了 offline RL 和 online RL 的优点。基本思想是先使用 offline RL 来提高模型的能力，降低训练成本。然后再使用 online RL 进一步提高模型的 reasoning 表现。\noffline RL 阶段使用的是 InternVL-MPO 算法，其损失函数为\n$$ \\mathcal{L}_{MPO} = w_p\\mathcal{L}_p+w_q\\mathcal{L}_q+w_g\\mathcal{L}_g $$其中，$\\mathcal{L}_p$ 为 DPO 的损失函数，$\\mathcal{L}_q$ 为 Quality loss, $\\mathcal{L}_g$ 为 SFT 使用的 next-token-prediction loss.\nonline RL 阶段使用的是 GSPO 算法，核心思想是在 sample 层面做归一化。这个阶段的目的是进一步提高模型的表现。\nCascade RL 的优势在于：\n训练更稳定：offline RL 可以避免模型产生 reward hacking 现象，online RL 阶段可以进一步提高模型的表现 训练效率更高：offline RL 可以有效提高采样效率 表现更好：先 MPO 再 RL 可以达到更好的表现 Visual Consistency Learning (ViCO) 的目的是降低 InternVL3.5 的推理开销。这个阶段主要训练 ViR 模块\nViCO 包括两个 stage:\nStage 1: Consistency training\n这一阶段的目的是使得不同压缩率处理的视觉 token 对应的输出与 ground truth 尽可能一致，作者使用另外一个 InternVL3.5 模型作为 reference model, 然后损失函数为\n$$ \\mathcal{L}_{ViCO} = \\mathbb{E}_{\\xi\\sim\\mathcal{R}}\\left[\\frac{1}{N}\\sum_{i=1}^N\\mathrm{KL}\\left(\\pi_{\\mathrm{ref}}(y_i\\mid y_{","date":"2025-09-01T11:30:50+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-internvl3.5/","title":"Notes on InternVL3.5"},{"content":"作者提出了 Ovis2.5, 一个基于 Ovis 改进的多模态大模型系列，包括 2B 和 9B 两个 size，Ovis2.5 主要强调了支持不同分辨率图片输入以及深度思考这两个 feature\nIntroduction 作者首先回顾了 Ovis, Ovis 主要是解决 text embedding 以及 visual embedding 对齐程度比较低的问题。\n接下来，作者介绍了以下 Ovis 的两个问题：\n只能支持固定大小的图片输入 缺乏深度思考能力 为了解决这两个问题，作者提出了 Ovis 2.5, Ovis 主要做出了两点改进：\n使用了 NaViT 来处理不同分辨率图片的输入 作者通过训练提高了模型的深度思考能力 最终 Ovis2.5 主要有以下 feature\n支持动态分辨率图片输入 深度思考能力 SOTA 的表现 高效的训练方式 Method Architecture Ovis2.5 的架构如下所示\nOvis 包括三个模块：\nvisual tokenizer： ViT 架构， visual embedding table: 类似 LLM 中的 text embedding table, 见 Ovis LLM: 基于 Qwen3 作者在架构上进行了如下改进：\n动态分辨率图片输入处理：作者使用了 NaViT 来支持动态分辨率图片输入 LLM: 作者使用了 Qwen3 来进一步提高模型的表现 Training 模型训练包括 pre-training 和 post-training 两个大的 stage, 其中 pre-training 又包含 3 个小的 stage, post-training 包含 2 个 stage. 训练过程如下所示\npre-training 阶段的数据包括 COYO, Laion, Wukong, DataComp, SAM 等。作者介绍了几个部分的数据：\nOCR 数据，作者基于 MLLM 来标注数据和合成 QA Grounding 数据，作者使用了 RefCoCo 等数据集以及先进的 MLLM 来标注数据 Reasoning 数据，作者收集了数据然后使用 MLLM 来合成 Reasoning path 训练时，\nVET pretraining: 训练 VET, 作者基于 SigLIP 来初始模型的参数，然后仅训练最后一层 ViT layer, visual head 以及 VET, 图片精度为 448-896. 作者采用了动态 position embedding Multimodal pretraining: 这阶段全量微调所有参数，主要目的是使用对话格式的数据。图片精度为 448-1792 multimodal instruction tuning: 这阶段训练所有参数，主要提高模型跟随多模态指令的能力 post-training 包括 DPO 和 GRPO 两个阶段。\nDPO: 训练所有参数，使用 pre-training checkpoint 来多次采样 GRPO: 使用 RLVR 数据集进行训练 Infra infra 方面，作者主要强调了 data packing 以及多种并行策略融合。\nConclusion 作者在本文中提出了 Ovis2.5, 一个基于 Ovis 架构的多模态大模型，作者主要强调了模型的动态图片输入处理能力以及深度思考能力。\n作者提出了几个未来的方向：\n将输入图片精度提升到 4K 处理长视频输入并进行 temporal reasoning 在 Reasoning 过程中加入 tool-use. References arxiv ","date":"2025-08-30T17:34:44+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/ovis2.5-mllm-with-stronger-perception-and-reasoning-capability/","title":"Ovis2.5 MLLM with stronger perception and reasoning capability"},{"content":"作者提出了 Ovis，一个离散化表示 visual encder 输出特征的方法，来更好对齐 LLM 的视觉输入和文本输入\nIntroduction 作者分析了已有多模态大模型的架构，已有多模态大模型的输入对于文本来说是离散的 (text token), 对于图片来说是连续的 (visual embedding)。作者认为这种连续 - 离散的输入可能会影响模型最终的表现。\n为了解决这个问题，作者构建了一个 visual embedding table, 将 visual embedding 也转换成离散的 token 表示形式，进而统一 LLM 输出的粒度。\nMethod 模型的架构如下图所示\n我们首先会构建一个 visual vocabulary $\\{e_k\\}_{k=1}^K$, 其大小为 $K$, 然后对于 ViT 输出的 $n$ 个 visual feature $\\{r_i\\}_{i=1}^n$, 我们会加入一个 linear head 以及一个 softmax 来构建一个 vocabulary 上的分布，即\n$$ v_i = \\mathrm{softmax}(Wr_i), W\\in\\mathbb{R}^{K\\times d} $$这里 $v_i\\in\\Delta^K$ 是 visual vocabulary 上的概率分布。最终，视觉模块的输入是 vocabulary 中 visual token 的一个加权求和\n$$ V_i = \\sum_{k=1}^K v_{i,k}e_k\\in\\mathbb{R}^{d'} $$训练分为三个阶段：\nStage 1: 训练 $W$, visual encoder 最后一个 block 以及 visual vocabulary Stage 2: 训练 $W$, visual vocabulary 以及 visual encoder Stage 3: multimodal SFT, 提高模型的指令跟随能力，模型所有参数都参与训练 训练数据分布如下表所示\nConclusion 作者提出了 Ovis，一个离散化表示 visual encder 输出特征的方法，来更好对齐 LLM 的视觉输入和文本输入\nReferences Arxiv ","date":"2025-08-30T17:32:22+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/ovis-discrete-visual-embedding/","title":"Ovis-discrete visual embedding"},{"content":"DeepSeek 在 2024 年 1 月发布了 DeepSeekMoE, 一个解决 MoE 模型 specialization 不足以及 redundancy 问题的大模型系列。\nIntroduction 作者首先回顾了已有 MoE 模型的不足，主要有两点：\nknowledge hybridity: 已有 MoE 模型的专家个数比较少，这就导致每个专家需要掌握更多样化的知识，提高了训练难度 knowledge redundancy: 不同专家掌握的知识可能有重叠，从而导致了模型参数的 redundancy 为了解决这两个问题，作者提出了 DeepSeek-MoE, DeepSeek-MoE 主要做出了两点改变：\nFine-Grained Expert Segmentation: 作者使用了更多的专家，来提高每个专家的 specialization, 降低训练成本 Shared Expert Isolation: 作者在 Routing expert 的基础上，加入了 shared expert 来学习 common knowledge. 作者在 2B-A0.6B 的模型进行了实验，结果显示模型表现超过了 GShard, 说明了 DeepSeekMoE 模型架构的有效性。\n作者还进一步将模型 scale 到了 16B-A2.8B 和 145B-A22B, 实验结果均验证了模型的 scaling 效果。\nMethod Preliminary 作者首先回顾了 transformer 架构，transformer 的第 $\\ell$ 层 decode layer 可以表示为\n$$ \\begin{aligned} u_{1:T}^{\\ell} \u0026= \\mathrm{self\\_attention}(h_{1:T}^{\\ell-1}) + h_{1:T}^{\\ell}\\\\ h_t^\\ell \u0026= \\mathrm{FFN}(u_t^{\\ell}) + u_t^{\\ell} \\end{aligned} $$其中 $T$ 是 sequence length, $h_{1:T}^{\\ell-1}$ 是第 $\\ell-1$ 层 decoder layer 输出的 hidden states.\n接下来，我们可以将 dense 架构转换为 MoE 架构，MoE 架构与 dense 架构不同的地方在与 $\\mathrm{FFN}$ 不再是 MLP, 而是一个 MoE 模块，其表达式如下\n$$ \\begin{aligned} h_t^\\ell \u0026= \\sum_{i=1}^N\\left(g_{i,t}\\mathrm{FFN}_i(u_t^{\\ell})\\right) + u_t^{\\ell}\\\\ g_{i,t} \u0026= \\begin{cases} s_{i,t,}, \u0026s_{i,t}\\in\\mathrm{Topk}(\\{s_{j,t}\\mid 1\\leq j \\leq N\\},K)\\\\ 0, \u0026\\text{otherwise} \\end{cases}\\\\ s_{i,t,} \u0026= \\mathrm{softmax}_i({u_t^{\\ell}}^Te_i^{\\ell}) \\end{aligned} $$这里 $N$ 是专家的总个数， $K$ 是激活专家个数， $e_{i}^{\\ell}$ 是 routing layer 的权重矩阵，$\\mathrm{FFN}_i$ 是每个专家对应的 FFN.\nDeepSeekMoE Architecutre DeepSeekMoE 架构如下图所示\n相比于其他 MoE 架构，DeepSeekMoE 主要做了以下几点改变。\nFine-Grained Expert Segmentation 作者首先解决了每个专家学习内容过多的问题。作者的做法就是将每个 expert FFN 分割为 $m$ 个更小的专家，具体做法就是将 FFN intermediate hidden size 降低为原来的 $1/m$. 这样的话就可以在不增加模型参数量的情况下提高模型的表现。修正后的 MoE 模块为\n$$ \\begin{aligned} h_t^\\ell \u0026= \\sum_{i=1}^{mN}\\left(g_{i,t}\\mathrm{FFN}_i(u_t^{\\ell})\\right) + u_t^{\\ell}\\\\ g_{i,t} \u0026= \\begin{cases} s_{i,t,}, \u0026s_{i,t}\\in\\mathrm{Topk}(\\{s_{j,t}\\mid 1\\leq j \\leq mN\\},mK)\\\\ 0, \u0026\\text{otherwise} \\end{cases}\\\\ s_{i,t,} \u0026= \\mathrm{softmax}_i({u_t^{\\ell}}^Te_i^{\\ell}) \\end{aligned} $$可以看到，现在我们一共有 $mN$ 个专家，激活专家个数为 $mK$ 个。作者认为，通过提高专家的粒度，我们可以有效增加专家组合的可能性，这就提高了最终的组合多样性。\nShared Expert Isolation 接下来，作者介绍了解决不同专家学习到重复知识的问题，作者的做法是在 routing Expert 的基础上加入 Shared expert. 也就是说，有固定几个专家始终都会被激活，这部分专家复杂学习通用知识，从而减少知识冗余。增加 shared expert 之后的 MoE 模块为\n$$ \\begin{aligned} h_t^\\ell \u0026= \\sum_{i=1}^{K_s}\\mathrm{FFN}_i(u_t^{\\ell})+\\sum_{i=K_s+1}^{mN}\\left(g_{i,t}\\mathrm{FFN}_i(u_t^{\\ell})\\right) + u_t^{\\ell}\\\\ g_{i,t} \u0026= \\begin{cases} s_{i,t,}, \u0026s_{i,t}\\in\\mathrm{Topk}(\\{s_{j,t}\\mid K_s+1\\leq j \\leq mN\\},mK-K_s)\\\\ 0, \u0026\\text{otherwise} \\end{cases}\\\\ s_{i,t,} \u0026= \\mathrm{softmax}_i({u_t^{\\ell}}^Te_i^{\\ell}) \\end{aligned} $$此时，模型中一共包含 $K_s$ 个共享专家，$mN-K_s$ 个 routing expert, 其中激活专家个数为 $mK-K_s$.\nLoad Balancing Loss 接下来，作者解决了训练时的 load imbalance 问题，作者提出了两个 loss 来分别解决不同层面的 load imbalance 问题。\n首先，在 expert 层面，作者使用了如下的 load balancing loss:\n$$ \\mathcal{L} = \\alpha_1\\sum_{i=1}^{N'}f_iP_i $$其中 $\\alpha_1$ 是超参数，\n$$ f_i = \\frac{N'}{K'T}\\sum_{i=1}^{N'}\\mathbb{1}(\\text{Token }i \\text{ selects Expert }i),\\quad P_i = \\frac{1}{T}\\sum_{t=1}^Ts_{i,t} $$分别为以及分配给第 $i$ 个专家的 token 比例以及概率之和。$N'=mN-K_s$, $K'=mK-K_s$. $\\mathbb{1}(\\cdot)$ 是 indicator function.\n其次，在 device 层面，作者也是用了 load balancing loss 来减少不同设备之间不必要的通信。作者将 routed experts 分为 $D$ 个 group $\\mathcal{E}_1,\\dots,\\mathcal{E}_D$, 然后每个设备部署一个 group, group level 的 load balancing loss 定义如下：\n$$ \\mathcal{L} = \\alpha_2\\sum_{i=1}^D f_i' P_i' $$其中 $\\alpha_2$ 是超参数，\n$$ f_i' = \\frac{1}{|\\mathcal{E}_i|}\\sum_{j\\in\\mathcal{E}_i}f_i,\\quad P_i' = \\sum_{j\\in\\mathcal{E}_i}P_i $$实际中，作者使用了一个较小的 $\\alpha_1$ 来避免 routing collapse, 使用了一个较大的 $\\alpha_2$ 来提高 Device 层面的负载均衡。\nTraining 作者使用了中英文数据进行训练，tokenizer 基于 BPE 算法，vocab size 为 8K. 模型训练基于 HAI-LLM.训练时使用了 TP, DP, PP, EP 等并行策略。\n2B, 16B, 145B 模型的参数如下表所示\nModel 2B 16B 145B total params 2B 16.4B 144.6B activated params 0.3B 2.8B 22.2B hidden size 1280 2048 4096 layers 9 28 62 attention heads 10 16 32 head dimension 128 128 128 routed experts 63 64 128 activated experts 7 6 12 shared experts 1 2 4 training tokens 100B 2T 245B Experiments Alignment 作者针对 DeepseekMoE 16B 进行了微调，微调使用了 1.4M 的训练样本，覆盖了 math, code, QA, reasoning 等任务。\nAblation Study 作者在 2B 的模型上进行了 ablation study.\n首先，作者探究了细粒度专家和共享专家的有效性，结果如下图所示\n实验结果显示，与 GShard 相比，使用共享专家可以有效提高模型的表现。并且，使用更细粒度的专家也可以进一步提高模型的表现\n作者还探究了共享专家与路由专家的比例，作者分别使用不同的比例进行实验，结果发现共享专家：路由专家个数为 1：3 的时候模型效果最好。\n作者还探究了模型的泛化性，作者 mask 掉一部分概率最高的 routing expert, 然后从剩下的专家里进行 topK 的挑选，然后作者比较模型和 GShard 的表现，结果如下图所示\n实验结果显示，DeepSeekMoE 对于 mask 操作更敏感，这说明 DeepSeekMoE 模型中专家的 specialization 更强。\n作者还探究了 mask 掉共享专家对模型表现的影响，结果显示共享专家与路由专家之间的 overlap 很小，去掉共享专家之后，模型表现会变差。\n作者进一步分析了共享专家与路由专家组合的有效性。作者探究 DeepSeekMoE 是否可以使用更少的路由专家来获取知识。作者通过使用不同的 activated routed experts 来进行实验，实验结果如下图所示\n实验结果显示，DeepSeekMoE 仅需激活 4 个路由专家，就可以达到与 GShard 相同的表现。这说明了 DeepSeekMoE 模型中每个专家可以学习到更准确的知识。\nConclusion 作者在本文中提出了 DeepSeekMoE 架构，来提高 MoE 模型中专家的利用效率。为了达到这一点，作者首先使用了更细粒度的专家，降低每个专家的学习成本，然后，作者使用了共享专家，来降低不同专家之间的知识冗余。作者先在 2B 的模型上验证了方法的有效性，然后作者将模型 scale 到了 16B 和 125B。结果显示模型的效果均超过了以前的工作。\nReferences arxiv ","date":"2025-08-29T11:03:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseekmoe/","title":"Notes on DeepSeekMoE"},{"content":"DeepSeek 在 2024 年 1 月 5 日发布了 DeepSeek LLM, 包括 7B 和 67B 两个 size, 作者主要强调了对于 scaling law 的探究\nIntroduction 已有的 scaling law 如 Kaplan 和 Chinchilla 介绍了 model size, dataset size, compute budget 与模型表现之间的关系。在本文中，作者进一步探究了 learning rate 和 batch size 等超参数与模型表现之间的关系。基于发现的 scaling law, 作者为不同大小的模型设置了最优的超参数。并且，作者还发现不同数据集与模型表现之间的关系。\n最终，基于这些实验结果，作者提出了 DeepSeek LLM, 模型使用 2T token 进行预训练，使用 1M samples 进行后训练，后训练包括 SFT 以及 DPO.\nPre-training Architecture DeepSeek-LLM 的架构与 LLaMA 基本相同，作者在 67B 的模型上使用了 GQA 来提高 inference 效率。最终模型的配置如下表所示\nParams 7B 67B $n_{\\text{layers}}$ $30$ $95$ $d_{\\text{model}}$ $4096$ $8192$ $n_{\\text{heads}}$ $32$ $64$ $n_{\\text{kv\\_heads}}$ $32$ $8$ Context Length $4096$ $4096$ Sequence Batch Size $2304$ $4608$ Learning Rate $4.2e-4$ $3.2e-4$ Tokens 2T 2T Data 作者主要从 Common Crawl 构建预训练数据，数据处理过程包括：去重，过滤以及 remixing 三个步骤。\n对于 tokenizer, 作者使用了 BBPE 算法，tokenizer 的大小设置为 100,000, 最终的 tokenizer 大小为 102400.\nHyper Parameters 作者主要对比了一下不同 learning rate schedule 的表现：\ncosine learning schedule multi-step learning rate schedule: 包含三个 Stage, 第一个 stage 保持最大学习率，第二个 stage 将学习率降低为最大学习率的 $31.6\\%$, 第三个 stage 降低为最大学习率的 $10\\%$. 对比的实验结果如下图所示\n实验结果显示，multi-step learning rate scheduler 的表现与 cosine learning rate 表现差不多。并且，multi-step learning rate scheduler 对于 continue pretraining 支持更好。因此在本文中作者使用了 multi-step learning rate scheduler.\nInfra 作者使用了数据并行，张量并行，序列并行以及 1F1B pipeline 并行。作者还使用了 flash attention 来提高硬件利用率。\nScaling Law 本节中，作者分析了 scaling law, 主要有以下三点：\n构建了针对 learning rate 和 batch size 的 scaling law 作者使用 non-embedding FLOPs/token $M$ 来表示 model scale 预训练数据的质量对最后中的 scaling 影响很大 作者首先构建了针对 batch size 和 learning rate 的 scaling law, 结果显示最优的 learning rate 和 batch size 范围都比较广，这个结论与 Kaplan 一致。\n接下来，作者构建了 batch size $B$, learning rate $\\eta$ 与 compute budget $C 之间的关系，实验结果如下图所示\n拟合得到的曲线为\n$$ \\begin{aligned} \\eta_{opt} \u0026= 0.3118* C^{-0.1250}\\\\ B_{opt} \u0026= 0.2920 * C^{0.3271} \\end{aligned} $$可以看到，随着 compute budget 增加，$B_{opt}$ 也逐渐增加，而 $\\eta_{opt}$ 逐渐减小。并且，最优参数的范围都比较广。\n接下来，作者进一步探究了 batch size 与 generalization error $L$ 之间的关系。作者希望找到 model scale $N$, data scale $D$ 与 compute budget $C$ 之间的关系，即\n$$ N_{opt} \\varpropto C^a,D_{opt} \\varpropto C^b $$compute budget 与 model scale, data scale 之间的关系可以近似表示为 $C=6ND$, 这个公式的推导见 LLM FLOPs computation。我们用 $N_1,N_2$ 分别表示模型的 non-embedding parameter 以及 complete parameters, 则我们可以用 $6N_1$ 或者 $6N_2$ 来近似 model scale, 但是 $6N_1$ 和 $6N_2$ 均没有考虑 attention 的计算开销，因此这两种近似的误差都比较大。\n为了解决这个问题，作者提出了一个新的 model scale 表示形式，即 non-embedding FLOPS/token $M$, 其中 $M$ 包含 attention 的计算开销但是不包含 vocabulary computation. 基于这种表示，compute budget 可以近似表示为 $C=MD$. $M$ 与 $6N_1,6N_2$ 的区别表示如下所示\n$$ \\begin{aligned} 6N_1 \u0026= 72nd^2\\\\ 6N_2 \u0026= 72nd^2 + 6Vd\\\\ M \u0026= 72nd^2+12ndl \\end{aligned} $$其中, $d$ 是 hidden size, $n$ 是 layers 个数, $V$ 是 vocabulary size, $l$ 是 sequence length. 作者在不同 scale 的模型上比较了三种表示方式，结果发现 $6N_1$ 和 $6N_2$ 要么低估，要么高估了模型的参数量。\n基于 model scale 的表示方式，作者构建了如下的优化问题\n$$ M_{opt}(C), D_{opt}(C) = {\\arg\\min}_{M,D\\ s.t.\\ C=MD} L(N,D) $$作者使用了 Chinchilla 提出来的 IsoFLOP 曲线进行拟合，实验结果如下图所示\n拟合的曲线为\n$$ M_{opt}(C) = 0.1715*C^{0.5243}, D_{opt}(C) = 5.8316*C^{0.4757} $$作者还进一步拟合了 compute budget 与 optimal generalization error 之间的关系，结果如下图所示\n实验结果显示，作者提出的 scaling law 可以很好预测模型的表现。\n最后，作者探究了以下不同数据集的 scaling law, 作者使用 early in-house data, current in-house data 以及 OpenWebText2 来将进行实验，结果如下图所示\n结果显示，scaling law 与数据质量高度相关。当数据质量提升时，model scaling exponent $a$ 逐步提升，data scaling exponent $b$ 逐步下降，说明 compute budget 更多由模型参数量决定。因此，作者认为提升 compute budget 之后，我们应该优先提高模型的 model size.\nPost-training 作者构建了 1.5M 的中英文指令数据。其中安全性的数据有 300K, 有帮助性的数据有 1.2M, 其中包括 $31.2\\%$ 的通用数据，$46.6\\%$ 的数学相关数据，$22.2\\%$ 的代码数据。\npost-training 包含两个阶段：\nSFT:7B 的模型训练了 4 个 epoch, 67B 的模型训练了 2 个 epoch, 作者发信进一步训练 67B 的模型会导致过拟合。作者发现，模型在训练过程中会出现重复输出的情况，特别是数学 SFT 数据，为了解决这个问题，作者使用了一个两阶段的 SFT 以及 DPO. DPO: 提高模型的能力，作者发现 DPO 可以提高模型 open-ended generation skill. Evaluation 我们主要关注一下消融实验。\n首先作者探究了分阶段 SFT 对模型表现的影响。作者发现，小模型在 math 和 code 数据集上需要训练更长时间，但是这也损害了模型的对话能力。为了解决这个问题，作者使用两阶段的训练模式，第一个阶段使用所有的数据进行训练，第二个阶段仅使用对话数据进行训练，实验结果如下表所示\nModel HumanEval GSM8K Repetition IFEval DeepSeek LLM 7B Chat Stage1 48.2 63.9 0.020 38.0 DeepSeek LLM 7B Chat Stage2 48.2 63.0 0.014 41.2 可以看到，经过第二阶段训练之后，模型的表现有所提升\n接下来，作者探究了 Multi-choice question 对模型表现的影响，MCQ 要求模型不仅需要有相关的知识，还要理解每个选项的含义。作者使用 20M 中文 MCQ 来进行消融实验，结果如下表所示\nModel MMLU CEval CMMLU TriviaQA ChineseQA DeepSeek LLM 7B Chat 49.4 47.0 49.7 57.9 75.0 DeepSeek LLM 7B Chat + MC 60.9 71.3 73.8 57.9 74.4 实验结果显示，MCQ 确实可以提高模型在上述几个 benchmark 上的表现，但是其泛化性会下降。因此，作者在 pre-training 和 fine-tuning 阶段并没有使用 MCQ 数据进行训练。\n作者还探究了在 pre-training 阶段加入 instruction data, 来提高 base model 在下游 benchmark 上的表现。结果发现，base model 的表现提升优先。作者认为，尽管 instruction data 可以提高 base model 表现，但是如果 Instruction data 数量过少，则模型表现不太可能学习到有用的知识。因此，作者的做法是不在 pretraining 阶段加入 Instruction data.\n最后，作者探究了 system prompt 对模型表现的影响。受 LLaMA 2 启发，作者也尝试在输入中加入 system prompt. 实验结果如下所示\nModel MT Bench DeepSeek LLM 7B Chat 7.15 DeepSeek LLM 7B Chat + System Prompt 7.11 DeepSeek LLM 67B Chat 8.35 DeepSeek LLM 67B Chat + System Prompt 8.58 可以看到，7B 的模型加入 system prompt 之后，模型表现有所下降；67B 的模型加入 system prompt 之后，模型表现有所提升。作者认为，大模型更容易理解 system prompt 的意图，而小模型的指令跟随能力则较差，因此 system prompt 反而会影响模型表现。\nConclusion 作者在本文中提出了 DeepSeek LLM 系列大语言模型，作者详细介绍了超参数的选择以及 scaling law 等。\nReferences arxiv ","date":"2025-08-26T10:53:10+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-deepseek-llm/","title":"Notes on DeepSeek-LLM"},{"content":"阶跃星辰等提出了 Multi-matrix Factorization Attention (MFA), 一个新型注意力机制，用于在 KV cache 限制下最大化模型的表现。\nIntroduction multi-head attention (MHA) 的问题在于，其 KV cache 的内存占用（memory footprint）随 sequence length 以及 batch size 线性增长，从而成为了 LLM 在 decoding 阶段的瓶颈。\n为了解决 MHA 的内存占用过高问题，已有的工作如 MQA, GQA 等通过共享 key, value projection 来降低 KV cache size. 而 DeepSeek-V3 提出的 MLA 则是通过对 key, value projection 进行 low-rank compression, 然后只存储 latents 的方法来降低 KV cache size.\n但是，已有的这些方法的问题在于，当我们设置 KV cache budget 之后，它们的表现就比标准的 MHA 要差。\n基于以上这些发现，作者首先分析了已有 attention 机制的 modeling capacity, 然后使用一个统一的框架来表示这些 attention 机制。作者发现，attention heads 的个数以及 dimension 对模型表现有较大影响。\n基于这个发现，作者提出了 Multi-matrix Factorization Attention (MFA), 以及其变体 MFA-Key-Reuse (MFA-KR). MFA 的主要目的是在有限的 KV cache size 下提高模型的表现。\nBackground 作者首先介绍了 GMHA 的概念，GMHA 由三部分组成：\nQK circuit: 决定了信息之间如何交互 valueoutput (VO) circuits：决定了信息如何传递 per-head softmax attention. 接下来，作者介绍了 Fully Parameterized Bilinear Attention (FPBA), FPBA 的定义如下：\n$$ O = \\sum_{c=1}^d\\left(\\sum_{j=1}^N\\phi\\left(\\frac{xW_cx_j}{H}\\right)x_jU_c\\right) $$其中 $\\phi$ 是 softmax 函数，$d$ 是模型的 hidden dimension, $N$ 是 sequence length, $W_c,U_c\\in\\mathbb{R}^{d\\times d}$ 每个 channel 上的参数矩阵\n每个 channel 都有各自的参数 $W_c, U_c$ 来获取 $x_i$ 与 $x_j$ 之间的信息 提高泛化性，所有 channel 的 $U_c$ 组合起来可以遍历 $d$ 维空间中的任意一个 permutation, 这样就避免来的信息损失 利用率高，FPBA 获取了 $x_i$ 与 $x_j$ 之间 $d$ 维空间可能的表示 基于以上这三个特点，作者认为 FPBA 是 GMHA 框架的一个 capacity upper bound. 此时每个 token 的 KV cache 占用为 $2d^2$ (key and value).\n然后，作者分析了 MHA 及其变体与 GMHA 的关系，MHA 可以写作如下形式\n$$ \\begin{aligned} O \u0026= \\sum_{c=1}^h\\left(\\sum_{j=1}^N\\phi\\left(\\frac{xQ_c(x_jK_c)^T}{\\sqrt{d}}\\right)x_jV_c\\right)O_c^T\\\\ \u0026= \\sum_{c=1}^h\\left(\\sum_{j=1}^N\\phi\\left(\\frac{x(Q_cK_c^T)x_j^T}{\\sqrt{d}}\\right)x_jV_cO_c^T\\right) \\end{aligned} $$其中 $Q_c,K_c,V_c\\in\\mathbb{R}^{d\\times h_d}$, $O_c\\in\\mathbb{R}^{d\\times h_d}$ 分别是 query, key, value, output projection layer 对应的权重矩阵，$n$ 是 attention head 的个数，令 $h_d$ 为每个 attention 的 head dimension，则我们有 $nh_d=d$.\n可以看到，MHA 实际上是一个特殊的 FPBA, 其中，$W_c$ 和 $U_c$ 分别由秩为 $h_d$ 的低秩分解 $Q_cK_c^T$ 以及 $V_cO_c^T$ 近似。此时每个 token 的 KV cache 占用为 $2d$ (key and value).\nMQA 可以看作是 GQA 的一个特殊情况。对于 GQA 来说，我们有一个 group size $g\\in[1, h]$, 当 $g=1$ 时，GQA 就是 MHA. 当 $g=h$ 时，GQA 就是 MQA, 通常 $g$ 满足 $h\\ \\%\\ g=0$. GQA 的表达式与 MHA 基本相同，只是多个 head 会共享一个 $K_c$ 以及 $V_c$. 此时，每个 token 的 KV cache 占用为 $2gh_d$. 对于 MQA，其每个 token 的 KV cache 占用为 $2h_d$.\n对于 MLA, 其表达式如下所示\n$$ \\begin{aligned} O \u0026= \\sum_{c=1}^m\\left(\\sum_{j=1}^N\\phi\\left(\\frac{xS_QQ_c(x_jS_KK_c)^T}{\\sqrt{d}}\\right)x_jS_VV_c\\right)O_c^T\\\\ \u0026= \\sum_{c=1}^m\\left(\\sum_{j=1}^N\\phi\\left(\\frac{x(S_QQ_cK_c^TS_K^T)x_j^T}{\\sqrt{d}}\\right)x_jS_VV_cO_c^T\\right) \\end{aligned} $$其中，$S_Q,S_K,S_V\\in\\mathbb{R}^{d\\times C}$ 在所有的 heads 中是共享的，$Q_c,K_c,V_c\\in\\mathbb{R}^{C\\times h_d}$ 是每个 head 的 query, key, value projection layer 的参数， 是 latent factorization 的维度。与 FPBA 相比，我们可以看到，MLA 实际上是在 $d/m$ 个 head 上共享了参数，其中，$W_c$ 和 $U_c$ 分别由秩为 的低秩分解 $S_QQ_cK_c^TS_K^T$ 以及 $S_VV_cO_c^T$ 近似。尽管模型中 $C\u003eh_d$, 但是最终的 rank 仍然是 $h_d$, 因此模型的表现也就受到了限制。\nMethod 对已有的 attention 分析之后，作者认为，要提高模型的表现，attention 需要做到亮点：\n最小化 KV cache 占用和参数量 attention 的 capacity 尽可能接近 FPBA 基于这两个原则，作者提出了 MFA, MFA 主要依赖三个策略：\n提升 attention heads 的 head dimension, 通过提高 head dimension, 我们可以有效提高 attention head 的表达能力 使用矩阵分解来降低参数量 使用单一的 KV head 来降低 KV cache 内存占用 最终，MFA 的表达式如下所示\n$$ \\begin{aligned} O \u0026= \\sum_{c=1}^m\\left(\\sum_{j=1}^N\\phi\\left(\\frac{xS_QQ_c(x_jS_K)^T}{\\sqrt{d}}\\right)x_jS_V\\right)O_c^T\\\\ \u0026= \\sum_{c=1}^m\\left(\\sum_{j=1}^N\\left(\\frac{x(S_QQ_cS_K^T)x_j^T}{\\sqrt{d}}\\right)x_jS_VO_c^T\\right) \\end{aligned} $$其中 $S_Q,S_K,S_V\\in\\mathbb{R}^{d\\times C}$ 是所有的 attention head 所共享的，$Q_c,O_c\\in\\mathbb{R}^{C\\times C}$ 是每个 head 的 query up projection 和 output projection, $C$ 是 latent factorization 的维度。\n在 inference 的时候，由于我们只需要保存 $x_jS_K$ 和 $x_jS_V$, 因此所需要的 KV cache size 为 $2C$. 与 FPBA 相比，MFA 分别使用 $S_QQ_cS_K^T$ 和 $S_VO_c^T$ 来近似 $W_c$ 和 $U_c$, 近似矩阵的 rank 为 $C$. 由于 $C\u003ed$, 因此其表达能力也更强，MFA 有如下优势：\nscalable head count: MFA 可以支持使用更多的 attention heads, 每增加一个 heads, 所需要的额外参数为 $2C^2$. 并且，增加 attention heads 个数不会增加 KV cache 占用 enhanced head expressiveness: MFA 近似矩阵的 rank 为 $C\u003ed$, 因此表达能力更强 Compatibility with position encodings: MFA 可以无缝集成 position encoding. 为了进一步降低 MFA 的 KV cache 占用，作者提出了 MFA-Key-Reuse (MFA-KA). 核心思想是使用 $S_K$ 来表示 $S_V$, 这样可以额外降低 $50\\%$ 的 KV cache 占用，表示方法如下所示\n$$ S_V = S_K + \\alpha\\odot NS_K = (I +\\mathrm{diag}(\\alpha)N)S_K $$其中 $N\\in\\mathbb{R}^{N\\times N}$, $\\alpha\\in\\mathbb{R}^C$.\n最终，MFA, MFA-KR 与 GQA 的对比如下图所示\n不同 attention 的量化对比如下表所示\nMethod KV Cache Parameter Heads Factor. rank per head Shared latent subspace Dim. Total effec. rank FPBA $2d^2$ $2d^3$ $d$ $d$ $d$ $d^2$ MHA $2d$ $4d^2$ $n$ $h_d$ $d$ $nh_d$ MQA $2h_d$ $(2 + 2/n)d^2$ $n$ $h_d$ $h_d$ $nh_d$ GQA $2gh_d$ $(2 + 2g/n)d^2$ $n$ $h_d$ $gh_d$ $nh_d$ MLA $2C$ $5dC + d^2$ $m$ $h_d$ $C$ $mh_d$ MFA $2C$ $3Cd + 2mC^2$ $m$ $C$ $C$ $mC$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 class Step3vAttention(nn.Module): def __init__(self, config: Step3VLConfig, layer_idx): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, \u0026#34;head_dim\u0026#34;, config.hidden_size // config.num_attention_heads) self.num_key_value_heads = 1 self.total_num_kv_heads = self.num_key_value_heads self.num_attention_heads = config.num_attention_heads self.num_key_value_groups = config.num_attention_heads // self.num_key_value_heads self.q_size = getattr(config, \u0026#34;share_q_dim\u0026#34;, self.head_dim) self.kv_size = self.num_key_value_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.is_causal = True self.q_proj = nn.Linear(config.hidden_size, self.q_size , bias=False) self.k_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) # query down projection normalization self.inter_norm = Step3vRMSNorm(self.q_size, eps=config.rms_norm_eps) # query up projection self.wq = nn.Linear(self.q_size, self.head_dim * self.num_attention_heads, bias=False) def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -\u0026gt; Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: input_shape = hidden_states.shape[:-1] query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2) value_states = self.v_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2) query_states = self.inter_norm(query_states) query_states = self.wq(query_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) ... Conclusion 作者在本文中提出了 MFA 以及 MFA-KR, 一个在 KV cache 有限的条件下最大限度提高 attention 表达能力的 attention 机制。\nReferences Multi-matrix Factorization Attention Step v3 code ","date":"2025-08-23T16:04:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-mfa/","title":"Notes on MFA"},{"content":"MX format 是一个表示数据的数据格式，在 LLM 中主要用于量化。相比于直接对整个张量进行量化，MX format 可以在更细粒度的层面控制量化，从而提高模型的表现\nMicroscaling Microscaling (MS) format 如下图所示\nMX format 包括三个部分：\nelements $P_1,\\dots,P_k$ 未 scale 的数据，要求 $P_1,\\dots,P_k$ 的数据类型相同 shared scale $X$, 对 element 进行的 scale 参数，所有的 $k$ 个 bits 共享一个 $X$ block size, 决定 element block 的大小 在存储时，我们只需要存储 $X$ 以及 $P_1,\\dots,P_k$, 我们假设 $X$ 需要 $w$ bits 来表示，$P_i$ 需要 $d$ bits 来表示，则我们一共需要 $w+kd$ bits 来表示这 $k$ 个元素。\nConcrete MX-compliant Formats MX-format 包含了一下几种数据格式：\nFormat Name Element Data Type Element Bits(d) Scaling Block Size(k) Scale Data Type Scale Bits(w) MXFP8 FP8 (E5M2) 8 32 E8M0 8 MXFP8 FP8 (E4M3) 8 32 E8M0 8 MXFP6 FP6 (E3M2) 6 32 E8M0 8 MXFP6 FP6 (E2M3) 6 32 E8M0 8 MXFP4 FP4 (E2M1) 4 32 E8M0 8 MXINT8 INT8 8 32 E8M0 8 GPT-oss Quantization gpt-oss 中使用了 MXFP4 来表示 MoE 中的 down projection 以及 up projection weight matrix 的权重。\n其具体操作过程如下：\n我们将参数分为大小为 32 的 block 每个 block 由一个 scale $X$ 来表示，其精度为 E8M0, 即 8bits, 表示范围为 $[-127,127]$, 以及 $32$ 个元素 $P_i$ 来表示，每个元素的精度为 E2M1, 即 4bits, 表示范围为 $[-6.0,6.0]$. 由于每个元素由 4bits 来表示，因此我们将两个元素合并在一起来表示 在加载时，我们可以用如下代码来恢复 $P_i$ 的值到 FP8\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 FP4_VALUES = [+0.0,+0.5,+1.0,+1.5,+2.0,+3.0,+4.0,+6.0,-0.0,-0.5,-1.0,-1.5,-2.0,-3.0,-4.0,-6.0] def convert_moe_packed_tensors( blocks, scales, *, dtype: torch.dtype = torch.bfloat16, rows_per_chunk: int = 32768 * 1024, ) -\u0026gt; torch.Tensor: import math # scales are represented with uini8 scales = scales.to(torch.int32) - 127 assert blocks.shape[:-1] == scales.shape, f\u0026#34;{blocks.shape=} does not match {scales.shape=}\u0026#34; lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device) *prefix_shape, G, B = blocks.shape rows_total = math.prod(prefix_shape) * G blocks = blocks.reshape(rows_total, B) scales = scales.reshape(rows_total, 1) # each byte representing 2 elements and represented with unit8 out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device) for r0 in range(0, rows_total, rows_per_chunk): r1 = min(r0 + rows_per_chunk, rows_total) blk = blocks[r0:r1] exp = scales[r0:r1] # nibble indices -\u0026gt; int64 idx_lo = (blk \u0026amp; 0x0F).to(torch.long) idx_hi = (blk \u0026gt;\u0026gt; 4).to(torch.long) sub = out[r0:r1] sub[:, 0::2] = lut[idx_lo] sub[:, 1::2] = lut[idx_hi] torch.ldexp(sub, exp, out=sub) del idx_lo, idx_hi, blk, exp out = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2) # to match for now existing implementation return out.to(torch.float8_e5m2) References gpt-oss code report ","date":"2025-08-21T18:23:03+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-mx-format/","title":"Notes on MX-format"},{"content":"作者提出了 flashattention, 一个通过降低 multi head attention 内存访问开销来提高 attention 计算效率的方法\nIntroduction Transformer 的 attention 是一个平方度复杂度的算法，这个平方复杂度既体现在时间复杂度上（矩阵乘法），也体现在空间复杂度上（需要存储中间结果）。因此，要降低 attention 的复杂度，我们有两种思路：\n从时间复杂度上入手，比如使用稀疏 attention 机制或者线性注意力机制 从空间复杂度上入手，比如使用 GQA, MQA 等减少内存的访问开销 本文提出的 flashattention 就属于降低空间复杂度的一种做法。作者认为，我们应该设计一种 IO-aware 的 attention 算法，来减少 attention 计算式的内存访问开销，进而提高 attention 的计算效率。\n作者首先提到，一个未解决的问题就是：\n降低 attention 的内存访问开销是否可以提高 attention 的计算效率？\n作者发现，已有的一些工作虽然在理论上降低了 attention 的计算效率，但是在实际中，他们的效果并没有提升太多。作者分析原因认为，已有工作主要关注于降低 FLOPs, 但是忽略了内存访问开销。\n因此，作者在本文中就提出了 flashattention, 一个 IO-aware 的 attention 算法，作者通过尽可能降低内存访问开销来提高模型的计算效率。具体做法就是，避免从内存中读写 attention matrix, 作者认为这个目标有两个挑战：\n计算 softmax 的时候不访问所有的输入 在反向传播时不存储中间的 attention matrix 作者提出了两个方法来分别解决这两个问题：\n作者使用了 tiling 技巧，将 input 分成多个 block, 然后分别进行处理，进而降低 softmax 的内存访问开销 作者使用了 recompute 技巧，在反向传播时，重新计算 softmax normalization factor 通过这些改进，我们可以让 attention 运行更快，并且降低内存访问开销。\n作者还从理论上分析了 flashattention 的复杂度，提供了理论基础。\n作者通过实验验证了 flashattention 的有效性，主要是三点：\n训练效率更高：相比于 Huggingface 和 Megatron, flashattention 的训练效率提升了 2-3 倍 模型的表现更好：相比于 GPT-2, 模型的 perplexity 提升了 0.7 个点左右 速度更快：flashattention 比标准的 attention 实现快 3 倍以上 Background Hardware Performance 作者首先介绍了以下 GPU 的内存架构，如下图所示\n可以看到，GPU 内存可以分为三个层级：\nSRAM: GPU 的寄存器，容量小，但是访问速度极快 High bandwith memory (HBM): GPU 的高速内存，访问速度较快，容量中等 DRAM: CPU 内存，容量最大，但是访问速度较慢 接下来作者介绍了 Execution model 的概念，GPU 有多个线程来执行同一个操作（SPMD），这个操作也被称为 kernel, kernel 会从 HBM 中加载输入到 SRAM 中进行计算，然后写回 HBM.\n对一个算法，我们可以将其归类为 compute-bound 和 memory-bound 两类， 我们可以用 arithmetic intensity 来进行区分，arithmetic intensity 定义为 arithmetic operations 与 memory access 的比率。\ncompute bound: 算法的瓶颈在于算力，由于算力不足导致运行时间慢，比如矩阵乘法 memory-bound: 算法的瓶颈在于内存访问效率，比如 element-wise 操作或者是 reduction 为了提高 memory-bound 类型算法的效率，我们进行 kernel fusion, 即把多个访问同一片内存的操作放在一起处理，避免多次读写内存\nStandard Attention Implementation 作者还回顾了一下标准化的 attention 实现。\nForward Pass 给定 $Q,K,V\\in\\mathbb{R}^{N\\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下\n$$ S = QK^T\\in\\mathbb{R}^{N\\times N},\\quad P = \\mathrm{softmax}(S)\\in\\mathbb{R}^{N\\times N},\\quad O = PV\\in\\mathbb{R}^{N\\times d} $$这里 $\\mathrm{softmax}$ 是逐行计算的。\n算法的执行过程如下\n我们有第一个结论\nProposition 1 标准化 attention 前向传播时访问 HBM 的内存访问开销为 $\\mathcal{O}(Nd+N^2)$.\n证明：对于 attention, 我们需要从 HBM 中加载 $Q,K,V\\in\\mathbb{R}^{N\\times d}$, 然后输出 $O\\in\\mathbb{R}^{N\\times d}$ 并保存到内存中。\n首先我们需要计算 $S = QK^T$, 这一步需要加载 $Q,K$ 并将 $S$ 保存到 HBM 中，内存访问量为 $\\mathcal{O}(Nd+N^2)$.\n接下来，我们需要计算 $P = \\mathrm{softmax}(S)$, 这一步需要加载 $S$ 然后将 $P$ 保存到 HBM 中，内存访问量为 $\\mathcal{O}(N^2)$.\n最后，我们需要计算 $O = PV$, 这一步需要加载 $P$ 和 $V$ 然后将 $O$ 保存到 HBM 中，内存访问量为 $\\mathcal{O}(Nd+N^2)$.\n总的来说，标准化 attention 的内存访问开销为 $\\mathcal{O}(Nd+N^2)$.\nBackward Pass 标准 attention 反向传播过程如下图所示\nProposition 2 标准化 attention 反向传播时访问 HBM 的内存访问开销为 $\\mathcal{O}(Nd+N^2)$.\n证明：对于标准化 attention 的反向传播，我们需要从 HBM 中加载 $Q,K,V,dO\\in\\mathbb{R}^{N\\times d}$ , 然后输出 $dQ,dK,dV$ 并保存到 HBM 中。\n首先我们计算 $dV=P^TdO$, 这一步需要加载 $P,dO$ 并将 $dV$ 保存到 HBM 中，内存访问开销为 $\\mathcal{O}(Nd+N^2)$.\n接下来我们计算 $dP=dOV^T$, 这一步需要加载 $dO, V$ 并将 $dP$ 保存到 HBM 中，内存访问开销为 $\\mathcal{O}(Nd)$.\n然后我们计算 $dS$, 这一步需要加载 $P$ 并将 $dS$ 保存到 HBM 中，内存访问开销为 $\\mathcal{O}(N^2)$.\n对于 $dQ$ 和 $dK$ 的计算，内存访问开销都是 $\\mathcal{O}(Nd+N^2)$.\n因此，标准化 attention 的内存访问开销为 $\\mathcal{O}(Nd+N^2)$.\nMethod 作者在本节首先介绍了 flashattention 算法，然后作者证明了 flashattention 的正确性以及分析了复杂度。最后作者对 flashattention 进行扩展得到了 Block-sparse Flashattention.\nFlashattention attention 模块的输入是 $Q,K,V\\in\\mathbb{R}^{N\\times d}$, 输出是 $O\\in\\mathbb{R}^{N\\times d}$, 作者的目标是减少计算过程中的 HBM 访问次数\n作者分别使用了 tiling 和 recomputation 来解决 attention 前向传播和反向传播中的内存访问开销。flashattention 的核心思想是，我们将 $Q,K,V$ 分割成 block, 然后在 block 层面进行加载和计算。\nTiling 首先作者介绍了一下如何使用 tiling 来计算 softmax.\n给定一个向量 $x\\in\\mathbb{R}^{B}$, 其 softmax 计算方式如下\n$$ m(x) = \\max_i x_i,\\ f(x) = [e^{x_1-m(x)},\\dots,e^{x_B-m(x)}], \\ \\ell(x)=\\sum_if(x)_i, \\ \\mathrm{softmax}(x) = \\frac{f(x)}{\\ell(x)} $$如果我们现在有两个向量 $x^{(1)}, x^{(2)}\\in\\mathbb{R}^{B}$, 记 $x=[x^{(1)}, x^{(2)}]^T\\in\\mathbb{R}^{2B}$, 我们可以将 $\\mathrm{softmax}(x)$ 的计算分解为\n$$ \\begin{aligned} m(x) \u0026= \\max(m(x^{(1)}), m(x^{(2)}))， f(x) = [e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})]\\\\ \\ell(x) \u0026= e^{m(x^{(1)})-m(x)}\\ell(x^{(1)}) + e^{m(x^{(2)})-m(x)}\\ell(x^{(2)}), \\mathrm{softmax}(x) = \\frac{f(x)}{\\ell(x)} \\end{aligned} $$因此，如果我们额外记录 $m(x)$ 以及 $\\ell(x)$ 这两个量，那么我们可以每次仅计算 softmax 的一个 block. 具体细节见softmax.\nRecomputation 在反向传播过程中，一般我们需要存储 $S,P\\in\\mathbb{R}^{N\\times N}$, 需要的空间复杂度为 $\\mathcal{O}(N^2)$. 但是，通过存储 $O\\in\\mathbb{R}^{N\\times d}$ 以及 $(m,\\ell)$, 我们可以避免重新计算 $S,P$,这可以看做是 gradient checkpointing. 但是与 checkpointing 相比，因为 flashattention 减少了内存访问开销，因此其反向过程并没有变得更慢。\nAlgorithm 最终，flashattention 的算法如下图所示\nAnalysis Correctness 算法的正确性由定理 1 给出\nTheorem 1 flashattention (即算法 1) 输出 $O=\\mathrm{softmax}(QK^T)V$, 其时间复杂度为 $\\mathcal{O}(N^2d)$, 空间复杂度为 $\\mathcal{O}(N)$.\n证明：时间复杂度主要由矩阵乘法决定。在计算 $S_{ij}=Q_iK_j^T$ 时，所花费的 FLOPS 为 $\\mathcal{O}(B_rB_cd)$. 在计算 $\\tilde{P}_{ij}V_j$ 时，所花费的 FLOPS 为 $\\mathcal{O}(B_rB_cd)$. 循环一共执行了\n$$ T_cT_r = \\left\\lceil\\frac{N}{B_c}\\right\\rceil\\left\\lceil\\frac{N}{B_r}\\right\\rceil $$从而总的 FLOPS 为\n$$ \\mathcal{O}\\left(\\frac{N^2}{B_rB_c}B_rB_cd\\right) = \\mathcal{O}(N^2d) $$在 flashattention 的计算过程中，我们只需要保存 $(\\ell, m)$ 即可，因此需要的额外内存空间为 $\\mathcal{O}(N)$.\n接下来，我们可以证明 flashattention 的正确性，我们使用归纳法来证明。令 $j$ 满足 $0\\leq j\\leq T_c$, $K_{:j}\\in\\mathbb{R}^{jB_c\\times d}$, $V_{:j}\\in\\mathbb{R}^{jB_c\\times d}$ 分别为 $K$ 和 $V$ 的前 $jB_c$ 行。 $S_{:, :j}=QK_{:j}^T\\in\\mathbb{R}^{N\\times jB_c}$, $P_{:,:j}=\\mathrm{softmax}(S_{:,:j})\\in\\mathbb{R}^{N\\times jB_c}$, $m^{(j)}, \\ell^{(j)}, O^{(j)}$ 分别为 $m,\\ell, O$ 的第 $j$ 个元素。我们证明经过第 $j$ 次迭代后，HBM 中保存的是\n$$ m^{(j)}=\\mathrm{rowmax}(S_{:,:j})\\in\\mathbb{R}^N, \\ell^{(j)}=\\mathrm{rowsum}(\\exp(S_{:,:j}-m^{(j)}))\\in\\mathbb{R}^N, O^{(j)} = P_{:,:j}V_{:j}\\in\\mathbb{R}^{N\\times d} $$当 $j=0$ 时，上面的结果显然成立。现在我们假设对某个 $j=0,\\dots, T_c-1$ 上面的结果成立，我们需要证明对 $j+1$ 也成立。\n首先\n$$ m^{(j+1)}=\\max(m^{(j)}， \\tilde{m}) = \\max(\\mathrm{rowmax}(S_{:,:j}), \\mathrm{rowmax}(S_{:,j:j+1}))=\\mathrm{rowmax}(S_{:,:j+1}) $$接下来\n$$ \\begin{aligned} \\ell^{(j+1)} \u0026= \\exp(m^{(j)}-m^{(j+1)})\\ell^{(j)} + \\exp(\\tilde{m}-m^{(j+1)})\\tilde{\\ell}\\\\ \u0026=\\exp(m^{(j)}-m^{(j+1)})\\mathrm{rowsum}(\\exp(S_{:,:j}-m^{(j)})) + \\exp(\\tilde{m}-m^{(j+1)})\\mathrm{rowsum}(\\exp(S_{:,j:j+1}-\\tilde{m}))\\\\ \u0026= \\mathrm{rowsum}(\\exp(S_{:,:j}-m^{(j+1)})) + \\mathrm{rowsum}(\\exp(S_{:,j:j+1}-m^{(j+1)}))\\\\ \u0026= \\mathrm{rowsum}(\\exp(S_{:,:j+1}-m^{(j+1)})) \\end{aligned} $$最后，我们计算 $O^{(j+1)}$ 得到：\n$$ \\begin{aligned} O^{(j+1)} \u0026= \\mathrm{diag}(\\ell^{(j+1)})^{-1}(\\mathrm{diag}(\\ell^{(j)})\\exp(m^{(j)}-m^{(j+1)})O^{(j)}+\\exp(\\tilde{m}-m^{(j+1)})\\exp(S_{:,j:j+1}-\\tilde{m})V_{:,j:j+1})\\\\ \u0026= \\mathrm{diag}(\\ell^{(j+1)})^{-1}(\\mathrm{diag}(\\ell^{(j)})\\exp(m^{(j)}-m^{(j+1)})P_{:,:j}V_{:,:j}+\\exp(-m^{(j+1)})\\exp(S_{:,j:j+1})V_{:,j:j+1})\\\\ \u0026= \\mathrm{diag}(\\ell^{(j+1)})^{-1}(\\mathrm{diag}(\\ell^{(j)})\\exp(m^{(j)}-m^{(j+1)})\\mathrm{diag}(\\ell^{(j)})^{-1}\\exp(S_{:,:j}-m^{(j)})V_{:,:j}+\\exp(-m^{(j+1)})\\exp(S_{:,j:j+1})V_{:,j:j+1})\\\\ \u0026= \\mathrm{diag}(\\ell^{(j+1)})^{-1}(\\exp(-m^{(j+1)})\\exp(S_{:,:j}))V_{:,:j}+\\exp(-m^{(j+1)})\\exp(S_{:,j:j+1})V_{:,j:j+1})\\\\ \u0026= \\mathrm{diag}(\\ell^{(j+1)})^{-1}( \\begin{bmatrix} \\exp(S_{:,:j}-m^{(j+1)}) \u0026 \\exp(S_{:,:j}-m^{(j+1)}) \\end{bmatrix}\\begin{bmatrix} V_{:,:j} \\\\ V_{:,j:j+1} \\end{bmatrix}\\\\ \u0026= \\mathrm{softmax}(S_{:,:j+1})V_{:,:j+1} \\end{aligned} $$因此上面的结果对 $j+1$ 也成立，从而 flashattention 的结果对 $j=0,\\dots,T_c$ 都成立。\nForward Pass of flashattention 第一个问题是如何提高 softmax 计算的效率，作者的做法先先计算 normalization constant 然后再分别计算不同的 column.\n给定 $Q,K,V\\in\\mathbb{R}^{N\\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下\n$$ S = QK^T\\in\\mathbb{R}^{N\\times N},\\quad P = \\mathrm{softmax}(S)\\in\\mathbb{R}^{N\\times N},\\quad O = PV\\in\\mathbb{R}^{N\\times d} $$我们有 $S_{ij}=q_i^Tk_j$, 这里 $q_i$ 和 $k_j$ 分别是 $Q$ 和 $K$ 的第 $i$ 列以及第 $j$ 列， normalization constant 定义为：\n$$ L_i = \\sum_{j=1}^N \\exp\\left(q_i^Tk_j\\right) $$对任意 $i$, 计算 $L_i$ 只需要 $\\mathcal{O}(N)$ 的空间复杂度。\n令 $v_j$ 是 $V$ 的第 $i$ 列，则输出 $O$ 的第 $i$ 列 $o_i$ 为\n$$ o_i = P_{i:}V = \\sum_{j=1}^N P_{ij}v_j = \\sum_{j=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}v_j $$这个过程中，对任意 $i$, 计算 $o_i$ 也只需要 $\\mathcal{O}(N)$ 的空间复杂度。\n因此，在 $L_i$ 已经计算好的情况下，我们可以在 $\\mathcal{O}(N)$ 的空间复杂度下计算 $o_i$.\n最终，flashattention 的 forward pass 过程如下图所示\n接下来，作者分析了 flashattention 的内存访问开销。结论如下\nTheorem 2 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\\leq M\\leq Nd$. 则 flashattention 前向传播的内存访问开销为 $\\Theta(N^2d^2M^{-1})$.\n证明：由 Algorithm 1（或者 Algorithm 2）可以知道，$K$ 和 $V$ 的每一个元素都只需要从 HBM 中加载一次，而每一次外层循环都会从 HBM 中加载一次 $O$ 和 $Q$, 因此总的 HBM 访问次数为 $\\mathcal{O}(Nd+NdT_c)=\\mathcal{O}(NdT_c)$.\n接下来，我们给出每一次内层循环的内存访问开销，这是由 SRAM 的大小决定的。由于我们需要 SRAM 可以存储 $K_j\\in\\mathbb{R}^{B_c\\times d}$ 以及 $V_j\\in\\mathbb{R}^{B_c\\times d}$ ，我们的 block size 需要满足\n$$ B_cd = \\mathcal{O}(M) \\Rightarrow B_c = \\mathcal{O}\\left(\\frac{M}{d}\\right) $$同理，对于 $O$ 和 $Q$, 我们有\n$$ B_rd = \\mathcal{O}(M) \\Rightarrow B_r = \\mathcal{O}\\left(\\frac{M}{d}\\right) $$最后，我们还需要 SRAM 可以存储 $S_{ij}\\in\\mathbb{R}^{B_r\\times B_c}$, 因此\n$$ B_rB_c=\\mathcal{O}(M) $$这样，\n$$ B_c = \\mathcal{O}\\left(\\frac{M}{d}\\right), B_r=\\mathcal{O}\\left(\\min\\left(\\frac{M}{d},\\frac{M}{B_c}\\right)\\right)=\\mathcal{O}\\left(\\min\\left(\\frac{M}{d},d\\right)\\right) $$从而\n$$ T_c = \\frac{N}{B_c} = \\mathcal{O}\\left(\\frac{Nd}{M}\\right) $$最终，总的内存访问开销为\n$$ \\mathcal{O}(NdT_c) = \\mathcal{O}\\left(\\frac{N^2d^2}{M}\\right) $$一般来说, $d$ 的大小为 $64-128$, $M$ 的大小为 $100 KB$ 左右, $d^2\u0026laquo; M, 因此 flashattention 的内存访问开销远小于标准化 attention 的内存访问开销。\n作者还证明 flashattention 的内存访问开销是一个下界，即\nProposition 3 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\\leq M\\leq Nd$. 则不存在一个对任意 $M\\in[d,Nd]$ 都可以在 内存访问开销为 $\\Theta(N^2d^2M^{-1})$ 的条件下完成 attention 计算的算法。\n证明可以用反证法，基本思想是加载 $Q,K,V$ 的 HBM 访问次数至少为 $\\mathcal{O}(Nd)$.\nBackward Pass of flashattention 第二个问题是能否在线性空间复杂度下计算 attention 的反向传播过程。\n首先我们记损失函数为 $\\phi$, 然后令 $\\phi$ 对 $O,Q,K,V$ 的梯度分别为 $dO,dQ,dK, dV\\in\\mathbb{R}^{N\\times d}$, 我们的目标是计算 $dQ, dK, dV$.\n$dV$ 的计算是最容易的，我们有 $dV=P^TdO$, 因此\n$$ dv_j = \\sum_{i=1}^N P_{ij}do_i = \\sum_{i=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}do_i $$由于我们已经计算了 $L_i$, 因此，$dv_j$ 只需要 $\\mathcal{O}(d)$ 的空间复杂度。\n接下来，注意到 $dP=dOV^T$, 因此我们有\n$$ dP_{ij} = do_i^Tv_j $$计算的空间复杂度也是要 $\\mathcal{O}(N)$ 的\n注意到 $P_{i:}=\\mathrm{softmax}(s_{i:})$, 且 $y=\\mathrm{softmax}(x)$ 的 Jacobian 是 $\\mathrm{diag}(y)-yy^T$ (推导过程见 softmax), 我们有\n$$ dS_{i:} = (\\mathrm{diag}(P_{i:})-P_{i:}P_{i:}^T)dP_{i:} = P_{i:} \\odot dP_{i:} - (P_{i:}^TdP_{i:})P_{i:} $$我们定义\n$$ D_i = P_{i:}^TdP_{i:}= \\sum_{j=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}do_i^Tv_j = do_i^T\\sum_{j=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}v_j = do_i^To_i $$$D_i$ 的空间复杂度也只需要 $\\mathcal{O}(N)$.\n则\n$$ dS_{i:} =P_{i:} \\odot dP_{i:} - D_iP_{i:} $$我们有\n$$ dS_{ij} = P_{ij}dP_{ij} - D_iP_{ij} = P_{ij}(dP_{ij}-D_i) $$注意到 $S_{ij}=q_i^Tk_j$, 我们有\n$$ dq_i = \\sum_{j=1}^N dS_{ij}k_j = \\sum_{j=1}^NP_{ij}(dP_{ij}-D_i)k_j = \\sum_{j=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)k_j $$因此计算 $dq_i$ 的空间复杂度为 $\\mathcal{O}(d)$.\n同样的，\n$$ dk_j = \\sum_{j=1}^N dS_{ij}q_i = \\sum_{j=1}^NP_{ij}(dP_{ij}-D_i)q_i = \\sum_{j=1}^N\\frac{\\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)q_i $$其空间复杂度为 $\\mathcal{O}(N)$.\n总之，attention 的反向传播过程所需要的空间复杂度为 $\\mathcal{O}(N)$.\n作者发现有两点可以改进：\nattention mask 不需要存储，我们只需要保存 forward pass 时的输入，然后在 backward pass 时重新生成即可，这样只需要 $\\mathcal{O}(N)$ 的空间复杂度。 计算 softmax 的梯度是，如果使用公式 $D_i=P_{i:}^TdP_{i:}$ 来计算的话，由于 $P_{i:}\\in\\mathbb{R}^N$, 可能会导致超过 SRAM 的内存使用限制，因此，我们可以使用 $D_i=do_i^To_i$ 来避免这个问题，其中 $o_i\\in\\mathbb{R}^d$. 最终，flashattention 的 backward pass 过程如下图所示\n经过前面的分析，flashattention 的反向传播的时间复杂度为 $\\mathcal{O}(N^2)$, 空间复杂度为 $\\mathcal{O}(N)$.\nTheorem 5 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\\leq M\\leq Nd$. 则 flashattention 反向传播的内存访问开销为 $\\Theta(N^2d^2M^{-1})$.\n定理的证明与 Theorem 2 基本一致，我们此处不再赘述。\nBlock-sparse Flashattention 当 attention 具有 block sparsity 的性质时，作者提出了 blck-sparse flashattention 来进一步提高 attention 的计算效率。\n给定 $Q,K,V\\in\\mathbb{R}^{N\\times d}$, 以及一个 mask $M\\in\\{0,1\\}^{N\\times N}$, 我们需要计算\n$$ S = QK^T\\in\\mathbb{R}^{N\\times N},\\quad P = \\mathrm{softmax}(S\\odot \\mathbb{1}_{M})\\in\\mathbb{R}^{N\\times N},\\quad O = PV\\in\\mathbb{R}^{N\\times d} $$其中当 $M_{kl}=1$ 时， $(S\\odot \\mathbb{1}_ {M})_ {kl}=S_ {kl}$, 否则 $(S\\odot \\mathbb{1}_ {M})_{kl}=0$.\nBlock-sparse attention 的算法如下所示\nProposition 4 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\\leq M\\leq Nd$. 则 block-sparse attention 的内存访问开销为 $\\Theta(Nd+N^2d^2M^{-1}s)$, 其中 $s$ 是 block-sparse mask 中的非零 block 的比例\n证明与 Theorem 2 的证明是类似的，总的内存访问开销为 $\\mathcal{O}(Nd+NdT_c)$, 但是在计算的过程中，由于 mask 矩阵的 block-sparsity, 我们实际上只需要计算一小部分 $M_{ij}\\neq0$ 的情况，因此最终的内存访问开销为\n$$ \\mathcal{O}\\left(Nd+\\frac{N^2d^2}{M}s\\right) $$可以看到，attention mask 的 sparsity 越高，block-sparse flashattention 的效率也就越高。当 $N$ 非常大时，$s 通常为 $1/\\sqrt{N}$ 或者 $N^{-1}\\log N$, 从而最终的内存访问开销为 $\\mathcal{O}(N\\sqrt{N})$ 或者 $\\mathcal{O}(N\\log N)$.\n作者对比了以下 block-sparse flashattention 和 flashattention 的效率对比，结果如下图所示\nExperiment 作者通过实验验证了 flashattention 的有效性，如下表所示\nAttention Standard FlashAttention GFLOPs 66.6 75.2 HBM R/W (GB) 40.3 4.4 Runtime (ms) 41.7 7.3 可以看到，尽管 flashattention 相比于标准化 attention 需要更多的算力，但是由于其内存访问开销更少，所以最终的运行时间大有了大幅度降低\n作者还探究了 block size 对 flashattention 性能对的影响，实验结果如下图所示\n可以看到，随着 block size 增加，循环次数降低，内存访问开销也逐渐降低。但是当 block size 充分大 ( $\u003e 256$) 之后，运行时间就会被别的因素所限制，并且过大的 block size 可能会导致 SRAM 的内存溢出\n作者首先在 BERT 和 GPT-2 上验证了 flashattention 的表现，BERT 的实验结果如下表所示\nBERT Implementation Training time (minutes) Nvidia MLPerf 1.1 $20.0\\pm1.5$ FlashAttention (ours) $17.4\\pm1.4$ GPT-2 的实验结果如下表所示\nModel implementations OpenWebText (ppl) Training time (speedup) GPT-2 small - Huggingface 18.2 9.5 days (1.0 ) GPT-2 small - Megatron-LM 18.2 4.7 days (2.0 ) GPT-2 small - FlashAttention 18.2 2.7 days (3.5 ) GPT-2 medium - Huggingface 14.2 21.0 days (1.0 ) GPT-2 medium - Megatron-LM 14.3 11.5 days (1.8 ) GPT-2 medium - FlashAttention 14.3 6.9 days (3.0 ) 实验结果显示，flashattention 比 Huggingface 快 3 倍左右，比 Megatron 快 1.7 倍左右\n训练速度：实验显示，flashattention 在 BERT 上，比 MLPerf 1.1 快 $15\\%$, 在 GPT-2 上比 HuggingFace 快 3 倍，比 Megatron 快 1.8 倍 准确率：flashattention 是第一个在 Path-X 上比随机表现更好的 transformer 模型；block-sparse flashattention 是第一个在 Path-256 上比随机表现更好的的 sequence model Conclusion 作者提出了 flashattention, 一个通过优化标准 attention 内存访问效率来提高 attention 计算效率的方法，作者详细介绍了算法设计的原理与证明，并通过实验证明了结果的有效性。\nReferences arxiv ","date":"2025-08-21T11:32:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-flashattention/","title":"Notes on flashattention"},{"content":"作者提出了 StreamingLLM, 一个基于 attention sink 来提高 sliding window attention 在超长上下文场景下表现的方法。\nIntroduction 已有的基于 softmax attention 的架构的问题在于很难扩展到长上下文的场景，主要原因有两点：\nKV cache 会随着序列长度增加而商城，从而提高 decoding 的 latency 序列长度超过预训练的 context length 之后，模型表现会急剧下降 为了解决这个问题，已有的方法可以分为三类：\nlength extrapolation: 使用 RoPE 或者 AliBi 等方法来扩展 LLM 的 context length, 这类方法的问题是扩展的上下文长度仍然有限，对于 streaming 的场景作用有限 context window attention: 扩展 LLM 的上下文长度，如 flash attention 等来降低 attention 的计算和内存开销。这类方法也是只在有限的上下文场景下 work Improving LLMs’ Utilization of Long Text: 更好利用长上下文的数据 基于已有的工作的发现，作者提出了本文研究的核心问题：\n如何在不损失模型表现和效率的情况下，提高模型在无限长上下文场景下的表现。\n为了解决这个问题，作者首先分析了 sliding window attention 的不足，作者发现，sliding window attention 在超过 KV cache size 之后，表现也会急剧下降。作者通过实验发现，sliding window attention 表现急剧下降的原因在于 attention sink, 也就是模型损失了对于初始 token 的关注，从而导致模型表现下降。\n基于 attention sink, 作者设计了 StreamingLLM, 用于提高 sliding window attention 在长上下文场景下的表现，结果发现，模型的表现有了大幅度的提升。\n作者还进一步在预训练阶段加入了 sink token 充当初始 token, 进一步提高模型的表现。\nMethod Attention Sink 作者首先探究了一下 softmax attention 以及 sliding window attention 性能下降的节点，实验结果如下图所示\n可以看到，softmax attention 性能急剧下降的节点为 pre-training 的 context length; 而 sliding window attention 性能急剧下降的节点为 KV cache size.\n接下来，作者分析了一下不同 layer 的 attention 分布情况，如下图所示\n可以看到，初始的 2 层 layer 里 attention logits 的分布比较均匀。但是在后续的 layer 里，第一个 token 的权重都大幅度上升。\n作者分析原因认为，sliding window attention 在超过 KV cache size 之后性能急剧下降的主要原因是初始 token 不再参与 softmax 的计算，这导致了 softmax 的计算出现了比较大的变化，从而模型的表现开始下降。\n为了探究初始 token 对最终模型表现的影响因素是语义层面还是位置层面的，作者将初始的 token 替换为 \\n, 并比较了模型的表现，结果如下表所示\nLlama-2-13B PPL (↓) 0 + 1024(Window) 5158.07 4 + 1020 5.40 4\u0026quot;\\n\u0026quot;+1020 5.60 可以看到，把初始的四个 token 替换为 \\n, 并不影响模型最终的表现，这说明是初始 token 的位置信息在发挥作用。\n作者接下来探究了一下模型架构的影响，实验结果如下表所示\nCache Config 0+2048 1+2047 2+2046 4+2044 8+2040 Falcon-7B 17.90 12.12 12.12 12.12 12.12 MPT-7B 460.29 14.99 15.00 14.99 14.98 Pythia-12B 21.62 11.95 12.09 12.09 12.02 Cache Config 0+4096 1+4095 2+4094 4+4092 8+4088 Llama-2-7B 3359.95 11.88 10.51 9.59 9.54 可以看到，不同的模型架构都存在这个问题，这说明 sliding window attention 的影响与架构无关。并且，作者认为，使用初始 4 个 token 就可以有效的避免模型的性能下降，进一步增加初始 token 的数量不会有进一步提升。\n作者分析 attention sink 出现的原因在于，\n初始的 token 对于后续所有的 token 都是可见的，因此其会携带一些信息 在预训练阶段，模型并没有一个一致的初始 token 来标注起始信息，这导致模型会默认使用第一个 token 来储存一些信息。 为了解决这个问题，作者就提出了缓存初始 token 的方法，具体做法就是，在 sliding window attention 的基础上，我们还会加上初始 token 的信息，作者展示示意图如下所示\n也就是说，我们初始 token 始终会参与计算（论文中初始 token 数量为 4），然后我们会维持一个大小为 3 的 KV cache 队列来进行最终 sliding window attention 的计算，这样，每次计算 attention 的时候，我们就会使用 $\\# \\text{iniital token} + \\# \\text{sliding window token}$ 这么多的 token 来计算 attention. 作者对比了不同 attention 的计算方式，如下图所示\n前面是在 inference 阶段进行优化的，作者现在进一步探究在 pre-training 阶段加入 attention sink 参与训练对模型表现的影响。\n[[softmax-off-by-one]] 提出了我们应该加入一个 zero sink token, 其计算公式如下\n$$ \\mathrm{softmax}_1(x)_i = \\frac{\\exp(x_i)}{1 + \\sum_{j=1}^N \\exp(x_j)} $$这里 $x\\in\\mathbb{R}^N$ 是输入的序列。我们可以将 sink token 视为一个 key 以及 value 都是 0 向量的特殊 token.\n在本文中，作者使用了一个可学习的 sink token. 作者对比了原始 softmax attention, 使用 zero sink attention, learnable sink attention 三种方法的表现，结果如下表所示\nCache Config 0+1024 1+1023 2+1022 4+1020 Vanilla 27.87 18.49 18.05 18.05 Zero Sink 29214 19.90 18.27 18.01 Learnable Sink 1235 18.01 18.01 18.02 可以看到 zero sink 仍然需要一部分初始 token 来维持模型的表现。作者在论文中推荐使用 learnable sink.\nExperiments 作者首先验证了 StreamlingLLM 在不同架构上的表现，结果如下图所示\n实验结果显示，StreamingLLM 可以扩展到 4M 的上下文\n接下来，作者探究了以下在 Pretraining 阶段加入 learnable sink token 对模型表现的影响，结果如下图所示\n可以看到，加入 sink token 之后对模型的表现没有显著影响。并且，模型在下游任务上的表现与标准的 softmax attention 表现差不多。\n作者还对 StreamlingLLM 进行了可视化，结果如下图所示\n作者进一步评估了 StreamingLLM 在下游任务上的表现，我们主要关注一下 ARC 上的表现，结果如下图所示\n可以看到，full attention 出现了 OOM error, 而 sliding window attention 虽然避免了 OOM 的问题，但是其表现非常差。而 StreamingLLM 则进一步提高了 Sliding Window attention 的表现。\nConclusion 作者在本文中提出了 StreamingLLM, 一个在 Sliding window attention 中加入 sink token 来避免超过 cache size 之后模型表现急剧下降的问题。作者详细介绍了 attention sink 现象以及解决方法。\nReferences Efficient Streaming Language Models with Attention Sinks ","date":"2025-08-20T10:16:35+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-streamingllm/","title":"Notes on StreamingLLM"},{"content":"openAI 发布了 gpt-oss 大语言模型，包含 120B-A5.1B 以及 20.9B-A3.6B 两个 size, 作者强调了模型的 instruction following, tool use, 以及 adaptive thinking 能力\nMethod Architecture gpt-oss 系列是一个基于 MoE transformer 架构的 LLM. 架构中交替使用 sliding window attention 和 full attention, sliding window size 为 128 token, 架构图如下所示\n模型的配置如下表所示\nModel 120B 20B Date 2025/8/5 2025/8/5 # Total Parameters 116B 20B # Activated Parameters 5.13B 3.61B # MoE Layers 36 24 Hidden Dim 2880 2880 MoE Intermediate Dim 2880 2880 Attention GQA+sliding window GQA+sliding window Attention Head Dim 64 64 Attention bias True True # Attention Heads 64 64 # Key-Value Heads 8 8 # Experts (total) 128 32 # Experts Active Per Token 4 4 # Shared Experts 0 0 QK-Norm No No 在架构上，gpt-oss 做的主要改变有：\nQ, K, V projection layer, expert layer, routing layer 都使用了 bias 修改了 expert layer 中 SwiGLU 的定义 attention 中额外使用了一个 attention sink SwiGLU 大多数模型使用的基于 SwiGLU 的 MLP 定义如下\n$$ y = W_2(W_3x \\odot \\mathrm{SwiGLU}(W_1x)) $$其中 $\\mathrm{SwiGLU}(x)=x\\odot\\mathrm{sigmoid}(x)$, 在 gpt-oss 模型中，作者首先定义了两个常数 $\\alpha=1.702$, $\\mathrm{limit}=7.0$, 然后 SwiGLU MLP 的定义如下\n$$ \\begin{aligned} o_1\u0026=W_1x+b_1,\\\\ o_3\u0026=W_3x+b_3\\\\ o_1\u0026=\\mathrm{clamp}(o_1,\\max=\\mathrm{limit})\\\\ o_3\u0026=\\mathrm{clamp}(o_3,\\min=-\\mathrm{limit},\\max=\\mathrm{limit})\\\\ o_3\u0026= o_3\\odot \\mathrm{sigmoid}(\\alpha\\cdot o_3)\\\\ o_3\u0026= (o_1+1)\\odot o_3\\\\ y \u0026= W_2o_3 \\end{aligned} $$Attention Sink 作者使用了 attention sink 来避免 window attention 在超过 kv cache size 之后，表现大幅度下降的问题。\nQuantization 为了降低模型的内存占用量，作者使用了 PTQ 来训练 MoE 的权重，使用的精度为 MXFP4, 这样每个参数由 4.25 bits 来表示。最终，模型的参数存储格式如下：\n通过这个流程，gpt-oss-120B 可以部署在 80GB 内存的 GPU 上，gpt-oss-20B 可以部署在 16GB 内存的 GPU 上。模型各部分参数量如下表所示\nComponent 120b 20b MLP 114.71B 19.12B Attention 0.96B 0.64B Embed + Unembed 1.16B 1.16B Active Parameters 5.13B 3.61B Total Parameters 116.83B 20.91B Checkpoint Size 60.8GiB 12.8GiB 这里在计算激活参数的时候，没有没有考虑 embedding 的参数量。\nPre-training 预训练细节不多，主要是使用了 flash attention 进行加速计算，使用了 Triton 进行了 kernel 的优化，gpt-oss-120b 训练了 120M H100-hours, gpt-oss-20B 的训练时间是 gpt-oss-120Bd 的十分之一左右\nPost-training post-training 的数据包括 coding, math 以及 science 等，主要使用 RL 进行训练\n作者介绍了以下 post-training 使用的格式，即 harmony chat format. 角色的优先级如下所示\n1 System \u0026gt; Developer \u0026gt; User \u0026gt; Assistant \u0026gt; Tool 作者还加入了 channels 来限制可以使用的信息，比如使用 analysis 来表示 CoT tokens, 使用 commentary 来表示 function calling 等，一个具体的例子如下：\nEvaluation gpt-oss 系列的表现如下表所示\nConclusion 作者提出了 gpt-oss 系列大语言模型，gpt-oss 在架构上与已有的主流模型架构如 Qwen3, DeepSeek-V3 等都有一定区别\nReferences technical report ","date":"2025-08-19T16:14:56+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gpt-oss/","title":"Notes on gpt-oss"},{"content":"作者提出了 QK norm, 一个解决 softmax 注意力权重不稳定的 scaling 算法。\nProblem Definition Softmax 可以用于将 logits 转化为一个概率分布，但是 softmax 问题是输入的微小差别会对输出产生巨大影响，甚至会 mask 掉其他信号。因此我们需要选取合适的缩放因子，来解决 softmax 的极端值问题。SDPA 的可视化结果如下图所示\nMethod 作者首先提出了 QKNorm, 其使用了一个可学习的 scaling 参数来控制 $QK^T$ 的范围，进而让 attention 的 pattern 更加分散。\n作者首先回顾了以下已有的进展，主要是三点：\nFixNorm: 将 word embedding 限制为单位长度 PerNorm: 使用 Pre-norm 替换 Post-norm ScaleNorm: 使用 $\\ell_2$ normalization 替换 LayerNorm, 并乘以一个可学习的 scaling 参数。 作者基于这三点进行了改进，改进后的 attention 定义如下\n$$ \\mathrm{softmax}\\left(g\\cdot \\hat{Q}\\hat{K}^T\\right)V $$其中,\n$$ \\hat{Q} = [\\frac{q_1}{\\|q_1\\|_2},\\dots,\\frac{q_m}{\\|q_m\\|_2}], \\hat{K} = [\\frac{k_1}{\\|k_1\\|_2},\\dots,\\frac{k_n}{\\|k_n\\|_2}] $$是对原始的 $Q, K$ 按列进行 $\\ell_2$ normalization 得到的结果， $g$ 是一个可学习的参数，其初始化值为\n$$ g_0 = \\log_2(L^2-L) $$这里 $L$ 是训练数据 $97.5$ 分位。\n使用这种动态缩放之后，attention 的分布变得更加分散了，结果如下图所示：\nReferences Query-Key Normalization for Transformers ","date":"2025-08-13T16:12:11+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qk-norm/","title":"Notes on QK-Norm"},{"content":"智谱 AI 提出了 GLM4.5, 包含 GLM4.5 和 GLM-4.5-Air,两个 MoE LLM. 模型大小分别为 355B-A22B 和 106B-A12B, GLM4.5 主要关注 agentic, reasoning 以及 coding 三个领域。\nIntroduction 作者认为，通用模型有三个关键能力，即 ARC：\nAgent: 与外部工具以及真实世界进行交互 Reasoning: 解决数学和科学领域的复杂问题 Coding: 解决真实世界软件工程相关问题 已有的商业模型如 o1/o3, Claude Sonnet 4 已经在 ARC 上达到了非常好的表现，但是开源模型仍然比较稀缺\n基于这个目标，作者就提出了 GLM4.5 和 GLM-4.5-Air, 来统一完成三个不同的目标。\nMethod Pre-training Architecture GLM-4.5 是一个基于 MoE 架构的 LLM, 架构与 DeepSeek-MoE 相似，作者做了如下几点改变：\n在 MoE layer 中，使用了 loss-free balance routing, 然后使用了 sigmoid function 作为 routing score 的 normalization. 与 Kimi-k2 和 DeepSeek-V3 相比，作者降低了 head dimension, 提升了 number of layers. 作者认为更深的模型更有利于提高模型的 Reasoning 表现 attention 上，作者使用了 GQA, 对于 #RoPE, 作者使用了 partial RoPE, 只旋转每个 token 的前半部分， 作者还将 attention heads 的个数增加到了 2.5 倍，作者发现增加 attention heads 可以提高模型的 Reasoning 表现 作者还使用了 QK-Norm 来防止 attention logits 爆炸 作者还使用了一个 MoE layer 作为 MTP layer 来支持 speculative decoding. 模型与 DeepSeek-V3 和 Kimi-k2 的对比如下\nModel GLM-4.5 GLM-4.5-Air Step 3 Kimi K2 Date 2025/8/8 2025/8/8 2025/7/25 2025/7/28 # Total Parameters 355B 106B 316B 1043B # Activated Parameters 32B 12B 38B 32B # Dense Layers 3 1 5 1 # MoE Layers 89 45 56 60 # MTP Layers 1 1 0 0 Hidden Dim 5120 4096 7168 7168 Dense Intermediate Dim 12288 10944 18432 18432 MoE Intermediate Dim 1536 1408 5120 2048 Attention GQA GQA MFA MLA Attention Head Dim 128 128 256 192 # Attention Heads 96 96 64 64 # Key-Value Heads 8 8 1 64 scoring sigmoid sigmoid softmax softmax # Experts (total) 160 128 48 384 # Experts Active Per Token 8 8 3 8 # Shared Experts 1 1 1 1 QK-Norm Yes No No No Pre-training Data 预训练数据包括四个方面\nWeb: 过滤低质量数据和使用模版产生的数据 Multilingual: 基于 webpages 和 Fineweb-2 Code: 基于 GitHub 和其他代码平台，作者使用了 [[Fill in the middle]] 来训练模型。 Math \u0026amp; Scirence: 训练一个 classifier 来给数据进行打分。 最终，预训练数据一共包括 23T token.\nPre-training Recipe 预训练包括 2 个阶段:\nPre-training: 使用网页数据进行训练 Mid-training: 加入 code, math, science 数据进行训练，在这个阶段，作者使用了 repo-level 的 code 数据，合成的 reasoning 数据以及长上下文数据。作者将模型上下文从 4K 扩展到 32K，然后在扩展到 128K. 作者在 pre-training 的时候使用了 random truncation, 在 mid-training 的时候使用了 best-fit packing 技巧\n训练时，与 Kimi-k2 一样，作者使用了 Muon 作为优化器。作者使用了 cosine decay schedule. batch size 从 16M token 到 64M token.\nPost-training Post-training 分为两个阶段：\nStage 1, Expert Training. 构建 agent, reasoning, General chat 三个 domain 的专家模型 Stage 2, Unified Training. 使用 self-distillation 来汇总多个模型的能力 训练框架如下图所示\nSFT 两个 stage 都由 SFT 开始，\n在 Stage 1 里，SFT 的目标是让 expert model 掌握初步的 chat, reasoning 以及 tool-use 的能力。作者使用了一小部分包含 CoT 的 SFT 数据进行训练 在 Stage 2 中，SFT 的目标是将不同的 expert model 蒸馏到一个模型中，作者使用了百万级的数据，包含 reasoning 任务和通用的 chat 数据，来训练模型的 hybrid reasoning 能力 在训练模型的 tool-use 能力是，作者发现，function call 在 code 场景下会出现混淆，提高了模型的学习成本。因此，作者的解决方法是使用了类似 XML 的 special token tags\nRecall 与之相反，Kimi-K2 认为模板应该尽可能简洁，因此 Kimi 采取了 TypeScript 作为 function call 的语言\n从专家模型进行采样是，作者进行了数据过滤。还对数据进行了分级结果发现，使用难题进行训练可以提升模型 $2\\%\\sim4\\%$ 的表现，多次采样也可以提高模型的表现\nAgentic SFT 数据的构建包括四个步骤：\nAgentic Framework and Tool Collection: 收集 MCP 和 tool API Task Synthesis: 合成不同的 agentic 任务 Trajectory Generation: 采样生成的 rollout Quality Filtering: 过滤低质量的数据 RL Reasoning RL 这个阶段使用了 GRPO 算法进行训练，与 DAPO 一样，作者去除了损失函数中的 KL divergence。\n首先，作者探究了课程学习对模型表现的影响，结果发现，课程学习可以有效提高模型的性能。因此，作者构建了一个 2 阶段的课程学习框架。实验结果如下图所示\n可以看到，在第二个阶段，模型可以进一步通过更难的题目获得提升。\n其次，作者探究了以下渐进式扩展模型上下文对模型表现的影响。DeepScaleR 认为，逐步提高模型的上下文长度，可以有效提高模型的表现。但是，本文确认为这种方法会损害模型的性能，原因在于，模型在 SFT 阶段的上下文长度就是 64K, 如果我们降低模型的上下文长度，这会导致训练数据分布不一致，从而影响模型的长上下文表现。因此作者直接在 64K 的上下文上进行训练。\n接下来，作者探究了以下采样温度对模型表现的影响，温度太低会导致模型探索能力下降，太高的话会导致输出质量下降。因此作者动态调整采样温度来平衡模型的性能以及探索能力。\n[!tip] Kimi-K2 认为随着 RL 训练的进行，我们应该逐步降低采样温度来稳定模型的表现\n最后，作者分析了以下 code 以及 Science RL 中的一些问题。对于 code RL, 作者发现，我们应该在 sequence 层面而不是 token 层面进行平均。对于 Science RL, 作者强调了高质量数据的重要性。实验结果如下图所示\nAgent RL 作者主要关注 web-search 以及 code generation 两个任务。对于 web-search, 作者构建了一个数据合成 pipeline, 用于生成 multi-step reasoning 的 QA 数据。构建过程包括基于知识图谱的 multi-hop reasoning 和 human-in-the-loop 的内容提取。对于 code generation, 作者基于 GitHub 的 PR 以及 issues 构建了 benchmark\nRL 的训练目标如下\n$$ \\mathcal{L}(\\theta) = \\mathbb{E}_{x\\sim\\mathcal{D}}\\left[\\frac1K\\sum_{i=1}^K(r(x,y_i) - \\bar{r}(x))\\right] $$其中 $(x,y_i)$ 是基于 $\\pi_{\\mathrm{old}}$ 采样的 trace, $\\bar{r}(x) = 1/k\\sum_{i=1}^Kr(x,y_i)$ 是平均的 reward. 计算损失时，只有模型的回答参与计算。\n作者发现，通过训练模型的 web-search 以及 code generation 能分，模型在 tool-use 以及 coding 任务上的表现也有了提升。作者还是用了 format penalty 来保证模型输出格式的正确性。如果格式不对的话，模型获得的奖励是 0\nRecall 在 GLM-4.V-Thinking 中，作者认为应该在 cold-start SFT 阶段让模型学会格式要求，而不是在 RL 阶段加入 format reward\n由于 agent RL 的训练比较耗时，为了提高训练效率。作者首先基于 SFT 模型进行 agent RL 训练，训练到一定步数之后，作者使用 self-distillation 来将能力蒸馏回 SFT model, 接下来再基于 Self-distillation 后的 SFT 模型来进行 agent RL 训练\nRecall Seed1.5-VL 和 LlaMA3.2 都提到了使用 multi-round SFT-RL 的形式来提高模型的表现\n作者还发现，随着交互轮数的提升，模型的表现也有相应提升。实验结果如下图所示\nGeneral RL General RL 用于提高模型的整体表现，解决潜在的问题以及提升关键能力，作者主要使用了 RLHF 和 RLAIF 两种方法\n对于 Holistic RL, 作者收集了 5000 条 prompt, reward 基于人类反馈和 AI 反馈。人类反馈用于训练一个 reward model, 对于 AI 反馈，作者构建了 scoring rubrics. 然后作者将两种反馈结合在一起\n对于 Instruction following RL, 作者构建了基于规则的奖励，reward model 的奖励以及 critical model 的奖励。实验结果显示，这种奖励方式可以有效降低模型的 reward hacking\n对于 function calling RL, 作者使用了 step-wise rule-based RL 来提高模型表现。对于 end-to-end multi-turn RL, 作者训练了一个 expert model 来蒸馏专家到模型。\n最后，对于 Pathology RL, 作者希望通过 RL 来解决潜在的问题，比如语言混合输出，重复输出以及格式错误等。作者构建了一批模型容易出错的数据，然后来训练模型。\nInfra 作者针对不同任务分别构建了不同的 scheduling 模式：\n对于通用 RL 任务，作者将 training engine 和 inference engine 放在一个 worker 来提高效率 对于 agentic RL 任务，作者将 training 和 inference engine 分开，来提高 data throughput 在训练时，作者使用了 BF16 精度，在推理时，作者使用了 FP8 精度来提高推理效率。\n针对 agentic RL 任务，作者还进行了优化。与 Kimi-k2 类似，作者让 inference engine 持续产出 rollout, 然后让 training engine 来更新模型权重，最后同步到 inference engine 上\nExperiments 整体表现如下图所示，GLm4.5 在 ARC benchmark 上的平均表现达到了第三名。\n具体来看，\n在 agentic benchmark 上, GLM4.5 仅次于 o3 的表现 在 coding benchmark 上，GLM4.5 次于 Claude Opus 4 和 Claude Sonnet 4, 排第三名 在通用能力上，GLM 人工对比 coding agent 能力的结果如下图所示\nConclusion 作者提出了 GLM4.5， 一个基于 MoE 架构的大语言模型系列，包含 GLM4.5(355B-A22B) 和 GLM4.5-Air(106B-A12B) 两个模型，作者详细介绍了模型的架构，训练，数据和评估。\nReferences arxiv ","date":"2025-08-13T12:27:48+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-glm-4.5/","title":"Notes on GLM-4.5"},{"content":"腾讯 ARC LAB 提出了 ARC-Hunyuan-Video-7B, 一个针对短视频理解和推理的视频多模态大模型。\nIntroduction 作者首先提出了 Structured video comprehension 的概念\nthe ability to decompose a video into its constituent events and narrative elements with temporal precision.\n视频信息包含了 dense visual elements, 比如文字等信息，还有 rich audio streams, 比如音效，音乐，再就是 rapid narrative pacing, 用于强调情感。\n已有的多模态大模型对于视频的细粒度理解和推理能力存在不足。尽管 Keye-VL 也是一个短视频理解多模态大模型，但是其并没有利用 audio 模态的信息。其他的 audio-visual LLM 虽然可以处理 audio 模态的信息，但是他们主要集中于通用场景的视频，这类视频的特点是：叙述节奏慢，信息密度低。因此，对于叙述节奏快，信息丰富的短视频，这些模型表现就比较差。\n基于以上这些问题，作者提出了 ARC-Hunyuan-Video-7B, 一个基于 Hunyuan-7B vision-language model 的短视频多模态大模型，其主要做了两点改进：\n加入了一个 audio encoder 用于提取 audio 信息，然后对 audio 信息与视觉信息进行同步 使用了一个 timestamp overlap 机制，将时间戳印在视频的帧上，帮助模型理解事件的时序信息 作者还涉及了多阶段的训练策略，来提高模型的表现。最后在测试时，对于一个 1 分钟的视频，模型在 H20 GPU 上进需要 10 秒就可以处理完毕。\nMethod Architecture 模型基于 Hunyuan-7B-VLM 开发得到，\nVisual Encoding: 作者将时间戳以 HH:MM:SS 的形式印在视频的帧上。视频的采样率为 1FPS, 超过 150s 的视频会进行均匀采样降低至 150s. 每一帧会 resize 到 $640\\times 640$, 最后每一帧输出 112 token Audio Encoding: 使用 Whisper 作为编码器。小于 300s 的视频，会按照 30s 为单位切分为不同的 chunk, 大于 300s 的视频，会分割为 150 个 chunk, 然后每个 chunk 只保留开头 2s 的内容。最后，使用 Whisper 进行特征提取，最后再通过 MLP 与 LLM 的特征空间进行对齐 Visual-audio Synchronization: 先对 audio token 进行 zero padding 与 Visual token 的个数进行对齐，然后将两者相加 Training Pre-training 作者首先保住了一些数据。具体流程就是，使用 Whisper-v3 来转录带时间戳的音频，然后使用 InternVL-2.5-8B [[InternVL-2.5]] 来生成细粒度的 caption. 作者还使用 CoT prompt 策略，让模型一步步生成时间的描述，创作者的想法以及如何吸引用户的 tag 等。\n预训练数据如下\nCategory data Video description and summary 4.5M short-form video\n0.2M public video Image caption and OCR 4.7M image-text pairs ASR 3.2M audio-text pairs Video temporal grounding 0.5M temporally grounding instances Video multi-granular caption 50K high-quality samples\n80K in-house videos 训练包含两个 stage:\nStage 1: 使用 ASR 数据训练模型接受音频输入的能力，为了避免视觉模态能力下降，作者还加入了 Image-text pair data 来训练，缺失的模态用 zero padding 来补齐 Stage 2: 使用全部数据进行训练，仅训练 MLP 和 LLM Post-training 作者首先进行了消融实验，探究 human-annotated 数据对模型表现的影响。作者收集了 140 条人类标注的数据，然后基于这些数据进行消融实验：\n直接使用人类标注数据进行 SFT, 模型表现变化不大 直接使用人类标注数据作为正样本，合成数据作为负样本进行 DPO, 模型表现变化也不大 作者分析原因认为，人类标注数据和合成数据之间存在 distribution shift.\n受 DeepSeek-R1 启发，作者决定使用 GRPO 算法来提高模型的表现，训练任务包含两个：\nmulti-dimensional Multi-choide QA: 提高模型的视频理解能力 Temporal video grounding: 提高模型的时序感知能力 Stage Data Module Description SFT MCQ:\n- 460K open-ended QA\n- 70K MCQA\n- 20K QA\nGrounding:\n- 10K academic\n- 5K real-world\nGeneral:\n- 45K description\n- 12K caption MLP+LLM 提高指令跟随能力 Cold Start SFT - 90K MCQA\n- 18K temporal grounding\n- 20K open-ended QA\n- 15K summarization\n- 3K chapter-level captioning MLP+LLM 初步激活模型的 reas 能力 RL - 100K MCQ\n- 35K temporal grounding LLM 提升模型的 reasoning 能力 SFT - 25K human-annotated subjective question\n- 100K MCQ with CoT\n- 50K temporal grounding with reasoning traces - 使用人类标注数据进一步提高模型的能力 Experiments 作者构建了一个评估短视频理解能力的 benchmark: ShortVid-Bench, 评估以下方面：\nTemporal Reasoning and Localization Affective Intent Classification Creator Intent Taxonomy Narrative Comprehension Humor \u0026amp; Meme Deconstruction Creative Innovation Analysis 对比的模型包括 Qwen2.5-VL-7B, Qwen2.5-omni-7B 和 Keye-VL-8B\n评估结果如下：\nConclusion 作者提出了 ARC-Hunyuan-Video-7B, 一个针对短视频的视频理解多模态大模型。作者详细介绍了模型的架构，训练数据以及评估。\nReferences arxiv Github ","date":"2025-08-12T10:57:57+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-arc-hunyuan-video-7b/","title":"Notes on ARC-Hunyuan-Video-7B"},{"content":"Google Research 在 23 年 12 月份提出了 Group Query Attention (GQA), 一个提升 multi-head attention 效率的方法。GQA 自 Qwen2 系列开始被应用。\nIntroduction Multi-head attention (MHA) 的问题在于 inference 阶段，每次 decoding，都需要重新加载 attention 模块中 query layer, key layer 和 value layer 的权重，而加载权重会受带宽限制。\n已有的工作有 MQA, 也就是我们把多个 head 的 key layer 以及 value layer 压缩成一个，这样对于 $h$ 个 head 的 attention，我们有 $h$ 个 query layer，$1$ 个 key layer 以及 1 个 value layer. 但是 MQA 的问题在于其会导致性能下降，而且训练过程会不稳定。\n因此，在本文中作者就作出了两点贡献：\n如何将一个 MHA 模型转化为一个一个 MQA 模型 提出了 Group Query Attention (GQA)，在保持模型性能的同时，提高计算效率 Method Uptraining 将 MHA 模型转化为 MQA 模型分为两步：\n将 MHA 权重转化为 MQA 权重 额外的预训练 具体来讲，作者使用了一个 mean pooling 的方法，来将不同 head 的 query layer 以及 key layer 的权重转化为 MQA 对应 layer 的权重。然后作者 pre-training 若干步来让模型适应新的结构。\nGQA GQA 的思路在于在 MHA 和 MQA 之间达到一个平衡，也就是说我们将 key layer 和 value layer 进行分组，每个组内共享一个 key layer 和 value layer, 我们假设有 $h$ 个 head，$G$ 个 group，那么\n$G=1$ 时，所有的 head 共享一个 key layer 和一个 value layer, 此时 GQA 等价于 MQA $G=H$ 时，每个 head 都有一个 key layer 和一个 value layer, 此时 GQA 等价于 MHA $1","date":"2025-08-07T18:08:36+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gqa/","title":"Notes on GQA"},{"content":"Google 在 2019 年提出了 multi-query attention (MQA), 用于解决 MQA 内存带宽瓶颈问题。\nMethod Background 对于 multi-head attention, 我们假设其 hidden size 为 $d$, 有 $h$ 个 heads, 每个 head 的 size 为 $d_h=d/h$, 输入 sequence 长度为 $n$, batch size 为 $d$. 则总的 arithmetic operations 为 $O(bnd^2)$. 总的内存访问量为 $O(bnd + bhn^2+d^2)$, 第一项是 $Q,K,V$ 的内存占用（$Q,K,V$ 分别是 query, key 和 value layer 的输出），第二项是 attention score 的占用，第三项是 query, key 和 value layer 的权重。\n因此，其 Memory Access Ratio (MAR), 也就是内存访问量 与 arithmetic operations 之比为\n$$ O\\left(\\frac1k + \\frac{1}{bn}\\right) $$对于现代的 GPU 来说，其一般算力比较强，但是内存访问带宽相对较慢，因此我们希望 MAR 越低越好，以充分发挥 GPU 的算力。\nMHA Analysis 在训练的时候，由于我们知道 ground truth sequence, 因此我们可以并行计算。但是在 inference 的时候，我们只能 token-by-token 进行计算，因此我们分析一下 token-by-token 场景下的 MAR\n我们整体的 arithmetic operations 还是 $O(bnd^2)$.\n但是，现在我们要调用 $n$ 次 multi-head attention, 因此我们总的内存访问量为 $O(bn^2d + nd^2)$, 第一项是 $K$ 和 $V$ , 第二项是 query, key 和 value layer 的权重。\n这种情况下，MAR 就变成了\n$$ O\\left(\\frac{n}{d} + \\frac{1}{b}\\right) $$当 $n\\approx d$ 或者 $b\\approx 1$ 时，MAR 就非常接近于 1，意味着内存带宽成了一个主要的瓶颈。为了解决这个问题，我们有两种做法：\n提升 batch size $b$, 也就是同时 inference 多次 降低 $K$ 和 $V$ 的大小 MQA MQA 的做法就是第二种，也就是降低 $K$ 和 $V$ 的大小，但是 $K,V$ 分别是 key 和 value layer 的输出，要降低输出大小，我们就必须改变 key 和 value layer 的 size。基于这个考虑，作者在所有的 head 上共享了一个 key 和 value layer，也就是说，原来\n1 2 self.k_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h) self.v_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h) 现在在 MQA 里，其变成了\n1 2 self.k_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h) self.v_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h) MQA Analysis 我们还是在 token-by-token 的场景下进行分析。\n我们整体的 arithmetic operations 还是 $O(bnd^2)$.\n调用 $n$ 次 multi-query attention 的总的内存访问量为 $O(bnd +bn^2d_h+ nd^2)$, 第一项是 $q$ , 第二项是 $K$ 和 $V$ , 第三项是是 query, key 和 value layer 的权重。\n此时，MAR 变成了\n$$ O\\left(\\frac{1}{d} + \\frac{n}{dh}+\\frac{1}{b}\\right) $$现在，我们就将 $n/d$ 这一项给降低了 $h$ 倍。如果我们的 batch size 足够大的话，理论上 MQA 应该能极大提高整体的计算效率。\nConclusion MQA 为了追求极致的内存带宽占用，选择使用单一的 key 和 value, 来极大提高 inference 的 decoding 效率，但是后来在 GQA 中验证发现，MQA 虽然非常高效，但是其表现比较差，这也是后来没有得以应用的原因。\nReferences Arxiv ","date":"2025-08-07T18:06:37+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-mqa/","title":"Notes on MQA"},{"content":"Kimi 提出了 Moonlight, 一个基于 Muon optimizer 训练得到的 16B-A3B MoE LLM. 作者详细介绍了如何 scale up muon optimizer.\nIntroduction Muon 验证了 Muon optimizer 在小语言模型 nanoGPT 上的表现，但是对于更大规模 LLM 的表现，尚未有人探究。因此 Kimi 就希望在大规模 LLM 上验证 Muon optimizer 的表现。作者主要进行了两点改进：\n加入 weight decay 调整了不同参数更新的 scale 基于改进后的 Muon optimizer, 其训练效率相比于 AdamW 提升了 2 倍。作者基于 Muon Optimizer 训练得到了 Moonlight, 一个 16B-A3B 的 MoE LLM.\n作者主要作出了三点贡献：\n探究了 weight decay 在 scaling Muon 时的作用 分布式 Muon optimizer 的实现 验证了 Muon optimizer 的 scaling law Method Background 作者首先介绍了一下 Muon optimizer, 给定步数 $t$, 参数矩阵 $W_{t-1}$, momentum $\\mu$, 学习率 $\\eta_t$ 以及目标函数 $\\mathcal{L}_t$, Muon optimizer 的更新方式如下：\n$$ \\begin{aligned} M_t \u0026= \\mu M_{t-1} + \\nabla\\mathcal{L}_t(W_{t-1})\\\\ O_t \u0026= \\mathrm{Newton-Schulz}(M_t)\\\\ W_t \u0026= W_{t-1} - \\eta_t O_t \\end{aligned} $$这里 $M_t$ 是 gradient 的 momentum, 初始化为 $M_0=0$. 在上面的更新公式中，Newton-Schulz 的作用是求解 $(M_tM_t^T)^{-1/2}M_t$. 令 $M_t=U\\Sigma V^T$ 为 SVD 分解， 我们有\n$$ (M_tM_t^T)^{-1/2}M_t = UV^T $$这是一个半正交矩阵，即 $(UV^T)^T(UV^T)=I$.\nNewton-Schulz 迭代的具体公式如下：\n$$ X_0 = \\frac{M_t}{\\|M_t\\|_F},\\quad X_k = aX_{k-1} + b(X_{k-1}X_{k-1}^T)X_{k-1} + c(X_{k-1}X_{k-1}^T)^2X_{k-1} $$其中，normalization 是为了保证 Newton-Schulz 的收敛性。 $a,b,c$ 是三个超参数，在 Muon 中设置为 $(a,b,c)=(3.4445, 4.7750, 2.0315)$.\nScaling up Muon 作者发现，尽管 Muon 在小规模场景下 work 的很好，但是大规模性场景下的收益就非常有限了。作者发现，这是因为模型的参数以及每一层输出的 RMS 变得很大，这可能会影响模型的性能。因此，作者就和 AdamW 一样使用 weight dacay 来避免这个问题，即\n$$ W_t =W_{t-1} - \\eta_t(O_t + \\lambda W_{t-1}) $$作者通过实验对比了 AdamW, vanilla Muon 和 Muon w/ weigth decay 三者的表现，实验结果如下图所示\n实验结果显示，尽管 vanilla Muon 手链最快，但是由于其权重增长很快，因此最后模型的表现不如 AdamW 和 Muon w/ weigth decay.\n接下来，作者分析了以下更新矩阵的 Root Mean Square (RMS), 结论是 Muon optimizer 的 RMS 与参数矩阵的形状相关：\nLemma For a full-rank matrix parameter of shape $[A, B]$, its theoretical Muon update RMS is $\\sqrt{1/\\max(A, B)}$.\n证明如下：通过 Newton-Schulz 迭代，我们得到 $O_t=UV^T$, 其中 $M_t=U\\Sigma V^T$ 是 SVD 分解，我们有\n$$ \\mathrm{RMS}(O_t) = \\sqrt{\\frac{\\sum_{i=1}^A\\sum_{j=1}^BO_{t,i,j}^2}{AB}}=\\sqrt{\\frac{r}{AB}} $$其中, $r=\\mathrm{rank}(M_t)$ , 这样就完成了证明。\n而 Adam 和 AdamW 的 RMS 都在 $1$ 附近。作者认为 RMS 也会影响模型表现：\n当 $\\max(A,B)$ 过大时，如 dense MLP matrix, 其更新就会变得很小，限制了模型的表现 当 $\\max(A,B)$ 过小时，如 GQA 中的 KV head 或者 DeepSeek-V3 中的 MLA, 更新又会变得很大，导致训练不稳定。 因此，作者就提出了一个 rescaling 的技巧，来消除 Muon optimizer 的影响。\n作者通过实验发现，AdamW 的 RMS 通常在 $0.2\\sim0.4$ 左右，因此，作者将 Muon optimizer 的更新设置如下\n$$ W_t = W_{t-1} - \\eta_t(0.2\\cdot O_t\\cdot \\sqrt{\\max(A,B)} + \\lambda W_{t-1}) $$基于这个改变， Muon 和 AdamW 可以共享学习率以及 weight decay 参数。\nDistributed Muon ZeRO-1 天然适合 AdamW, 因为 AdamW 都是 element-wise 进行计算的。但是 Muon 则需要梯度矩阵的全部信息。因此，作者就针对 ZeRO-1 进行适配， 提出了 Distributed Muon, 分布式版本将优化器的状态进行切分，然后加入了两个额外的操作：\nDP gather: 将 ZeRO-1 切分的梯度矩阵 gather 为一个完整的矩阵 Calculate Full Update: 对完整的梯度矩阵执行 Newton-Schulz 迭代 最终，Distributed Muon 的算法如下图所示\n最后，作者分析了一下 distributed Muon 和 distributed AdamW 的内存和算力占用：\n内存开销：Muon 只有一阶矩，而 AdamW 有二阶矩，因此 Muon 的额外内存开销为 AdamW 的一半。 通信开销：对于 ZeRO-1，通信开销来源于三个过程：All-Gather 参数 $P$ 用于前向传播, Reduce-Scatter 梯度 $G$ 用于反向传播, All-Gather 更新后的参数 $P$ 用于下一轮的前向传播。AdamW 不引入额外通信，所以其每个参数的通信量为 $4+4=8$, 分别代表 $G$ 和 $P$ 的通信量。而 Muon 则需要额外的一次通信来得到 full matrix, 因此每个参数通信量为 $4+4+2=10$, 分别代表 $P, G$ 和 full matrix. 也就是说，分布式 Muon 的通信量最高为 AdamW 的 $1.25$ 倍。实际上由于我们使用 multiple DP, 这个比例会更接近于 $1.0$. latency：Distributed Muon 相比于 AdamW latency 更高，这是因为 Muon 需要进行 DP gather 以及计算 Newton-Schulz 迭代。但实际上，latency 很小，因为 Newton-Schulz 迭代只需要迭代 5 次，并且 optimizer 的 end-to-end latency 相比于 forward-backward 过程是可以忽略的。一些额外的技巧也可以降低 latency. 实际在训练的过程中，作者发现 Distributed Muon 相比于 AdamW 并没有太明显的 latency.\nExperiments Scaling Law of Muon 作者分析了一下 Muon Optimizer 的 scaling law, 实验结果如下图所示\n实验结果表明，在最优设置下，Muon Optimizer 只需要 $52\\%$ 的 FLOPs 就可以达到 AdamW 的表现\nPretraining with Muon 作者分贝使用 AdamW 和 Muon 训练模型，然后评测了以下模型在不同 benchmark 上的表现，结果如下图所示\n可以看到，在相同的设置下，Muon optimizer 的表现更好。\nDynamics of Singular Spectrum Muon optimizer 的核心思想就是让比较难更新的方向也能被更新到，本节作者就探究了 Muon 是否满足这个性质，作者对参数矩阵进行 SVD 分解，然后定义 SVD entropy 如下\n$$ H(\\sigma) = -\\frac{1}{\\log n}\\sum_{i=1}^n\\frac{\\sigma_i^2}{\\sum_{j=1}^n\\sigma_j^2}\\log\\frac{\\sigma_i^2}{\\sum_{j=1}^n\\sigma_j^2} $$作者对 SVD entropy 可视化如下\n可以看到，Muon optimizer 的 SVD entropy 比 AdamW 更大，这说明 AdamW 的更新方向更多更广，验证了 Muon optimizer 的核心思想\nSFT with Muon 作者还在 SFT 阶段验证了 Muon optimizer 的有效性。实验结果如下图所示\n结论主要有两个：\n预训练阶段与 SFT 阶段使用不同的优化器时，模型表现没有明显区别 SFT 阶段使用 Muon 可以达到与 AdamW 差不多的表现，但是最好还是在 pre-training 阶段使用 Muon Conclusion 作者探究了如何 scale up Muon Optimizer. 通过改进，作者在 16B-A3B 的 MoE LLM 上验证了 Muon Optimizer 的性能。实验结果发现，Muon Optimizer 的训练效率比 AdamW 提升了 2 倍左右。\n作者提出了三个未来可行的研究方向：\n目前 Muon 只能针对 2D 参数进行优化，其他参数仍然依赖于 AdamW 优化器，是否可以使用 Muon 优化所有参数？ Muon optimizer 可以理解是 spectral norm 下的 steepest descent 方法，如何将其扩展到 Schatten norm 是一个可以研究的方向 实验里提到，预训练和 SFT 阶段使用不同的 optimizer, 表现不是最优的，如何解决这个因为不同 optimizer 导致的性能差距是一个需要解决的问题。 References Muon is Scalable for LLM Training ","date":"2025-08-07T10:49:32+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-moonlight/","title":"Notes on Moonlight"},{"content":"腾讯混元提出了 Hunyuan-Large, 一个 389B-A52B 的 MoE LLM, 上下文长度为 256K.\nIntroduction Hunyuan-Large 主要在三个方向进行了改进：\n使用了更高质量的合成数据：模型使用了 7T 的预训练数据，其中包含了 1.5T 的合成数据 优化了模型的架构：作者提出了 KV cache compression, recycle routing, expert-specific learning rate scaling 策略来提高模型的表现 探究了 MoE 模型的 scaling law: 作者探究了 MoE 模型的 scaling law Pre-training Architecture Hunyuan-Large 是一个基于 MoE 的 transformer 架构，attention 部分使用了 GQA, position encoding 使用了 RoPE, MLP 的激活函数为 SwiGLU. 在 MoE layer 中，Hunyuan-Large 使用了 shared experts. 最终，模型的配置如下图所示\nKV Cache Compression 为了减少 KV cache 的内存开销，作者使用了两个技巧：\nGQA: 通过共享 KV projection 的参数，来减少内存访问次数 [[CLA]]: 在相邻的 layer 中共享 KV cache, 来进一步压缩 KV cache 在 Hunyuan-Large 中，作者将 GQA 的 group size 设置为 8, 然后相邻的 2 层 layer 共享 KV cache.\n假设输入的 batch size 为 $B$, sequence 长度为 $L$, layers 个数为 $\\ell$, attention heads 个数为 $h$, KV heads 个数为 $h_{kv}$, 每个 head 的 hidden size 为 $d_h$, 则每一层的 GQA 需要缓存 $K,V\\in\\mathbb{R}^{B\\times _{kv}\\times L\\times d_h}$， KV cache 的总占用为\n$$ 2\\times B\\times h_{kv}\\times L\\times d_h \\times \\ell \\times 2=4BLh_{kv}d_h\\ell $$第一个 $2$ 是因为同时缓存 K 和 V, 第二个 $2$ 是因为一般使用 bfloat16 数据格式。\n对于 CLA, 因为连续两层共享相同的 KV cache，因此结果除以 2; 对于 MHA, $h_{kv}=h$; 对于 MQA, $h_{kv}=1$. 最后，KV cache 的内存占用如下表所示\nAttention Mechanism KV Cache Memory MHA $4BLhd_h\\ell$ GQA $4BLh_{kv}d_h\\ell$ MQA $4BLd_h\\ell$ CLA $2BLhd_h\\ell$ GQA+CLA $2BLh_{kv}d_h\\ell$ 可以看到，使用 GQA+CLA 之后，模型的 kv cache 占用相比于 MHA 变成了\n$$ \\frac{2BLh_{kv}d_h\\ell}{4BLhd_h\\ell}=\\frac{1}{16} $$也就是说，Hunyuan-Large 的 KV cache 内存占用下降到了 MHA 的 1/16.\nExpert Routing Strategy 作者采用了 shared expert + activated expert 的形式，其中包含 1 个 shared expert, 然后从 16 个专家里激活 1 个专家。\n为了解决 MoE 中 expert capacity 难以设定的问题，作者提出了一个 recycle routing 的策略，基本思想就是，当 activated expert 的容量超出限制时，会从其他没有超出容量限制的专家里重新进行激活。\nExpert Specific Learning Rate Scaling 作者使用 AdamW 作为优化器，作者探讨了如何设定学习率。基于之前的工作，最优的学习率与 batch size 相关：\n$$ \\epsilon_{\\mathrm{opt}}(B) = \\frac{2\\epsilon_{\\max}}{\\sqrt{\\frac{\\mathcal{B}_{\\mathrm{noise}}}{B}}+\\sqrt{\\frac{B}{\\mathcal{B}_{\\mathrm{noise}}}}} $$这里 $\\epsilon_{\\max}$ 是 AdamW 的学习率, $\\mathcal{B}_{\\mathrm{noise}}$ 是训练速度与数据使用效率的一个平衡因子。\n但是，在 MoE 模型中，不同专家处理的 token 是不一样的。基于 load balancing loss, shared expert 和 activated expert 处理的 token 个数比例大概是 $n :1$, 其中 $n=16$ 是总的专家个数。因此，对于 shared expert, 作者使用 $\\epsilon_{\\mathrm{opt}}(B)$ 作为学习率，然后对于 activated expert, 作者使用 $\\epsilon_{\\mathrm{opt}}(B/n)$ 作为学习率。\nData 预训练数据包括收集和合成。收集的数据主要来自互联网，覆盖中英文两种语言。\n合成数据包括 4 个步骤：\ninstruction generation: 作者使用高质量的语料作为 seed, 然后生成多样的 instruction 覆盖不同的 domain Instruction evolution: refine 上一步生成的 instruction Response generation: 使用 specialized model 来生成回答 response filtering: 对生成的回答进行过滤 数据合成的流程如下图所示\ntokenizer 大小为 128K, 由 tittoken tokenizer 和额外的 28K token 组成。\nPre-training Recipe 作者首先探究了一个针对 MoE 模型的 scaling law. 结果发现，最优的激活参数量为 58.1B, training token 个数为 5.6T. 经过平滑之后，作者最终将模型的激活参数两定为 52B, 训练 token 数定为 $7T$.\n在训练时，作者将学习率分为了 3 个 stage:\nwarmup phase gradual decay phase concise annealing phase 上面的三个 stage 结束之后，作者加入了两个 stage 来扩展模型的上下文长度从 32K 扩展到 256K. 训练的数据包括 75% 的短文本和 25% 的长文本。两个 stage 训练的 token 数均为 $10B$ 左右。\nPost-training post-training 分为 SFT 和 RLHF 两个阶段。\nSFT SFT 数据副高 math, coding, logical reasoning 等 domain, 包含超过 1M 的数据。\nSFT 训练了 3 个 epoch, 学习率从 2e-5 降低到 2e-6, 为了避免 overfitting, 作者使用了 0.1 的 attention dropout 和 0.2 的 hidden dropout.\n[!tip] 作者发现，MoE 模型可以从 dropout 中学习到更多\nRLHF 作者使用 DPO 来进行 RLHF, 作者同时使用了 offline 和 online 的数据来进行训练，前者是收集的数据，后者是当前 policy 生成的数据。与 LLaMA 3 和 Nemotron-4 一样，为了提高训练稳定性，对于 chosen reponse, 作者使用了 SFT loss.\n作者还是用了 exponential moving average 策略来减少 reward hacking 现象，以及降低 alignment tax.\nExperiment 对于 base 版本，作者对比了 LLaMA 3, Mixtral, DeepSeek-V2, 实验结果如下图所示\nInstruction 版本的表现如下图所示\nConclusion 作者提出了 Hunyuan-Large, 一个 389B-A52B 的 LLM, 上下文长度为 256K. 作者详细介绍了模型的架构，数据和训练方式。\nReferences arxiv ","date":"2025-08-06T16:46:32+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-hunyuan-large/","title":"Notes on Hunyuan-Large"},{"content":"Qwen 提出了 Group Sequence Policy Optimization (GSPO), 一个针对 GRPO 进行改进的 RL 算法。GSPO 在 sequence 层面计算 importance ratio, 避免了 token-level 计算带来的训练不稳定性。\nIntroduction GRPO 的问题在于训练超大规模的 LLM 时，会出现 model collapse, Qwen3 和 MiniMax-01 均对 GRPO 算法进行了改进。\n在本文中，作者认为 GRPO 算法的问题在于 Importance sampling weight 的计算是有问题的，这导致了误差随生成回答长度累积，最后被 clipping 机制放大，从而导致模型崩溃。\n为了解决这个问题，作者提出了 GSPO, 一个在 sequence 层面进行 Importance sampling 的 RL 算法，GSPO 还对 reward 进行 normalization 来保持训练的稳定性\nMethod Preliminary 对于一个大语言模型 $\\pi_{\\theta}$, 我们将 $x$ 记为 query, 将 $y$ 记为 $\\pi_{\\theta}$ 针对 $x$ 的 response, 即\n$$ \\pi_{\\theta}(y\\mid x) = \\prod_{t=1}^{|y|}\\pi_{\\theta}(y_t\\mid x, y_{","date":"2025-08-06T11:26:26+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gspo/","title":"Notes on GSPO"},{"content":"Muon (MomentUm Orthogonalized by Newton-Schulz) 是一个针对二维神经网络的优化器，它基于 SGD-momentum 改进，增加了一个 Newton-Schulz 的后处理步骤\nMethod Newton-Schulz (NS) 的目的是用一个正交矩阵近似一个给定矩阵，即\n$$ \\mathrm{Ortho}(G) = \\arg\\min_{O} \\{\\|O-G\\|_F: \\text{either } O^TO=I\\text{ or } OO^T=I\\} $$也就是说，NS iteration 将 SDG-moment 的更新矩阵替换为了“最近的” semi-orthogonal matrix. 这等价于将更新矩阵替换为 $UV^T$, 其中 $USV^T$ 是更新矩阵的 SVD 分解。\n[!tip] 作者观察到，对于 SGD-momentum 和 Adam 来说，其在基于 transformer 的神经网络里有非常高的 condition number, 也就是 optimizer 仅在少数几个方向上进行优化。作者认为，通过正交化，可以有效提高模型在其他方向上的更新速度，进而提高模型表现\nNewton-Schulz 作者提到，正交化矩阵的方法有很多，比如 SVD 分解，但是其问题是非常慢，还有 Coupled Newton iteration, 但是其精度要求非常高，必须要在 float32 以上。\n作者因此使用了 Newton-Schulz iteration.\n令 $G=USV^T$ 是 SGD-momentum 更新矩阵的 SVD 分解，则基于系数 $(a,b,c)$ 的 NS iteration 定义如下：\n$$ \\begin{aligned} G' \u0026= aG + b(GG^T)G + c(GG^T)^2G\\\\ \u0026= (aI+b(GG^T)+c(GG^T)^2)G\\\\ \u0026= (aI+bUS^2U^T+cUS^4U^T)USV^T\\\\ \u0026= U(aS+bS^3+cS^5)V^T \\end{aligned} $$也就是说，如果我们定义五次多项式函数 $\\phi(x)=ax+bx^3+cx^5$, 然后执行 $N$ 次 NS iteration, 则我们得到 $U\\phi^N(S)V^T$, 其中 $\\phi^N$ 代表 $\\phi$ 复合 $N$ 次。\n为了保证 NS iteration 收敛到 $\\mathrm{Ortho}(G) = UV^T$, 我们必须保证两点：\n$S$ 的值，也就是 $G$ 的奇异值必须在区间 $[0,1]$ 上 $\\phi$ 必须满足 $\\phi^N\\to 1$, $N\\to\\infty$, $\\forall x\\in[0,1]$. 为了满足第一个条件，我们可以对 $G$ 进行 rescale, 即 $G\\gets G/\\|G\\|_F$, rescale 不影响最终的结果，即 $\\mathrm{Ortho}(G) = \\mathrm{Ortho}(cG)$.\n对于 $\\phi(x)$, 我们有很多选择，比如我们定义 $(a,b,c):=(2,-1.5,0.5)$ 就得到如下结果\nCoefficient Optimization 尽管 $(a,b,c):=(2,-1.5,0.5)$ 已经满足了第二个条件，但是我们还是想进一步优化，优化的方向主要有两个：\n让 $a$ 尽可能大，这是因为 $\\phi'(0)=a$ 控制了较小奇异值的收敛速率。 对于所有的 $x\\in[0,1]$, 我们希望 $\\phi^N(x)\\in[1-\\epsilon, 1+\\epsilon]$, $N\\to\\infty$. 这样 NS iteration 的结果与 $\\mathrm{Ortho}(G)$ 不会相差太远。 作者发现， $\\epsilon$ 可以设置为 $0.3$ 而不影响 Muon optimizer 的收敛性。因此，作者的目标现在是\n$$ \\begin{aligned} \\max\\quad \u0026a\\\\ \\mathrm{s.t.}\\quad \u0026\\lim_{N\\to\\infty}\\phi^N(x)\\in[0.7, 1.3] \\end{aligned} $$作者通过 ad-hoc gradient 方法求解得到一组数值解为 $(a,b,c)=(3.4445, 4.7750, 2.0315)$, 作者将这组数值应用于 Muon optimizer 中。迭代结果如下图，可以看到，当 $x\\approx0$ 时，函数变得更加陡峭。\n实验中，作者发现，仅需迭代五次，最终的结果就 work 的很好。作者还尝试了不同的多项式，结果发现并没有太大的提升。\nAlgorithm 最终，Muon Optimizer 的算法如下\n其中, NewtonSchulz5 算法伪代码定义如下\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 def newtonschulz5(G, steps=5, eps=1e-7): assert G.ndim=2 a, b, c = (3.4445, -4.7750, 2.0315) X = G.bfloat16() X /= (X.norm() + eps) if G.size(0) \u0026gt; G.size(1): X = X.T for _ in range(steps): A = X @ X.T B = b * A + c * A @ A X = a * X + B @ X if G.size(0) \u0026gt; G.size(1): X = X.T return X Analysis 本节作者分析了以下 Muon 的内存占用和算力开销。\n在 NS iteration 之前，Muon optimizer 和 SGD-moment 是一样的。\n对于 $n\\times m$ 的矩阵（假设 $m\\leq n$）， 首先 NS iteration 会进行转置，NS iteration 的每一步需要 $2(2nm^2+m^3)$ FLOPs, 其中括号前面的系数 $2$ 代表精度。因此，Muon 相比于 SGD momentum 需要的额外 FLOPs 为 $2T(2nm^2+m^3)$, 其中 $T$ 是迭代次数。\n使用 baseline 进行一次训练（前向 + 后向），所需要的 FLOPS 为 $6nmB$, 其中 $B$ 是 batch size. 因此，Muon 的 FLOP 开销至多为 $Tm/B$, 其中 $m$ 是模型的 hidden size, $B$ 是 batch size, $T$ 是 NS iteration 的步数。\n作者分别基于 nanoGPT 和 LLaMA-405B 进行验证，结果发现，Muon optimizer 带来的额外开销不足 $1\\%$.\n作者发信啊，使用 Nesterov-style momentum 可以比普通的 SGD-momentum 效果更好，因此作者在 muon 中使用了前者。\n作者还发现，对于 QKV layer，分别进行优化效果会更好。\nExperiments Limitation and Future Work Muon 仅被设计用于优化 2D 参数（因为涉及矩阵计算），其余的参数仍然需要 AdamW 等优化器参与。\n作者认为未来的工作有：\n能否 scale up Muon Optimizer 分布式优化 在 fine-tuning 和 RL 阶段使用 Muon Optimizer Conclusion 作者提出了 Muon optimizer，该优化器在 nanoGPT speedrun 上取得了 SOTA 的结果，作者详细介绍了优化器的工作原理。\nReferences Muon: An optimizer for hidden layers in neural networks ","date":"2025-08-05T11:10:51+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-muon-blog/","title":"Notes on Muon blog"},{"content":"Apple 在 7 月份发布了 AFM 技术报告，包括两个多语种多模态大模型，分别为 3B 和 xB, 一个面向 device, 另一个面向 server， 前者主要集中于效率，后者集中于表现。\nPre-training Architecture On-Device Model 对于 on-device model, 作者将模型分为两个 block, Block1 占 $62.5\\%$ 的 transformer layers, Block2 占 $37.5\\%$ 的 transformer layers. 但是，对于 Block2, 作者移除了 key, value projection, 对应的 KV cache 则直接从 Block1 中获取。通过这种方式，作者将 KV cache memory usage 减少了 $37.5\\%$. 并且，由于 Block2 不产生任何 key values, prefill stage 可以跳过这些计算，这样 TTFT 也可以减少 $37.5\\%$.\nServer Model 对于 server model, 作者对架构进行了改进，来提高效率。架构如下图所示\nParallel Track Transformer 作者提出了 Parallel Track (PT) Transformer 架构，PT-Transformer 将 transformer 模型分割多个小的 transformer, 作者将这些小的 transformer 称之为 track. 每个 track 包含多个 transformer block. 不同的 track 只会在输入和输出的时候进行交互，这样就能够减少同步的开销。作者讲这种模式称为 track parallelism.\nPT-MoE 为了进一步提高 server model 的效率，作者将 MoE 和 PT-transformer 结合在一起。具体的做法就是，每两个 transformer block 为一组，每组里包含一个 dense layer 和一个 MoE layer.\nInterleaving Global and Local Attention Layers 作者还设计额 interleaved attention 机制，也就是，将 transformer block 按照四个为 1 组，前面 3 个 block 使用 window attention, window size 为 4096 和 RoPE. 最后一个 block 使用 global attention layer 以及 NoPE. 作者认为，使用 NoPE 可以提高模型对长上下文的泛化性。\nRecall Qwen2.5-VL 的 ViT 使用的类似的做法，即 8 个 block 为一组，前面 7 个 block 使用 window attention, 最后一个 block 使用 full self attention.\nVision Encoder Vision encoder 包含 ViT 和 adapter 两个模块\n对于 ViT 来说，作者使用了 ViT 架构：\nserver model 使用了 1B 参数的 ViT-g on-device model 使用了 300M 参数的 ViTDet-L backbone 作者在 ViTDet 的基础上加入了 Register-Window 机制，这个机制用于编码一个 global register token 来与不同的 loca windows 进行交互。\n对于 adapter 来说，其包含了一个 transformer layer, 一个 linear projection layer, 一个 $3\\times 3$ 的 convolutional layer. 其中， linear projection 用于将 visual token 映射到 LLM 的特征空间，pooling layer 用于压缩 visual token 个数。\nData 主要包括 web data 和 image data 两部分\nimage data 部分：\nImage-Text Crawl Data: 包含 175M 图文交错数据，包含 550M images Synthetic Image Caption data: 5B image caption 数据 Text-Rich Image Data High-quality Domain-Specific Image-text Data: 包括 caption 数据， grounding 数据，table, chart, plots 数据以及 knowledge-required domains 的数据 Training Recipe text tokenizer 大小为 150K.\nVision encoder 的训练包含两个 stage:\n基于 CLIP 的方法，使用 6B的 image-text pair 数据进行训练，图片精度为 448, 作者还使用了 FLIP 来提高训练效率 使用一个 compact LLM, 同时训练 vsion encoder, adapter 和 compact LLM. 加入了更高质量的数据，图片精度为 672. LLM 的训练使用了 13.4T token\nPost-training SFT SFT 数据包括：\nGeneral knowledge Reasoning: 纯文本包括 math 和 reasoning, 多模态包括 STEM, math, CoT 数据 Text-Rich Image understanding: chart, table 数据 Multilingual OCR: OCR 相关数据 Text and visual grounding: grounding 数据 Multi-image reasoning: 多图推理数据 作者还基于 retrieval-based 方法来收集数据，具体做法就是给定一些 prompt, 然后通过一个 Image search pipeline 来进行检索。\n训练的时候，作者将图片精度从 672 提升到 1344, 处理方式就是将图片切分为四个子图，然后作者还加入了一个总蓝图。这样，vision encoder 的输入包括四个子图和一个 thumbnail 图.\n为了提高 on-device model 的效率，作者设置了三种模式：\nrapid mode: 图片精度为 224 balanced mode: 只有 thumbnail 图 high-resolution mode: 四个子图和一个 thumbnail 图 对于不同的 mode, 如果输入的是低精度图片，则 $50\\%$ 概率为 rapid mode; 如果输入的是高精度图片，则 $1\\%$ 的概率为 rapid mode. 对于其他数据，作者将 $20\\%$ 的数据设置为 balanced mode.\nRLHF 作者使用 RLOO 作为 RLHF 的算法。\nRL 的 infra 如下图所示\ninfra 主要由两个部分组成：\nTrajectory Generators: 生成轨迹并提供反馈 Policy updater: 更新 policy 训练时，作者首先训练了一个 reward model, 与 AFM-2024 相似，作者使用了一个 preference loss function 以及一个 single-sided grading 作为 regularization.\n数据包括以下类别:\ntext-only prompts Image-text prompts Math prompts Image-text STEM reasoning prompts 其中，前面两个使用 reward function 进行打分，后面两个基于 ruled-based verifier 进行打分\n作者还发现，人类的打分和 reward model 的发奋可能会出现 $20\\%\\sim30\\%$ 的偏差。为了解决这个问题，作者训练了一个单独的 reward model, 专门用于 prompt selection.\nTool Use 工具调用数据由于 Multi-turn 和依赖软件工具，比较难以收集。为了解决这个问题，作者设计了一个交互式标注平台，包括一个 agent 和一个工具执行环境。环境包括了工具和数据库，能执行工具调用并反馈。\n标注时，用户发起一个请求，然后 agent 自动执行工具调用，最后平台返回反正的轨迹。\nMultilingual 作者逐步增加模型对于新语言的理解能力。默认情形下，输入和输出的语种一致，但是包含 $0.4\\%$ 的跨语种数据。在 SFT 和 RLHF 阶段，英语和多语种数据的比例为 $80\\%:20\\%$.\nOptimization 作者使用了 QAT 来将 on-device model 压缩到 2 bits-per-weight, 使用 Adaptive Scalable Texture Compression (ASTC) 来 post-training 3.56 bits-per-weight 版本的 server model.\nQAT QAT 是一个在模型训练过程中模拟量化误差，从而提升模型量化后表现的方法。它解决了传统后量化方法精度损失较大的问题，是平衡模型性能用户效率的关键手段。\n训练时，作者通过修改权重 $W$ 来模仿量化：\n$$ \\tilde{W} = s\\left(\\mathrm{clamp}(\\lfloor \\frac{W}{s}+z\\rceil, q_{\\min}, q_{\\max}) - z\\right) $$其中, $s$ 是 scaling factor, $z$ 是 zero point, $q_{\\min}$, $q_{\\max}$ 是 quantization 的 range. 为了解决 rounding operation 不可微的问题，作者使用了 straight-through estimator 的方法来近似梯度。\n作者还提出了一个可学习的 scaling factor $f$ 用于计算 quantization scale, 计算方法如下所示\n$$ s = \\frac{f\\cdot \\max(|W|)}{q_{\\max}} $$作者通过精细设计 $f$ 的初始化来保证模型训练的 robust.\nASTC 对于 server model, 作者使用了 ASTC, 一个针对 GPU 图形纹理压缩的技术，来压缩模型权重。具体做法就是，模型训练好之后，作者对模型权重应用 ASTC, 然后对每个块进行预处理。存储时，每个块用 ASTC-HAR-ch 模式压缩为 128 位。最小值单独存储为 float16.\n推理时，GPU 硬件自动解压缩 ASTC 块，然后解压的权重最小值相加参与矩阵计算\nQuality Recovery Adapters 作者还是用 LoRA 来恢复量化模型的精度，并通过选择性压缩策略优化 ASTC 过程，在极小的算力开销下实现了接近全量微调的性能。\nEvaluation On-device model 表现如下\nModel MMLU MMMLU MGSM AFM On-Device 67.85 60.60 74.91 Qwen-2.5-3B 66.37 56.53 64.80 Qwen-3-4B 75.10 66.52 82.97 Gemma-3-4B 62.81 56.71 74.74 Gemma-3n-E4B 57.84 50.93 77.77 server model 表现如下\nModel MMLU MMMLU MGSM AFM Server 80.20 74.60 87.09 LLaMA 4 Scout 84.88 80.24 90.34 Qwen-3-235B 87.52 82.95 92.00 GPT-4o 85.70 84.00 90.30 Conclusion 作者提出了 AFM-2025 多模态多语种大语言模型系列，包括 on-device 和 server 两个版本，作者介绍了模型的架构，训练数据和训练方式。\nReferences Publication blog ","date":"2025-07-29T12:36:28+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-afm2025/","title":"Notes on AFM2025"},{"content":"Kimi-k2 是一个总参数为 1T, 激活参数为 32B 的 MoE 大语言模型，模型使用 15.5T token 进行训练，optimizer 使用了 MuonClip. 作者主要关注模型的 agent 能力\nIntroduction 作者首先强调现在 LLM 发展主要是 Agentic Intelligence, 也就是让模型自主感知，规划，思考和与环境交互。\n基于这个目标，作者就提出了 Kimi K2, 一个 1.04T 总参数，32B 激活参数的 MoE 模型，用于解决实现 agent Intelligence 中遇到的问题。作者主要进行了三点改进：\nMuonClip, 一个基于 Muon 的优化算法，来提高 Kimi K2 对 token 的利用效率以及提高训练的稳定性 大规模的 agentic 数据合成 pipeline: 作者构建了一个用于合成工具调用，agent 数据的 pipeline 通用的 RL 框架，作者将 RLVR 和 self-critic rubric reward mechanism 结合起来，用于提升模型的表现 Pre-training Architecture Kimi-K2 的架构与 DeepSeek-V3 相似，配置如下表所示\n指标 DeepSeek-V3 Kimi K2 $\\Delta$ # Layers 61 61 = Total Parameters $671\\text{B}$ $1.04\\text{T}$ $\\uparrow 54\\%$ Activated Parameters $37\\text{B}$ $32.6\\text{B}$ $\\downarrow 13\\%$ Experts (total) 256 384 $\\uparrow 50\\%$ Experts Active per Token 8 8 = Shared Experts 1 1 = Attention Heads 128 64 $\\downarrow 50\\%$ Number of Dense Layers 3 1 $\\downarrow 67\\%$ Expert Grouping Yes No - 与 DeepSeek-V3 相比，模型主要进行了以下改动：\n作者认为提高专家的稀疏性，可以有效提高模型表现。因此作者将专家个数提升了 50%. 为了降低模型在 inference 阶段的算力开销，作者将 attention heads 的个数降低了 50%. Sparsity Scaling Law 作者首先构建了基于 Muon 的 sparsity law, 这里 sparsity 定义为激活专家个数与总专家个数之比，即：\n$$ \\mathrm{sparsity} = \\frac{\\# \\mathrm{activated\\ experts}}{\\# \\mathrm{total\\ experts}} $$作者在小规模上的实验结果如下图所示\nObservation 实验结果表明，在相同算力（激活专家个数）下，模型的表现随 sparsity 提高而增加。\n但是，进一步提高 sparsity, 会让 infra 变的难以优化，因此在 Kimi-K2 里，将 sparsity 定义为 $48$.\nNumber of attention heads 作者还分析了探究了 attention heads 的最优配置。DeepSeek-V3 将 attention heads 的个数设置为 layers 层数的 2 倍，来提高带宽利用率以及提高计算效率。但是，当上下文长度增加之后，attention head 变成了 computational bound. 作者对比了不同配置模型的表现，实验结果如下：\nObservation 实验结果表明，使用双倍 attention head 带来的收益是比较小的，只有 $0.5\\%$ 到 $1.2\\%$ 左右\n因此，作者在 Kimi-K2 中，奖 attention head 的个数设置为 $64$.\nData Kimi-K2 主要强调了 token efficiency, 即每个 token 对模型表现提升的贡献。token efficiency 越高，说明每个 token 对最终模型的贡献也就越高\n相比于 Kimi-k1.5, Kimi-K2 使用了一个 rephrasing pipeline, 来提高高质量 token 的利用率，作者在 knowledge 和 mathematics domain 上进行了实验。\n最终，Kimi-K2 的预训练数据包括了 15.5T token, 主要覆盖 Web text, code, mathmatics 以及 Knowledge 四个 domain. 大部分的数据处理与 Kimi-k1.5 相同。\nKnowledge Data Rephrasing 作者构建了一个 synthetic rephrasing 框架来提高 knowledge token 的利用率，框架主要包含以下几个模块：\nStyle and perspective-diverse prompting: 作者构建了一系列 prompt, 来让 LLM 从多个角度和风格来重新表述原始文本 Chunk-wise auto-regressive generation: 为了保持模型在长文档中的 coherence 以及避免信息损失，作者使用了一个 chunk-based rewriting 策略，也就是对每个 chunk 分别进行改写，然后将它门汇总在一起 Fidelity verification: 作者对原始文本和改写文本进行了 fidelity 检查来保证语义的一致性。 作者对比了以下三个设置对模型在 SimpleQA benchmark 上表现的影响，这三个设置分别是：\n原始数据集训练 10 epoch 改写数据一次，然后训练 10 epoch 改写数据一次，训练 1 epoch 实验结果如下表所示\n# Rephrasings # Epochs SimpleQA Accuracy 0 (raw wiki-text) 10 23.76 1 10 27.39 10 1 28.94 可以看到，改写的策略可以有效提高模型在 SimpleQA benchmark 上的表现\nMath Data Rephrasing 对于数学相关数据，作者基于 SwallowMath, 将数据改写成了学习笔记 (learning-note) 的风格。然后，作者还将其他语言的高质量文本翻译为英文。\nRecall 个人认为，数据改写在一定程度上也算是合成数据，对于合成数据这一块，微软提出的 phi 系列是一个比较好的参考\nMuonClip Optimizer Kimi-K2 使用了 Muon optimizer, 之前的实验结果说明了在相同的 compute budget 下，Muon optimizer 的表现超过了 AdamW. 也就是说，Muon optimizer 更加高效。\n但是，Muon optimizer 的问题是训练不稳定。为了解决这个问题，作者提出了 QK-clip, 一个 weight clipping 机制，来显示约束 attention logits. QK-clip 的机制是当输出的 logits 超过某一个阈值之后，就对齐进行截断。\n每个 head 的 attention 的计算公式如下\n$$ O = \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d}}\\right)V $$其中, $d$ 是 hidden size, $Q, K,V$ 分别是 query, key, value, 定义如下：\n$$ Q = XW_Q, K=XW_K, V=XW_V $$这里 $W_Q,W_K,W_V$ 是模型可学习的参数。\n作者定义每个 head 的 max logit 如下：\n$$ S_{\\max}^h = \\frac{1}{\\sqrt{d}} \\max_{X\\in\\mathcal{B}}\\max_{i,j} [QK^T]_{ij} $$最简单的做法就是直接进行截断，也就是\n$$ W_Q\\gets \\gamma^\\alpha W_q, W_K\\gets \\gamma^{1-\\alpha}W_K $$其中 $\\gamma=\\min(1, \\tau S_{\\max})$, 这里 $S_{\\max}=\\max_h S_{\\max}^h$ 是所有 head 对应 $S_{\\max}^h$ 的最大值。\n但是，实际中，作者发现只有少部分 head 会出现 logits 爆炸现象。为了提高计算效率，作者针对每个 head 单独进行 scaling, 也就是 $\\gamma_{h}=\\min(1, \\tau S_{\\max}^h)$. 对于 MLA 架构，作者仅在 unshared 模块使用 clipping:\n$q^c$ 以及 $k^c$, scaling factor 为 $\\sqrt{\\gamma_h}$ $q^R$, scaling factor 为 $\\gamma_h$ 最后，作者将 Muon, weight decay, RMS matching 和 QK-clip 汇总在一起，得到 Kimi-k2 使用的 MuonClip optimizer, 算法如下所示：\n接下来，作者对比了以下 Muon 和 MuonClip 两个优化器的表现，作者分别使用这两个优化器训练 53B-A9B 的 MoE 模型，实验结果如下\n实验结果表明，Muon 优化器在 1000 步左右就出现了 logits 爆炸的现象，但是 MuonClip 通过优化可以避免这个问题。作者还发现，MuonClip 可以减少 loss spike 的产生，整体的 loss 变化情况如下图所示\n作者还在附录中说明 QK-clip 并不影响最终的收敛性，并且，为了尽可能减少对模型训练的干扰，作者发现，仅在训练初期 QK-clip 被激活:\n在初始的 70, 000 步里，有 12.7% 的 attention heads 至少触发了一次 QK-clip, 并且将它们的 $S_\\max$ 降到了 100 以下 接下来的 70, 000 步里，QK-clip 就不再被激活 Infra Kimi-k2 使用了 16-way 的 PP, 16-way 的 EP 以及 ZeRO-1 Data Parallelism. 具体流程如下\n作者发现通过增加 warm-up mircro-batches, 我们可以有效重叠 EP 的 all-to-all communication 以及 computation. 但是，对于 [[DualPipe]], 其要求双倍的存储来保存参数和梯度，并且提高 PP 粒度会引入更多的 bubble, 因此 Kimi-K2 没有使用 DualPipe.\n为了减少 [[1F1B]] 的 PP 通信开销，作者在反向传播的过程中同时进行 PP 通信，来进一步重合通信与计算。\n作者还发现，使用更小的 EP group 可以提高整体的训练速度，因此作者将 EP 设置为 16.\n作者发现，剩余的 GPU 内存不足以存放 MoE 的 activation, 因此作者采取了三个办法解决这个问题：\nSelective recomputation: 作者对 LayerNorm, SwiGLU, MLA 的 up-projection 和 MoE 的 down-projections 等进行重新计算。 FP8 storage for insentive activations: 作者使用了 FP8 精度来存储 MoE up-projections 以及 SwiGLU. 作者发现使用 FP8 进行计算可能会导致性能下降，因此作者并没有使用 FP8 进行计算。 Activation CPU offload: 作者将其余的 activation 放在 CPU RAM 上，在 1F1B 的过程中，作者再进行加载 Training Recipe 模型上下文长度为 4096, 使用 MuonClip 进行优化，使用 WSD lr Scheduler.weight decay 为 0.1, global batch size 为 67M tokens.\n预训练结束之后，作者加入了一个 long-context activation stage. 这个阶段，作者使用了 400B 的 4K 上下文的 token 和 60B 的 32K 上下文的 60B token. 最后作者使用 YARN 将模型上下文长度扩展到 128K.\nPost-training SFT 数据的构建主要是基于：\nprompt 的多样性 Response 的质量 作者构建了一系列的数据生成 pipeline, 然后使用 Kimi-k1.5 来生成多样化的回答，最后再进行质量评估和过滤。这里，作者主要介绍了一下 tool-use 数据的构建\n受 ACEBench 启发，作者构建了一个模仿真实世界 tool-use 的数据合成 pipeline. pipeline 如下图所示\npipeline 主要包含三个阶段：\ntool spec generation: 作者基于 MCP tools 和 self-evolved 的数据合成策略来构建 tool repository. MCP tools 包括 3000+ 的 tools,合成了 20，000 个 tools Agent and task generation: 作者还用不同的 system prompt 以及不同的 tools combination 来构建对应的 agent. 然后对不同的 agent, 作者构建了对应的成功标准均，工具调用模式以及评估方式 Trajectory generation: 作者首先构建不同风格的 LLM, 然后作者使用一个 simulator 来执行 tool call 并给予反馈。 最后，作者对数据进行了过滤。\n作者还加入了 coding 以及软件工程等任务相关数据来进一步提高模型的 agent 能力\nRL RL 与 Kimi-K1.5 差不多。作者进一步提高了 RL 阶段的算力并做出了亮点改进：\n作者构建了类似 Gym 的框架，用于扩展 RL 的能力 作者加入了更多 RLVR 的任务 Data 数据主要包括以下几类：\nMath, STEM and logical tasks: 数据构建的原则为多样化和中等难度 Instruction following: 作者基于 hybrid rule verification 和 multi-source instruction generation 来合成复杂的指令跟随数据 Faithfulness: 作者训练了一个 judge model 来提供 reward Coding \u0026amp; Software Engineering: 作者从开源数据收集并合成了代码相关数据 Safety. 提高模型的安全性，防止 jailbreak Reward 作者使用了 self-critique rubric reward 的奖励机制。\n首先，对于 Kimi-k2 的回答，作者会使用另一个 kimi-k2 作为 critic，基于三个方面进行排序：\ncore rubric: AI 的核心价值观 prescriptive rubric: 避免 reward hacking human-annotated rubric: 特定的上下文 在训练的过程中，critic 也会基于 verifiable signals 进行 refine\nRL training RL 的训练目标与 Kimi-k1.5 相同\n$$ \\mathcal{L}(\\pi_\\theta) = \\mathbb{E}_{x\\sim \\mathcal{D}}\\left[\\frac1K\\sum_{i=1}^K \\left(r(x,y_i)-\\bar{r}(x) -\\tau\\log\\frac{\\pi_{\\mathrm{\\theta}}(y_i\\mid x)}{\\pi_{\\mathrm{old}}(y_i\\mid x)}\\right)^2\\right] $$其中 $\\bar{r}(x)=1/K\\sum_{i=1}^Kr(x,y_i)$ 是 sample response 的平均奖励。\n作者做了以下几点改进来提高模型在不同 domain 上的表现：\nBudget control: 作者针对不同任务分别设置了最大 token 限制，模型超过这个限制会受到惩罚。猜测应该是 length penalty 或者类似 Qwen3 一样，直接中断思考过程输出最终答案 PTX loss: 作者使用了 PTX loss 来提高模型对于高质量数据的利用率以及降低模型的过拟合 Temperature Decay: 作者发现，训练后期保持较高的采样温度会影响模型的表现，因此作者设置了一个 schedule, 来逐步降低采样温度。 RL Infra RL infra 与 Kimi-k1.5 类似。主要包含三个模块。其中 train engine 和 inference engine 两者互相切换，一个 engine 进行 training 的时候，另一个 engine 就进行 inference, 为了提高 engine 切换的效率，作者使用了 checkpoint engine 来传输更新后的参数。\nEvaluation 模型评估结果如下图所示\n评估结果显示，模型在 coding 和通用任务上的表现仅次于 Claude 4, 大部分 benchmark 上都是第二名甚至是 sota.\nConclusion 作者在本文中提出了 Kimi-K2, 一个总参数为 1B 的 MoE 大语言模型。作者在架构，数据和优化器上进行了创新。并且通过 post-training 显著提高了模型的 agent 以及 tool-use 能力\n作者发现模型主要存在的问题有：\nreasoning 任务过难或者 tool 的定义不清晰的时候，模型会使用很多 token 有时候工具调用可能会降低模型的表现 模型在 agentic coding 任务上的能力需要进一步提升 References Github ","date":"2025-07-24T10:56:50+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-kimi-k2/","title":"Notes on Kimi-k2"},{"content":"Keye-VL 是快手在 25 年 7 月份提出的一个 8B 的多模态大模型，其亮点为短视频理解能力。预训练包括 4 个 stage，使用了 600B token，后训练包括 2 个 stage，用于提升模型的 reasoning 和 non-reasoning 能力。\nIntroduction 作者首先回顾了已有的 MLLM 工作，然后强调理解短视频仍然是一个很难的任务，特别是要求模型基于 video 和 audio 来理解视频。因此，在本文中，作者提出了 Kwai Keye-VL，一个 8B 的多模态大模型，主要用于短视频理解任务。\nMethod Architecture Keye-VL 是一个标准的 ViT-MLP-LLM 的架构，其中，ViT 是 SigLIP-400M-384-14, MLP 是一个基于 SwiGLU 的 2 层 MLP，使用了和 Qwen2.5-VL 一样的 patch merge 方法，LLM 使用的是 Qwen3-8B，模型架构示意图如下\n作者针对 ViT 和 visual encoding 分别做了如下改进\nNaViT 作者实现了 native resolution ViT，来处理不同分辨率的图片。\n具体做法为，作者基于 SigLIP-400M-384-14 来初始化 ViT。\n然后，作者首先采用了 interpolation 来将 ViT 的 position encoding 扩展到不同的图片精度下面去。\n接下来，作者提出了 2D RoPE 来进一步提升 position encoding 的表现。\n最后，作者加入了 NaViT 的 packing 技巧来继续预训练 ViT.\n在 ViT 预训练的过程中，作者使用了 500B 的 token\nVisual Encoding 为了提升模型理解图片和视频的能力，作者针对图片和视频进行了进一步的处理。\n对于不同精度的图片，作者将最大 token 个数设置为 16384。\n对于视频，作者将每帧的 token 数限制在 $[128,768]$, 每个视频的最大 token 个数设置为 24576\n对于提取的 frames，作者重新计算了 FPS, 然后在 3D RoPE 中让时间维度与真实时间严格对齐。\nPre-training Data 预训练数据一共包括 600B token，覆盖了 6 个类别：\nImage caption: 包括中英文数据，来源为 LAION, DataComp 以及 Coyo. 作者基于 CLIP 来计算相速度，然后过滤掉相似度比较低的数据，作者还对数据进行了 re-caption，实验发现 re-caption 可以提高模型的细粒度图片理解能力 OCR \u0026amp; VQA: 数据包括开源数据和合成数据。合成数据包括 Synthesis 和 Rendering 两个方法，第一是基于 text-dense image 构建 OCR 数据以及基于 image-caption pair 数据使用 MLLM 构建 VQA 数据。第二是使用字体渲染工具合成高质量的 OCR 数据 Grounding \u0026amp; Counting: Grounding 数据主要包括 RefCoCo, VisualGenome, TolokaVQA, Counting 数据包括 PixMo. 作者仍然使用 CLIP 对数据进行过滤 Interleaved text-image data: 作者发现图文交错数据可以提供通用知识，并且还可以提高模型的视觉语言对齐能力，第三就是提升模型的泛化能力。作者主要从 academic PDF 以及结构化知识中提取对应的数据。作者基于 Garbled character recognition, low-resolution/broken image filtering 以及 text-image similarity validation 来保证数据的质量 Video understanding: 作者使用 Qwen2.5-omni 从 interleaved video-asr 来将视频数据转化图文交错数据， 然后基于 ASR 的结果进行 recaption，最后对每一帧进行 OCR. 作者还构建了 Frame-level re-ordering 以及 multiple video matching 两个任务来提高模型的上下文理解能力 Pure Text: 未提及 对于源数据，作者进行了数据清洗：\n使用 CLIP 对数据进行打分，然后过滤掉低质量的数据 使用开源的 MLLM 作为 discriminator 来选择高质量的数据 去重 Training Recipe 预训练包括 4 个 stage：\nStage 0: 使用 SigLIP 损失函数来继续训练 ViT Stage 1: cross-modal Alignment，仅训练 MLP Stage 2: multi-task pre-training, 解冻所有参数，使用 Multi-task 数据来训练模型 Stage 3: annealing, 在高质量数据集上进行 fine-tune，进一步提升模型的能力 作者发现，预训练后的模型在下游任务上的表现对训练数据配比非常敏感。为了解决这个问题，在最后一个训练阶段，作者使用了一个 merging 的技巧，来保持模型的能力。\nPost-training post-training 阶段一共包含了 2 个 step, 5 个 stage, 第一个 step 包含 2 个 stage，用于提升模型的 non-reasoning 能力。第二个 step 包含 3 个 stage, 用于提升模型的 reasoning 能力\nNo-reasoning Training 第一个 step 是 non-reasoning training, 包含了 SFT 和 MPO 两个 stage, 训练 pipeline 如下图所示\nSFT SFT 阶段一共使用了 5M 的多模态 QA 样本，为了提升数据的多样性，作者使用 TaskGalaxy 来将数据分类为 70,000 种任务类型，然后作者对每条数据，使用 MLLM 评估问题的难度，来过滤掉过于简单的问题。最后，作者请人类来进行标注，保证数据的可靠性。\nMPO 训练方面，作者使用了 MPO 进行训练。 数据方面，作者使用了：\n400,000 开源的数据，作者主要进行了去重以及过滤低质量的数据 50,000 偏好数据，基于 MM-RLHF 和 MMPR 等数据构建，然后构建高质量的 negative examples 10,000 条 self-improvement 样本：基于 benchmark 和人类反馈，使用 SFT model 的回答作为 chosen samples, 然后基于 reward model 或者 rule-based rewards 评估模型输出，选择分数最低的座位 rejected samples 90,000 纯文本样本： in-house data 30,000 人类标注样本：使用开源和闭源模型进行回答，然后请人类进行排序 Reasoning Training 第二个 step 是 reasoning training, 包含了 CoT Cold-Start, Mix-Mode RL 和 Iterative Alignment 三个 stage, 训练 pipeline 如下图所示\nCoT cold-start 作者收集了如下数据：\n330, 000 条 non-reasoning 样本：和第一个 step 的数据分布类似，但是不重合 230,000 reasoning 样本：作者构建了一个 data construction pipeline,用于保证 long-CoT 的正确性 20,000 automatic reasoning 样本：作者基于 MMPR, MM-Eureka 以及 OpenR1-math 来收集数据，基于这些模型来训练模型自动化决定是否要进行 reasoning 100,000 agentic reasoning 样本：训练模型的 \u0026ldquo;think with image\u0026rdquo; 能力。基于 3M QA pairs，作者使用 Qwen2.5-72B 来识别需要进行 manipulating 的问题，然后生成对应的代码。接下来作者构建了一部分 OCR 数据，让模型学会 constrast enhancement 或者 rotation 操作。最后对于数学问题，作者让 Gemini-2.5-Pro 生成对应的思考过程，再让 GPT-4o 将对应的计算转换为可执行的代码。 32, 000 Video data: 包含了 24,000 条 thinking samples 和 80,000 条 non-thinking samples 训练时，所有样本混在一起进行训练。作者认为，将日常使用的 SFT 数据和 reasoning 数据放在一起训练，可以保持模型在通用场景下的能力。\nMix-Mode RL 训练数据主要包括 4 个任务：\nMultimodal perception: 复杂文本识别和 counting 任务 Multimodal reasoning: MMPR 和 MM-Eureka Text-based mathematical reasoning: 数学推理问题 Agentic reasoning: 从 DeepEyes 中获取的 47,000 条样本 作者使用了 GRPO 来训练，reward 基于 MLLM 进行，包括最终结果和思考过程。\n作者还使用 RL 来提高模型的短视频理解能力。作者发现 RL 训练之后，模型的短视频理解能力有了大幅度的提升。\nIterative Alignment 这一步主要解决模型的重复性输出，或者 reasoning logic 不对的问题。作者使用了 rejection-sampling 数据，包括 instruction following, OCR, mathematics, charts, counting 等。\n作者基于 Rule-based score 和 model-based score 来进行打分，最后使用 MPO 算法进行训练。通过这个过程，模型的输出格式和动态思考能力都有了提升。\nEvaluation 作者首先评估了 ViT 的表现，主要有两点：\n在 SigLIP 的基础上加入 1D interpolation 之后，模型的表现所有下降，作者认为这是由于 1D 的 position encoding 无法识别 2D 的 patch 排列导致的 加入 2D RoPE 之后，ViT 与 SigLIP 的表现持平 接下来是 Keye-VL 在公开 benchmark 上的表现，如下图所示\n作者还构建了一个内部的 benchmark, 用于进一步评估模型的能力。\n已有 benchmark 的问题：\ncontamination 多语种覆盖不足：大部分 benchmark 都是英文的 任务和 domain 覆盖不足：大部分 benchmark 只考虑基本的 perception 和 reasoning 能力 任务难度和评估格式单调 构建 benchmark 的原则：\n在中文场景下的真实用户需求，open ended QA, 包括短视频理解能力 细粒度的评估 多样性高 没有 contamination 多角度评估策略: 正确性，相关性，理解性，流畅性和创造性 结果如下：\n分析：\nKeye-VL 在 OCR 等任务上的表现有所不足，其细粒度的识别能力也有所不足，会认错人。有时候还会忽略掉一些信息 描述 temporal action 时会出现不稳定性，模型对镜头的感知能力不足。需要准确定位时间等 在需要逻辑链条和数学计算时，模型能力不足，对于特定 domain 上的任务会出现事实性错误。写作时，模型倾向于输出通用的回答，而不是定制化的回答。 Discussion 作者讨论了两点关键发现：\nreasoning 和 non-reasoning 的数据可以互相促进彼此的表现，这与 ERNIE 4.5 的发现一致。 作者认为通过 mix-mode 的训练，模型在简单和复杂任务上的表现都可以提升，因此作者使用了混合数据来进行训练，结果发现效果很好。 Conclusion 本文中，作者提出了 Keye-VL 8B，一个短视频理解能力出色的多模态大模型，作者详细介绍了 pre-training 和 post-training. 其中，mix-mode training 可以有效提高模型的表现。\n作者认为 Keye-VL 有如下改进的地方：\n并没有优化 video encoder 或者是改进 video encoding 的策略 Keye-VL 的视觉感知能力有进一步的提升空间，其 \u0026ldquo;reasoning with image\u0026rdquo; 能力依然落后于领先的 reasoning model 使用一个额外的 MLLM 作为 reward model 会极大消耗算力，如何构建一个更可靠更高效的 reward model 需要进一步探索。 Reference Arxiv Github ","date":"2025-07-23T11:11:43+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-keye-vl/","title":"Notes on Keye-VL"},{"content":"本文中，我们介绍一下如何计算 LLM 的参数量。我们将基于 Qwen3 模型架构出发，对模型架构进行拆解，然后给出 LLM 参数量计算公式。\nDense Model 我们首先来看一下 Qwen3 的架构，如下图所示\n这里，Qwen3ForCausalLM 就是我们的的 LLM, 其关键代码如下\n1 2 3 4 class Qwen3ForCausalLM: def __init__(self, config): self.model = Qwen3Model(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 我们假设 vocab_size, 也就是词表大小为 $|V|$ (我们用 $V$ 表示词表), hidden_size 为 $d$, 则总参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3ForCausalLM}) =d|V| + \\mathrm{parameter}(\\texttt{Qwen3Model}) $$Qwen3Model 包含三个（含参数的）模块，分别是 nn.Embedding, Qwen3DecodeLayer 以及 Qwen3RMSNorm, 分别代表了输入 token 的 embedding layer, Transformer block 和对输出的 normalization. 其关键代码如下：\n1 2 3 4 5 6 7 class Qwen3Model: def __init__(self, config: Qwen3Config): self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) 其中，nn.Embedding 参数量与 lm_head 一样，都是 $d|V|$.\n对于 normalization, 现在大部分 LLM 用的都是 RMSNorm, 其定义如下：\n$$ \\mathrm{RMSNorm}(x) = \\frac{x}{\\sqrt{\\|x\\|_2^2+\\epsilon}}\\odot \\gamma $$其参数量为：$d$.\n如果说我们使用的是 LayerNorm, 则其定义如下：\n$$ \\mathrm{LayerNorm}(x) = \\frac{x-\\mathbb{E}[x]}{\\sqrt{\\mathrm{var}[x]+\\epsilon}}\\odot \\beta + \\gamma $$其参数量为： $2d$.\n因此，Qwen3Model 的参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3Model})=d|V| + N*\\mathrm{parameter}(\\texttt{Qwen3DecoderLayer}) + d $$这里第一项为 nn.Embedding, 第三项为 Qwen3RMSNorm， 第二项里，$N$ 代表 decode layer 的个数，也就是 config.num_hidden_layers.\nQwen3DecoderLayer 包含了四个模块，其关键代码如下\n1 2 3 4 5 6 class Qwen3DecoderLayer: def __init__(self, config: Qwen3Config): self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx) self.mlp = Qwen3MLP(config) self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) 因此，Qwen3DecoderLayer 的参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3DecoderLayer}) = 2d + \\mathrm{parameter}(\\texttt{Qwen3MLP}) + \\mathrm{parameter}(\\texttt{Qwen3Attention}) $$其中，第一项是两个 Qwen3RMSNorm 的参数。\n对于 Qwen3MLP, 其定义如下：\n$$ y = W_2(W_3x\\odot \\mathrm{SwiGLU}(W_1x)) $$这里 $W_3,W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$, $d_{ff}$ 是 MLP 的 hidden size, 代码中用 intermediate_size 来表示，因此\n$$ \\mathrm{parameter}(\\texttt{Qwen3MLP}) = dd_{ff} + dd_{ff} + d_{ff}d=3dd_{ff} $$这里三项分别代表 $W_1,W_3,W_2$ 的参数量。\n如果说，我们使用原始 transformer 的 MLP, 也就是\n$$ y = W_2\\max(0,W_1x+b_1)+b_2 $$其中 $W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$, $b_1\\in\\mathbb{R}^{d_{ff}}$, $b_2\\in\\mathbb{R}^d$, 则总参数为\n$$ \\mathrm{parameter}(\\texttt{TransformerMLP}) = d_{ff}d + dd_{ff} + d_{ff} +d = 2d_{ff}d + d_{ff} + d $$这里的四项分别代表了 $W_1,W_2,b_1,b_2$.\nQwen3Attention 接下来，就是 Attention 部分的参数，Qwen3Attention 的关键代码如下：\n1 2 3 4 5 6 7 8 9 10 class Qwen3Attention: def __init__(self, config: Qwen3Config, layer_idx: int): self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias) self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) 这里，我们先定义几个量：\n我们将 head 的个数记为 $h$, 即 num_attention_heads 我们将每个 head 的 hidden size 记为 $h_d$, 即 head_dim 我们将 key 和 value head 的个数记为 $h_{kv}$ , 即 num_key_value_heads Qwen3Attention 的参数由以下几个部分组成：\nQuery projection: $W_{Q}\\in\\mathbb{R}^{hh_d\\times d}$, $b_Q\\in\\mathbb{R}^{hh_d}$ Key projection: $W_K\\in\\mathbb{R}^{h_{kv}h_d\\times d}$, $b_K\\in\\mathbb{R}^{h_{kv}h_d}$ Value projection: $W_V\\in\\mathbb{R}^{h_{kv}h_d\\times d}$, $b_V\\in\\mathbb{R}^{h_{kv}h_d}$ output projection: $W_O\\in\\mathbb{R}^{d\\times hh_d}$ RMSNorm：前文已经提到过，两个 normalization (query norm 以及 key norm) 的总参数量为 $2h_d$. 因此， Qwen3Attention 部分的总参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3Attention}) = hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d = 2hh_dd + 2h_{kv}h_dd + 2h_d $$分别代表 $W_Q, W_O, W_K,W_V$ 和 两个 normalization layer 的参数量。\n注意，这里我们没有加入 bias, 这是因为 QKV bias 在 Qwen3 中被取消，取而代之的是两个 normalization.\n如果我们查看 Qwen2Attention 的代码，我们可以得到 Qwen2Attention 的总参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen2Attention}) = 2hh_{d}d + 2h_{kv}h_dd +h_{kv}h_d+hh_d+h_{kv}h_d $$分别代表 $W_Q, W_K, W_V,W_O$ 和 $b_Q,b_K, b_V$ 的参数量。\n我们将计算结果汇总在一起就得到：\n$$ \\begin{aligned} \\mathrm{parameter}(\\texttt{Qwen3ForCausalLM}) \u0026=d|V| + \\mathrm{parameter}(\\texttt{Qwen3Model})\\\\ \u0026=2d|V| + N*\\mathrm{parameter}(\\texttt{Qwen3DecoderLayer}) + d\\\\ \u0026= N*(2d + \\mathrm{parameter}(\\texttt{Qwen3MLP}) + \\mathrm{parameter}(\\texttt{Qwen3Attention})) + d(2|V|+1)\\\\ \u0026= N*(2d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd + 2h_d) + d(2|V|+1)\\\\ \\end{aligned} $$这是针对 Qwen3ForCausalLM 的参数量计算。这里的变量定义如下：\nMath Variable Code Variable Description $N$ num_hidden_layers Transformer block 个数 $\\vert V\\vert$ vocab_size 词表大小 $d$ hidden_size token embedding 的维度 $d_{ff}$ intermediate_size MLP 的中间层的维度 $h_d$ head_dim 每个 head 的维度 $h$ num_attention_heads query head 的个数 $h_{kv}$ num_key_value_heads key 和 value head 的个数 Verification 接下来，我们就可以基于 Qwen3 的模型来验证了，比如，Qwen3-32B 的配置如下:\nField Value $N$ 64 $\\vert V\\vert$ 151936 $d$ 5120 $d_{ff}$ 25600 $h_d$ 128 $h$ 64 $h_{kv}$ 8 根据上式，最终的参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3-32B}) = 32.762123264*10^9\\approx 32.8B $$我们使用 Qwen3-32B 的 index.json 可以得到其真实的参数量为\n$$ \\texttt{total\\_size}/\\texttt{precision} = 65524246528/2 = 32762123264 $$与我们计算的结果一致（这里除以 2 的原因是其表示模型权重文件的总大小，以 bytes 为单位，一般模型都是 bfloat16, 大小为 2 个 bytes, 因此总参数量为总大小除以权重的精度）\nMoE Model MoE model 与 Dense model 不同的地方在于每一层的 FFN, 因此，其总参数计算方式为：\n$$ \\mathrm{parameter}(\\texttt{Qwen3MoeForCausalLM})=N*(2d+\\mathrm{parameter}(\\texttt{Qwen3MoE})+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2d) + d(2|V|+1) $$对于 MoE layer, 其关键代码为：\n1 2 3 4 5 6 class Qwen3MoeSparseMoeBlock: def __init__(self, config): self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) self.experts = nn.ModuleList([Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]) 我们记 $n$ 为总专家个数，即 num_experts， 记 $k$ 为激活专家个数，即 num_experts_per_tok 或者 top_k,\n首先 gate 的参数量为 $dn$, 接下来每个 expert 都是一个 Qwen3MLP, 因此 experts 总参数量为 $n * 3dd_{ff}$. 这样 MoE layer 的总参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3MoE}) = nd + 3ndd_{ff} $$在推理时，只有一部分专家，也就是 $k$ 个专家会参与计算，此时激活参数量为\n$$ \\mathrm{parameter\\_activated}(\\texttt{Qwen3MoE}) = nd + 3kdd_{ff} $$我们带入到 Qwen3MoeForCausalLM 中就得到：\n模型总参数量为：\n$$ \\mathrm{parameter}(\\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3ndd_{ff}+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$模型激活参数量为：\n$$ \\mathrm{parameter\\_activated}(\\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3kdd_{ff} + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$这里的变量定义如下：\nMath Variable Code Variable Description $N$ num_hidden_layers Transformer block 个数 $\\vert V\\vert$ vocab_size 词表大小 $d$ hidden_size token embedding 的维度 $d_{ff}$ intermediate_size MLP 的中间层的维度 $h_d$ head_dim 每个 head 的维度 $h$ num_attention_heads query head 的个数 $h_{kv}$ num_key_value_heads key 和 value head 的个数 $n$ num_experts 总专家个数 $k$ top_k (num_experts_per_tok) 激活专家个数 MoE Verification 我们用 Qwen3-235B-A22B 来验证，其配置如下：\nField Value $N$ 94 $\\vert V\\vert$ 151936 $d$ 4096 $d_{ff}$ 1536 $h_d$ 128 $h$ 64 $h_{kv}$ 4 $n$ 128 $k$ 8 计算得到总参数量为\n$$ \\mathrm{parameter}(\\texttt{Qwen3-235B-A22B}) = 235.09363456*10^9\\approx 235B $$计算得到激活参数量为\n$$ \\mathrm{parameter\\_activated}(\\texttt{Qwen3-235B-A22B}) = 22.14456064 * 10^9 $$实际总参数量为\n$$ 470187269120/2 = 235093634560.0\\approx 235B $$可以看到，计算结果与实际相符。\nExtension 通过以上计算过程，我们可以很轻松将上述公式扩展到混合架构或者是 MLA 上，对于混合架构，我们分别计算不同 attention 的 layer 个数，然后分别计算。对于 MLA, 我们可以替换 Attention 的计算逻辑。\n对于小语言模型系列，比如 0.6B 和 1.7B, 模型的大部分参数集中在 embedding 上，因此 Qwen3 采取了 tie embedding 的方式来减少参数量，具体做法就是 nn.Embedding 和 lm_head 共享参数。\nVisualization 接下来，我们来可视化一下不同大小模型不同模块的参数量占比。计算的代码如下：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 def compute_param_distribution( N: int, V: int, d: int, d_ff: int, h_d: int, h: int, h_kv: int, tie_word_embeddings: bool) -\u0026gt; dict: \u0026#34;\u0026#34;\u0026#34; Compute parameter distribution across different components of the model. Args: N: Number of layers V: Vocabulary size d: Hidden dimension d_ff: Feed-forward dimension h_d: Head dimension h: Number of attention heads h_kv: Number of key/value heads tie_word_embeddings: Whether input and output embeddings are tied Returns: dict: Parameter counts for each component \u0026#34;\u0026#34;\u0026#34; # Embedding parameters embedding = V * d if not tie_word_embeddings: embedding *= 2 # Attention parameters per layer attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d attention_total = attention_per_layer * N # FFN parameters per layer ffn_per_layer = 3 * d * d_ff ffn_total = ffn_per_layer * N # RMSNorm (2d), QK-Norm (2d), output normalization (d) others = N * (4 * d) + d return { \u0026#34;Embedding\u0026#34;: embedding, \u0026#34;Attention\u0026#34;: attention_total, \u0026#34;FFN\u0026#34;: ffn_total, \u0026#34;Others\u0026#34;: others, } 对于 dense 模型，每部分的参数量可视化如下：\n可以看到，随着模型 size 增加，模型大部分参数量都集中在 FFN 上。\n接下来，我们可视化一下 MoE 模型的参数分布，核心计算代码如下：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 def compute_moe_param_distribution(N, V, d, d_ff, h_d, h, h_kv, n, k, activate=False) -\u0026gt; dict: \u0026#34;\u0026#34;\u0026#34; Compute parameter distribution across different components of the MoE model. Args: N: Number of layers V: Vocabulary size d: Hidden dimension d_ff: Feed-forward dimension h_d: Head dimension h: Number of attention heads h_kv: Number of key/value heads n: Number of experts k: Number of experts per token activate: Whether to compute activated parameters Returns: dict: Parameter counts for each component \u0026#34;\u0026#34;\u0026#34; # Embedding parameters embedding = 2 * V * d # Attention parameters per layer attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d attention_total = attention_per_layer * N # FFN parameters per layer if activate: moe_per_layer = n * d + 3 * k * d * d_ff else: moe_per_layer = n * d + 3 * n * d * d_ff ffn_total = moe_per_layer * N # RMSNorm (2d), QK-Norm (2d), output normalization (d) others = N * (2 * d + 2 * h_d) + d return { \u0026#34;Embedding\u0026#34;: embedding, \u0026#34;Attention\u0026#34;: attention_total, \u0026#34;MoE\u0026#34;: ffn_total, \u0026#34;Others\u0026#34;: others, } 结果如下\n可以看到，MoE 模型的大部分参数还是集中在 MoE 模块上，但是由于其稀疏机制，在激活的参数里，MoE 占比从 95% 以上降低到了 60% 左右。\nConclusion 在本文中，我们基于 Qwen3 大语言模型系列，介绍了如何计算 dense 模型和 MoE 模型的参数量。模型的参数量计算为后面的显存占用以及优化提供了基础。\nReferences Qwen3 Collection ","date":"2025-07-22T10:50:47+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/llm-parameter-computation/","title":"LLM Parameter Computation"},{"content":"Seed 在 2025 年 6 月 25 号发布了 Seed 1.6，支持 adaptive deep thinking, multimodal understanding, GUI-based interaction, and deep reasoning with a 256K context window\nSeed 1.6 基于 Seed 1.5 开发，是一个 MoE 模型, 总参数为 230B, 激活参数为 23B.\nPre-training pre-training 分为 3 个 stage:\nText-only pre-training: 使用 web pages, books, papers, code 以及其他数据来进行训练，作者使用了 rule-based 以及 model-based 策略来提高数据的质量。 Multimodal mixed continue training: 提高学科，代码以及 reasoning 相关的数据，来进一步提高知识密度以及 reasoning 的复杂度。同时，加入高质量的视觉数据。 Long-context continual training: 将模型的上下文长度从 32K 扩展到 256K Seed 1.6 Base 的表现\nModel Seed1.6 Base LLaMA-4 -Maverick Base DeepSeek-V3Base Qwen3-235B-A22B Base Seed1.5 Base MMLU 88.83 85.16 87.19 87.81 88.35 MMLU-Pro 69.98 63.91 59.84 68.18 66.47 SuperGPQA 45.08 40.85 41.53 44.06 36.81 BBH 92.08 83.62 86.22 88.87 88.36 GPQA-Diamond 43.43 43.94 41.92 47.47 45.25 GSM8k 93.10 87.72 87.57 94.39 89.99 MATH 72.86 63.32 62.62 71.84 66.18 MBPP 83.60 75.40 74.20 81.40 81.40 Post-training 作者在 post-training 构建了两个变体：\nSeed1.6-Thinking: 先进的多模态 reasoning 能力 Seed1.5(Adaptive CoT): 使用 AdaCoT 来压缩 CoT 的长度，提高效率 Seed1.6-Thinking Seed1.6-Thinking 与 Seed-Thinking-v1.5 的训练方式差不多，作者使用了 multi-stage [[rejection fine-tuning]] (RFT) 和 RL 来提高模型的表现。在每一轮中，先进行 RL 的训练，然后进行 RFT 的训练。相比于 Seed-Thinking-v1.5, Seed1.6-Thinking 使用了更多的算力和更高质量的数据，包括数学问题，代码，谜题以及 non-reasoning 数据。\n表现如下：\n为了进一步提高模型的表现，作者在推理时使用了 parallel decoding, 来让模型在回答之前使用更多的 token, 作者发现 parallel decoding 可以显著提高模型在困难任务上的表现。\nSeed1.5(Adaptive CoT) reasoning model 的问题是会出现 overthinking,导致输出的过程中花费很多不必要的 token.\n为了解决这个问题，Seed1.6 提出了 AdaCoT, 来训练得到三个模型：FullCoT, NoCoT 和 AdaCoT\n通过 AdaCoT, 我们可以显著降低 CoT 的长度。为了完成这个目标，作者在 RL 训练阶段构建了一个 reward, 来惩罚模型的 overthinking. 用户可以基于 prompt 来选择不同的模型：\nFullCoT: 对每个 prompt，都进行 reasoning, 模型表现与 Seed1.6-Thinking 差不多，但是 CoT length 更小 NoCoT: 直接回答，不进行 reasoning AdaCoT: 由模型自主决定是否需要进行 reasoning, 模型是 FullCoT 和 NoCoT 的一个 fusion 版本 结果发现，对于不同难度的任务，AdaCoT 的出发概率是不同的，说明模型会基于任务难度调整思考方式。\n作者还评测了以下模型在高考中的表现，结果如下图：\nConclusion 作者提出了 Seed1.6 系列多模态大模型，模型在多个 benchmark 上表现优异。作者还提出了 AdaCoT, 让模型基于问题难度动态选择 CoT length\n作者认为下一步工作是：\n构建更加高效的架构，来进一步提高 reasoning 的有效性 扩展模型的多模态能力 提升模型的 agent 能力，特鄙视 end-to-end task execution References blog ","date":"2025-07-18T14:59:35+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-seed1.6/","title":"Notes on Seed1.6"},{"content":"V-Triune 探究了如何使用一套统一的 RL 训练框架，来同时提高模型的 perception 和 reasoning 能力。\nIntroduction 已有的 RL 训练方法主要提升模型在特定 domain 上的能力，比如 MM-Eureka, Vision-R1 等专注于提升 MLLM 的 reasoning 能力，而 R1-V, DeepPerception 等专注于提升模型的 perception 能力。\n因此，本文的 motivation 就是，如何用一套统一的 RL 训练框架，来同时提高模型的 perception 和 reasoning 能力。\nV-Triune 包含 3 个关键模块：\nSample-Level data formatting: 即在 sample 的层面定义 reward 等信息 Verifier-Level Computation: 即分离 verifier 和 generator, 提高整体效率 Source-Level Metric Monitoring: 基于 data source 来监控模型的表现 作者还提出了 Dynamic IoU reward, 来提高训练的稳定性。\n基于 V-Triune 和 Qwen2.5-VL, 作者提出了 Orsta,一个在 reasoning 和 perception 任务上均有提升的多模态大模型。\n论文的主要贡献为：\n提出了 V-Triune, 一个统一的 RL 训练框架，来同时提高模型的 reasoning 和 perception 能力 在 infra 上进行了改进，提高整体的训练效率和稳定性 基于 V-Triune, 构建了 Orsta, 相比于 Baseline, 模型的 perception 和 reasoning 能力均有了提升。 Preliminary 作者首先回顾了一下 GRPO 和 DAPO , 我们这里就不再赘述。\n然后，作者介绍了一下 reward function, reward function 由两部分组成， 第一部分是 format reward, 第二部分是 accuracy reward,.\n对于 format reward, 作者采取了和 OpenR1 相同的做法，也就是要求输出仅包含 1 个给定的 token, 这里特定的 token 为 $s_1=\\texttt{}$, $s_2=\\texttt{}$, $s_3=\\texttt{}$, $s_4=\\texttt{}$, reward 的定义如下：\n$$ R_{\\mathrm{format}}(o_q) = 0.25\\sum_{i=1}^4 \\mathbb{1}(\\mathrm{count}(o_q,s_i)=1) $$这里 $\\mathbb{1}(\\cdot)$ 是示性函数。\n对于 accuracy reward, 作者根据不同的任务分别进行处理。\n对于 reasoning 任务，reward function 的定义如下\n$$ R_{\\mathrm{acc}}(\\hat{a},a) = \\begin{cases} 1, \u0026\\text{ if }\\texttt{verify(parse($\\hat{a}$), parse($a$))}\\text{ is True}\\\\ 0, \u0026\\text{ else} \\end{cases} $$这里 $\\hat{a}$, $a$ 分别是 prediction 和 ground truth.\n对于 perception 任务，reward function 与我们要处理的子任务相关。如果是 OCR 和 counting 任务，则我们采取和 reasoning 一样的 reward function. 如果我们要处理的是 grounding 和 detection 任务，则我们可以使用 IoU 以及 mAP 两种 reward 形式。\nIoU reward 的定义如下：\n$$ R_{\\mathrm{acc}}(\\hat{a},a) = \\begin{cases} \\mathrm{IoU}(\\hat{a}, a), \u0026\\text{ if }\\mathrm{IoU}(\\hat{a},a)\\geq \\epsilon\\\\ 0, \u0026\\text{ else} \\end{cases} $$这里 $\\epsilon\u003e0$ 为超参数， IoU 定义如下\n$$ \\mathrm{IoU}(\\hat{a},a) = \\frac{\\mathrm{Area}(\\hat{a}\\cap a)}{\\mathrm{Area}(\\hat{a}\\cup a)} $$mAP 的定义如下\n$$ \\mathrm{AP}_c = \\int_0^1 \\max_{\\tilde{r}\\geq r}P_c(\\tilde{r}) dr,\\quad mAP=\\frac1C\\sum_{c=1}^C \\mathrm{AP}_c $$最终的 reward 是 format reward 和 accuracy reward 的加权求和：\n$$ R = \\alpha_{\\mathrm{acc}}R_{\\mathrm{acc}} + \\alpha_{\\mathrm{format}}R_{\\mathrm{format}} $$V-Triune V-Triune 框架如下图所示，其包含三个模块\nSample-Level Data Formatting 核心思想就是不同的任务的 reward function 可能是不一样的，如果在 task 层面设计 reward function 的话，RL 训练就没有了 flexibility, 因此，在本文中，作者在 sample 层面设计 reward function. 具体做法就是，每个样本除了 QA 相关信息，还有 reward model 的相关参数，这样就可以更加灵活地调整不同 sample 的损失了。\nVerifier-Level Reward Computation 核心思想就是构建多个不同的 Verifier, 来分别处理不同的 sample, 因为我们 reward model 也是在 sample 层面构建的，因此我们可以基于 sample 的 metadata 来动态分配对应的 verifier, 这样就提升了 veriier 的可扩展性。本文作者使用了 MathVerifier 和 DetectionVerifier 两种。\nSource-Level Metric Monitoring 核心思想就是根据 sample 的 metadata 信息来监控不同的 metric, 这样可以防止训练不稳定。\nDynamic IoU Reward 核心思想就是根据训练进程，动态调整 IoU 的阈值。\n作者发现，如果将 IoU 阈值设置的比较低的话，模型很容易拿到比较高的 reward; 如果将 IoU 阈值设置的比较高的话，模型又很难拿到奖励。\n因此，作者的解决方案就是，使用课程学习，随着训练的进行，逐步提高 IoU 的阈值。\nData 数据集列举如下：\nMath: mm_math, geometry3K, mmk12 Puzzle: PuzzleVQA, AlgoPuzzleQA, VisualPuzzles Science: ScienceQA, SciVQA, ViRL39K (Broader STEM topics, (GradeSchool) Science) Chart: ChartQAPro, ChartX, Table-VQA, ViRL39K (Tables/Diagrams/Charts) Detection: V3Det, Object365 Grounding: D3 Counting: CLEVR OCR: LLaVA-OV (OCR questions), EST-VQA Rule-based filtering 对于 visual reasoning 数据，作者的过滤如下：\n多项选择题以及判断题，防止 reward hacking 答案包含特殊符号如 \u0026ldquo;=\u0026rdquo;, \u0026ldquo;[\u0026rdquo; 等的数据 答案字符数超过 20 个，防止答案太复杂 对于 visual perception 数据，作者过滤流程如下：\nDetection: 超过 10 个 bounding box 的，或者 box 占 50% 以上面积的数据，保持 single BB 和 multi BB 的比例为 1:2 Grounding: box 占 50% 以上面积的数据, label 过于复杂的数据 Counting: 保持类别平衡，仅使用英文数据 OCR: 保留英文数据，对 labels 进行 verify Difficulty-based filtering 主要是移除比较简单的问题。\n经过这两个阶段的过滤之后，得到了 20.6K perception samples 以及 27.1K reasoning samples. 数据保存在 Parquet 格式\nTraining Recipe Disable ViT Training 作者首先发现，全量微调会导致训练崩溃。作者分析原因发现，主要是 ViT 在反向传播过程中，存在 gradient amplify 的问题，第一层与最后一层的梯度的 norm 相差 5-10 倍。作者认为有两点原因：\nRL 会强制要求模型的不同模态之间进行对齐 对比学习训练得到的 VIT 可能使得其不适合 RL 的训练。 基于这两点考虑，作者再进行训练的时候，冻结了 ViT 模块，结果如下图\nSpurious Image Special Tokens 作者发现，模型的输出可能会包含一些 \u0026lt;image_pad\u0026gt; 等 token, 因此，作者对输出进行了过滤，然后再进行计算。\nCoT Prompt Pool 作者还发现，多样化的 CoT prompt 会损害模型的表现。因此，作者构建了 20 个 prompt 模版，每次随机挑选其中一个加入到 instruction 中。\nTraining Configuration 模型基于 Qwen2.5-7B-instruct 和 Qwen2.5VL-32B 来进行开发\n作者探究了 on-policy 和 off-policy 两种配置。roll-out batch size 设置为 1024, 使用 GRPO 进行训练，GRPO 的 group size 设置为 8\n作者还设置 $\\epsilon_{high}=0.28$ 以及 $\\epsilon_{low}=0.2$ 来提高 token 的采样率。作者去除了 KL divergence loss, 并且在 token 层面上进行平均\nEvaluation Performance Analysis on-policy v.s. off-policy 作者首先对比了以下 on-policy RL 和 off-policy RL 两种训练方式的表现， 结果如下图所示\n结果有两点发现：\non-policy 的表现基本上都是比 off-policy 要好的。 小模型通过 RL 带来的表现提升更明显，大模型的表现则更慢 generalization 作者还评估了以下模型的泛化性，结果如下图所示\n结果发现，RL 的训练更像是对齐，而不是泛化。也就是说，模型在 in-domain 的任务上表现较好，在 OOD 的任务上表现就一般。这与已有的结论是不一样的\n[!tip] Recall Magistral 发现模型在 math domain 上进行训练，在 code domain 上的表现也会提升\nTraining Dynamics 作者还分析了不同任务不同指标随训练进行的变化情况，如下图所示\n结果显示，reasoning 和 OCR 任务岁训练进行其输出长度和正确率都有所提升。但是 detection 相关任务则变化不是很明显。\nAblation Study 作者进行了三个消融实验：\n训练策略：冻结 ViT, 冻结 LLM, 两者都不冻结 任务分解策略：仅使用 reasoning 数据进行训练，仅使用 perception 数据进行训练，同时使用两种数据进行训练。 learning rate: 不同的学习率 实验结果如下图所示\n结果发现，\n仅训练 ViT 对表现没有任何帮助，只有训练 LLM 才会带来性能的提升。 同时使用两种数据进行训练的效果是最好的，其次是仅使用 reasoning 数据进行训练，最后是仅使用 perception 数据进行训练 较大的学习率会损害模型的表现。 Conclusion 本文中，作者提出了 V-Triune, 一个使用 RL 来同时提高 MLLM 的 reasoning 和 perception 能力的 RL 训练框架。\n作者发现 RL 对于 VLM 来说更像是一种 Alignment 策略，而不是一种 Generalization 策略，也就是模型并没有引入新的能力，而是提升模型本身的 utility 以及 robustness.\n作者认为本文有以下 Limitation:\n在 perception 任务上没有看到明显的 test-time scaling 现象。作者认为，探究使用 o3 类似的方法或许能提高模型的表现 探究 RL-zero 在 VLM 中的应用，也就是直接在 base model 上进行 RL 的训练。 References arxiv ","date":"2025-07-17T09:37:36+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-v-triune/","title":"Notes on V-Triune"},{"content":"Magistral 是 Mistral 提出的一个 reasoning model 系列，主要针对 math 和 code 两个 domain\nIntroduction Mistral 在 2025 年 6 月 12 日发布了 Magistral ，一个 reasoning model, 包含两个模型，一个是由纯 RL 训练得到的 Magistral Medium, 另一个是由 SFT 和蒸馏 Magistral Medium 得到的 Magistral Small\n作者首先介绍了一下本文的贡献：\n介绍了如何仅使用 RL (而不是用蒸馏) 来训练 Magistral Medium infra 上的改进，主要使用最新的权重来更新 generator 多语种能力，支持多种语言 系统性探究了 RLVR 的能力边界 开源了 Magistral small (24B) Method RL RL 算法基于 GRPO 改进，主要有以下几点：\n去掉了 KL Divergence loss, 这一点跟 DAPO 是一致的，提升模型的探索能力 Loss Normalization，在 sample 层面做平均，还是跟 DAPO 一致，减少不同长度输出对训练造成的影响 Advantage Normalization, 作者首先将每个 token 的 Advantage 定义为 $\\hat{A}_i=R_i-\\mu$, 其中 $\\mu$ 是每个 group 的 reward 均值， 然后在每个 mini-batch 里对 advantage 进行 Normalization (这里 $\\hat{A}_{mean}$ 和 $\\hat{A}_{std}$ 分别为 mini-batch advantage 的均值和方差)： $$ \\hat{A}_{i,t}^{norm}=\\frac{\\hat{A}_i-\\hat{A}_{mean}}{\\hat{A}_{std}} $$ CLIP-Higher, 跟 DAPO 一致，提高稀有 token 的被采样概率 Eliminating non-divsere groups, 跟 DAPO 一致，去掉过于简单和过于难的题目 最终 RL 阶段训练的损失函数为：\n$$ \\mathcal{L}(\\theta) = \\mathbb{E}_{q\\sim P(Q),\\{o_i\\}_{i=1}^G\\sim \\pi_{old}(\\cdot\\mid q)}\\frac{1}{\\sum_{i=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\min\\left[r_{i,t}(\\theta)\\hat{A}_{i,t}^{norm}, \\mathrm{clip}(r_{i,t}(\\theta), 1-\\epsilon_{low}, 1+\\epsilon_{high})\\hat{A}_{i,t}^{norm}\\right], \\mathrm{s.t.}, \\exists 1\\leq m \u003c n \\leq G, r_m\\neq r_n. $$其中\n$$ r_{i,t}(\\theta) = \\frac{\\pi_{\\theta}(o_{i,t}\\mid q, o_{i,","date":"2025-07-16T11:04:04+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-magistral/","title":"Notes on Magistral"},{"content":"Hugging Face 在 2025 年 7 月 8 号发布了 SmolLM3, 一个 3B 的，128K 上下文，支持 6 种语言，支持 dual mode reasoning 的小语言模型。\nPre-training Architecture SmolLM3 是一个基于 transformer 的小语言模型，其模型架构与 SmolLM2 相似，SmolLM3 做了以下改进：\n使用 GQA 代替 multi-head attention. 作者通过消融实验发现 GQA 和 MHA 的效果差不多，并且还可以减少 KV cache 的 size NoPE: 作者使用了NoPE, 来选择性（每 4 个 layer）移除 position embedding. 这个方法可以在不损害模型短文本能力的同时，提高模型长上下文的表现 Intra-Document Masking: 不同的文档之间使用 attention masking 隔开 Training stability: 与 olmo 2 一样，作者移除了 embedding layer 的 weight decay, 来提升训练的稳定性。 Data 与 SmolLM2 一样，作者使用了11.2T token 进行训练，训练包括 3 个 stage。作者针对数据混合策略进行了消融实验，实验配置如下：\nStage Stable phase Stable phase Decay phase Description Base training High quality injection LR Decay Tokens 8T 2T 1.1T Web 85% 75% 63% Code 12% 15% 24% Math 3% 10% 13% datasets web: FineWeb-Edu, DCLM, FineWeb2, Fineweb2-HQ\ncode: The Stack v2, StarCoder2 PRS, Jupyter and Kaggle notebooks, Github issues, StackExchange\nmath: FineMath3, InfiWebMath3+ Adding Stack-Edu, FineMath4+, InfiWebMath4+, MegaMath upsampling of high-quality code data\nupsampling math data and introducing instruction and reasoning datasets such as OpenMathReasoning 其中，Web data 包含 12% 的多语种数据。\nTraining Recipe 训练过程中，作者使用了 2.36M tokens 的 batch size, 上下文长度为 4096. 优化器为 AdamW\n作者使用了 nanotron 进行训练， datatrove 来处理数据， lighteval 来评估模型的表现。\n模型在 384 张 H100 上训练了 24 天。分布式训练的配置如下\nMid-training Mid-training 的主要目标为扩展模型的上下文以及提升模型的 reasoning 能力。\nLong Context Extension 在预训练阶段结束之后，作者使用了额外的 100B tokens 来扩展模型的上下文。作者将扩展过程分为两个 stage:\nStage 1: 将模型的上下文从 4K 提升到 32K. 具体做法是将 RoPE 的 base frequency 提升到 1.5M Stage 2: 将，模型的上下文从 32K 提升到 64K. 具体做法是将 RoPE 的 base frequency 提升到 5M 训练过程中，作者对 math, code 和 reasoning data 进行了上采样。作者发现，对长文本数据进行上采样并不会提高模型在 RULER 和 HELMET 上的表现。\n[!tip] Recall Qwen2.5-1M 里也分析了长文本数据的问题，也就是大部分长文本数据依然是局部相关性强，而全局相关性弱\n与 Qwen2.5 一样，作者还是用了 YARN 来在推理阶段进一步提高模型的上下文长度，作者发现模型可以处理 128K 上下文长度的文本。\nReasoning Mid-training 扩展模型上下文长度之后，作者还额外增加了一个 mid-training stage 来提高模型的 reasoning 能力。这个阶段的目标在于提升模型的通用能力。作者希望模型不针对特定的 domain, 如 math 或者 code 等。\n训练过程中，作者使用了 35B 的 token. 数据来源包括 Open-Thoughts3-1.2M 以及 NVIDIA 的 Llama-Nemetron-Post-Training-Dataset-v1.1. Reasoning trace 由 DeepSeek-R1 生成。作者使用了 ChatML 的格式，还使用了 Packing 来提升训练效率。训练持续了 4 个 epoch.\nPost-training 如何构建 dual instruction model 来同时支持 reasoning 和 non-reasoning 任务并没有一个共识，大部分模型的数据都是非公开的。因此，作者就构建了一个 training pipeline, 用于提升模型在两种模式下的能力。\nChat Template 作者首先构建了一个 chat template, 用于支持 reasoning 和 non-reasoning 两种模式。该 chat template 支持用户使用 /think 和 /no_think flag 来控制模型的思考模式。在 non-reasoning 模式下，作者还在模型输出中 prefill 了 \u0026lt;think\u0026gt;\\n\\n\u0026lt;/think\u0026gt;, 这一点与 Qwen3 一致。\nSmolLM3 还支持工具调用，其在 chat template 中加入了两种描述方式：XLM tools 和 Python tools.\nSmolLM3 还在 system prompt 中加入了 metadata, 如知识的截止时间，当前的 reasoning mode 等。\nSFT 作者针对 math, code, general reasoning, instruction following, 以及 multilinguality 这几个领域来提升模型的表现。\n作者首先使用 reasoning mode 的 Qwen3-32B 来合成数据，合成数据的过程如上图所示。\n最终，SFT 阶段的数据包括 1.8B token, 其中 1B 为 non-reasoning mode 的 token, 覆盖 12 个数据集， 0.8B 为 reasoning token, 覆盖 10 个数据集。作者训练了 4 个 epoch, 使用了 Packing 的技巧。\nAPO SFT 阶段之后，作者使用了 Tulu3 的 preference dataset 用于 non-reasoning mode 的训练，然后合成了一批数据用于 reasoning mode 的训练，这批合成数据使用 Qwen3 32B 和 Qwen3 0.6B 生成得到，具体做法就是 Qwen3 32B 输出的结果定义为正样本，Qwen3 0.6B 输出的结果定义为负样本。\n作者使用了 APO 算法来进行训练，APO 是 DPO 的一个变体， DPO 的目标函数为\n$$ \\mathcal{L}_{DPO}(x,y_w,y_{l}; \\theta) = -\\log \\sigma(r_\\theta(x,y_w) - r_\\theta(x, y_l)) $$其中\n$$ r_\\theta(x,y) = \\beta\\log \\frac{\\pi_{\\theta}(y\\mid x)}{\\pi_{ref}(y\\mid x)} $$$\\beta\u003e0$ 是一个超参数。APO 的目标函数如下\n$$ \\mathcal{L}_{APO}(x,y_w,y_l;\\theta) = -\\sigma(r_\\theta(x,y_w)) + \\sigma(r_\\theta(x,y_l)) $$作者发现，模型的 reasoning 能力提升之后，其长上下文能力反而下降了。作者认为这是因为在 reasoning mid-training stage 的训练中，提升 reasoning 能力损害了模型的 long context 能力。并且，APO 的训练数据上下文长度大多都在 24K 左右。为了解决这个问题，作者提出了 Model merging 的方法\nModel Merging 作者使用 MergeKit 来完成 model merging 的任务。merge 的过程包括两步：\n构造一个 model soup, 包括 APO 的每个 checkpoint 将 model soup 与 mid-training 的一个拥有强上下文能力的 checkpoint 结合起来，作者使用了 linear merge, APO model soup 和 mid-training checkpoint 的权重分别为 0.9 和 0.1. [!tip] Recall Kimi-k1.5 也使用了 model merging 的方法，来将 long CoT 模型的能力迁移到 short-CoT 模型上去\nEvaluation 作者首先评估了一下 base model 的表现，评估使用了 12 个 benchmark, 对比的模型包括 Qwen3, Gemma3, LLaMA 3.2. 结果如下图所示\n模型的多语种表现如下\nInstruction (non-reasoning) 版本的评估结果如下\nReasoning 版本的评估结果如下\n可以看到，Qwen3-4B 的表现是最好的。而 SmolLM3 的表现在 3B 左右也是非常强劲的\nConclusion 作者提出了 SmolLM3, 一个拥有长上下，文多语种以及 dual mode reasoning 能力的大语言模型，作者详细介绍了数据，训练以及 model merging 的技巧，来提高模型的表现。\nReferences blog ","date":"2025-07-15T11:01:13+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-smollm3/","title":"Notes on SmolLM3"},{"content":"智谱 AI 在 25 年 7 月份发布了 GLM-4.1V-Thinking, 一个 9B 的多模态大语言模型，其在多个 benchmark 上达到了相同大小 MLLM 的 SOTA\nIntroduction 已有工作如 MiMo-VL 和 V-Triune 主要集中在特定的领域，目前还没有一个在所有领域都能超过 non-reasoning model 表现的多模态 reasoning model.\n基于这个目标，作者提出了 GLM-4.1V-Thinking, 一个 9B 的，用于通用领域多模态 reasoning 的 MLLM. 在预训练阶段，作者通过构建多样化的数据来提升模型的基础能力。在 post-training 阶段，作者构建了 domain-specific 的数据来让模型学会 reasoning. 最后，作者提出了 Reinforcement Learning with Curriculum Sampling (RLCS) 来扩展模型的 reasoning 能力。\n作者主要发现如下：\nmulti-domain RL 的泛化性非常强，在某一个 domain 上进行 RL 训练也可以提高模型在其他 domain 上的表现 动态选择 RL 训练的数据可以有效提高模型的表现和训练效率 一个 robust 以及 precise 的 reward model 对于 multi-domain RL 是至关重要的 Method Architecture GLM-4.1V-Thinking 是一个标准的 ViT-MLP-LLM 架构，其中：\nViT: ViT 使用的是 AIMv2-Huge MLP: 是一个 3 层的 MLP，架构为 linear-LayerNorm-GELU-SwiGLU, 并且在进入 MLP 之前，GLM-4.1V-Thinking 还会在 spatial 上进行降采样，降采样率为 2 LLM: LLM 使用的是 GLM 与 Qwen2.5-VL 一样，作者将 encoder 的 2D convolution 替换为了 3D convolution，用于处理视频输入，然后对于 single-image inputs，GLM-4.1V-Thinking 会将输入复制，变成一个含有两帧相同图片的 video 输入\n为了支持不同分辨率图片输入，作者进行了两点改进：\n第一是使用了 2D RoPE 来处理不同分辨率的图片输入\n第二是使用 bicubic interpolation 来将不同分辨率的图片适应到 encoder 的 position embedding 上。具体来讲，对一个拥有 $H_p\\times W_p$ patches 的图片，每个 patch 的坐标 $(w,h)$ 首先会被 normalized 到 $[-1,1]$ 中：\n$$ g_{norm} = (w_{norm}, h_{norm}) = 2* \\left(\\frac{w+0.5}{W_p}, \\frac{h+0.5}{H_p}\\right) - 1 $$然后，我们再使用 bicubic interpolation $\\mathcal{I}_{bicubic}$ 将坐标映射到原始的 position embedding 上：\n$$ P_{adapted} = \\mathcal{I}_{bicubic}(P_{orig}, g_{norm}) $$对于视频输入，作者在每一帧之后插入了一个 time index token, 用每一帧的 timestamp string 来表示。这个做法与 Qwen2.5-VL 一样。\nPre-training Pre-training Data Pre-training 数据包括 image-caption data, interleaved image-text data, OCR data, grounding data, video data 和 instruction tuning data\nImage caption data Image caption data 用于提升模型的世界知识。作者从 LAION, DataComp, DNF 以及 Wukong 等数据集收集了 10B 的 image-text pairs, 然后进行数据清洗：\nHeuristic-based filtering: 基于规则，如 resolution, caption length 等过滤掉低质量的图片 Relevance filtering: 计算 CLIP score, 过滤掉低质量的数据 Concept-balanced resampling: 基于 MetaCLIP , 来提升数据的覆盖程度，解决长尾分布的问题。这里 Seed1.5-VL 有一个类似的消融实验。 Factual-centered re-captioning: 通过 rewrite caption 来提升 caption 的质量和信息密度。Idefics2 通过实验发现，rewriting caption 可以提高模型的表现。作者在最终的数据集里，混合了一部分 re-caption 的数据。 Interleaved image-text data 相比于 image-caption data, interleaved image-text data 体量更大，且逻辑关系更强。但是噪声也更多，因此作者针对不同来源的数据构建了不同的数据清洗策略：\nWeb data processing pipeline: 数据来源为 MINT, MMC4, OmniCorpus 等。作者首先基于 CLIP-score，去掉与上下文不相关的图片；然后，作者去除掉广告，二维码等低质量的图片；最后，作者将图多文少的样本也给去除掉，比如相册图片。作者还训练了一个 high-knowledge-density image classifier, 来识别信息密度高的样本。 Academic book processing pipeline: 作者从电子书中收集了 100M 的样本，这些电子书主要是 STEM domain，然后作者构建了一个 PDF parsing tool 来提取文档内容。 OCR data OCR 数据包含220M images, 数据集由三部分组成：\nSynthetic document images: 基于预训练语料，使用不同的格式来渲染文字生成不同的图片，然后基于 LAION 的图片作为背景 Natural scene text images: 使用 Paddle-OCR 来从图片中提取文字以及对应的 bounding box Academic documents: 使用类似 Nougat 的方法来构建 PDF page-Markdown 的数据。 Grounding data 主要包含自然场景和 GUI 两个 domain 的 grounding 数据：\nNatural image grounding: 基于 LAION-115M 得到。作者使用 GLIPv2, 来自动化解析 caption 中的实体，然后预测对应的 bounding box, 作者将 bounding box 较少的样本去除掉。最终得到40M高质量的自然场景数据 GUI grounding: 基于 CC 提取对应的 url, 然后使用 Playwright 来与网页进行交互，进而获取到所有的 DOM elements 与其对应的 bounding box. 最终收集到 140M高质量的 QA 样本 Video data Video data 从多个来源得到，作者使用了 fine-grained human annotation 来保证数据质量。作者还标注了 camera motion 等信息\n为了保证数据的质量，作者进行了去重，来减少数据语义的相似性。\nInstruction tuning data instruction tuning data 的策略如下：\nTask coverage and taxonomy: 优化数据分布 Complex scenario augmentation: 构建复杂的指令跟随数据 Data contamination check: 数据污染检查 最终一共收集到 50M样本，覆盖 visual perception, multimodal reasoning, GUI agent 等任务。\nPre-training Training Pre-training 由两个 stage 组成：\nMultimodal pre-training: 提高模型的通用多模态能力，上下文长度为 8192, global batch size 为 1536. 使用了 data packing 策略，使用了 2-way tensor parallelism Long-context continue training: 提升模型在高精度图片，长视频，长上下文场景下的表现，上下文长度为 32768, 使用了 2-way tensor parallelism 和 4-way context parallelism. SFT SFT Data 数据包括 long CoT reasoning 数据，来让模型以标准格式输出 multi-step solution.\n数据包括 verifiable tasks 以及 non-verifiable tasks, 覆盖中英两种语言。作者过滤了非常简单和非常困难的样本。\n数据格式如下\n1 \u0026lt;think\u0026gt; {think_content} \u0026lt;/think\u0026gt; \u0026lt;answer\u0026gt; {answer_content} \u0026lt;/answer\u0026gt; 对于 verifiable tasks, 输出格式为\n1 \u0026lt;answer\u0026gt; \u0026lt;|begin_of_box|\u0026gt; {answer_content} \u0026lt;|end_of_box|\u0026gt; \u0026lt;/answer\u0026gt; 数据清洗策略包括：\n严格遵守格式要求 reasoning style 不一致或者有噪声 混合语言输出。 作者还使用 RL checkpoints 来生成一部分高质量数据，加入到 SFT 阶段的数据集里面。\nSFT Training 全参数微调。上下文长度为 32768， batch size 为 32.数据包括 long-form reasoning data 以及纯文本数据，包括 math, multi-turn conversation, agent planning 以及 instruction following 等.\n[!tip] Observation 作者发现，在 SFT 阶段即使是使用带噪声的 reasoning data，模型在 RL 阶段也能继续训练。也就是说，不完美的 reasoning trace 也能提供 guidance. 但是，使用更高质量的数据可以让模型 RL 阶段的训练更稳定且表现更好。\nRL 与 Seed1.5-VL 一样，作者在 RL 阶段结合了 RLVR 和 RLHF 两个训练目标。任务包括 STEM problem solving (such as mathematics, physics, chemistry), grounding, optical character recognition (OCR), video understanding, GUI agents, chart and document understanding, logical reasoning, and instruction following.\nData 作者首先构建了一个任务集合，然后基于这些任务，来过滤或者生成对应的 QA pair, 接下来作者基于难度和质量来过滤，最后作者在每个 domain 上进行 RL 的训练来保证数据的有效性。\nReward modelling 作者构建了一个 reward system, 用于给不同 domain 的任务进行奖励。为了探究不同 domain 的 reward 对模型整体表现的影响，作者设计了一个 low-quality reward system，实验结果如下，可以看到模型在训练后期出现了 collapse 或者 reward hacking 的情况。\n[!tip] Observation 作者发现，不同任务的不同信号奖励会对模型最终表现产生巨大影响。\n在使用 LLM 提取模型回答的 answer 的时候，作者要求模型的最终答案为 \u0026lt;|begin_of_box|\u0026gt; {answer_content} \u0026lt;|end_of_box|\u0026gt; 这种形式，避免大语言模型提取出错。\n作者构建的 reward system 包含如下模块：\nshared verification functions: 包括格式检查等 domain specific modules: 与 domain 相关的模块 unit testing: 观测 domain 输出的分布，然后基于输出的分布来调整 reward logic 最终 reward 的 design 如下表所示\nRLCS 作者使用了 GRPO 来进行优化训练。作者发现，仅需 200 training steps 模型在一半以上问题的准确率就达到了 $90\\%$.\n为了提升模型的训练效率，作者提出了 Reinforcement Learning with Curriculum Sampling (RLCS), 也就是将课程学习与 RL 结合起来。作者基于模型训练情况动态调整训练样本的难度，来保证每个样本对模型训练的提升都是最大的。\n在 RLCS 训练之前，作者先使用模型评测得到每个样本的 pass@k, 然后根据结果将样本进行分类。在训练过程中，作者记录每个样本的 pass@k， 然后动态更新其分类结果。在采样的时候，作者基于难度来对样本进行加权，来降低太简单或者太难样本的采样概率。\n作者还提出了一些提高 RL 表现的方法：\nLarger batch size: 使用更大的 batch size 效果更好 Dynamic sampling expansion via ratio EMA: 作者定义了 naive 样本和高质量样本的比例，然后根据这个比例进行采样，来提升整体的采样效率。 Force answering: 与 Qwen3 一样，对于比较难的问题，当输出的 token 数过长时，作者对模型输出进行截断，然后让模型基于已有思考过程直接输出结果。 Discard KL loss: 与 DAPO 一样，作者也移除了 KL divergence loss Clip-higher: 与 DAPO 一样，作者通过修改超参数来提高模型的探索能力 [!tip] Observation\nRL 阶段的表现与 cold-start SFT 的表现并不完全相关。作者发现，cold-start SFT 取得更好的表现并不一定保证 RL 的表现就更好。 RL 阶段各个 domain 数据的影响是正交的，这与 cold-start SFT 阶段的结果不同。 作者还提出了一些提高训练稳定性的方法：\n提高 cold-start SFT 阶段数据的质量 移除 KL divergence loss 使用 top-p 为 1，而不是 0.9 (如 Qwen-LLM 中的设置)，top-p 为 1 可以 cover 整个 vocabulary，避免某些 token 没有参与训练 per-sample loss 和 per-token loss 没有显著区别，但是 per-sample loss 稳定性更高 在 cold-start SFT 阶段让模型学会格式要求，而不是在 RL 阶段加入 format reward Infra 在 infra 方面，作者做了如下改进：\nLoad balancing of sequence lengths across DP ranks. 平衡每个 rank 的 sequensequence length 以及 compute load Intra-rank training with sequence packing and gradient accumulation. 将 sequence packing 和 gradient accumulation 结合起来 Sample packing and reorganization within DP ranks. data packing Dynamic sampling expansion via ratio EMA. 平衡 naive 样本和高质量样本之间的比例 Evaluation Performance 可以看到，GLM-4.1V-Thinking 在多个 benchmark 上都达到了 SOTA.\nAblation Study 作者评估了一下不同 domain 的 RL 训练对模型整体表现的影响。作者使用不同的数据来进行训练，然后分析结果的相关性，结果如下\n实验结果发现：\n在某个 domain 上的训练会提高模型在其他 domain 上的表现 联合多个 domain 进行训练最终的表现会更好 作者还发现，不同的任务之间存在关联，如 GUI-agent 和 grounding 这两个任务是高度相关的。OCR \u0026amp; Chart 以及 GUI-agent 也是高度相关的。\nConclusion 在本文中，作者提出了 GLM-4.1V-Thinking, 一个 9B 的通用多模态 reasoning model, 作者通过构建高质量的数据，多阶段的训练，让模型在 reasoning 和 non-reasoning 任务上均超过了已有的 SOTA 模型的表现\n作者认为，模型有以下局限：\n模型并没有提升 reasoning 的质量，在某些情况下，模型结果正确，但是中间过程错误。作者分析这是因为目前只有针对结果的奖励机制。因此，未来需要能够评估中间过程的 reward model RL 的训练具有不稳定性，作者发现模型对设置比较敏感，这对 RL 训练的 scaling 提出了挑战。 模型在复杂场景下表现依旧不太好，比如图片中有物体遮挡的情况，这些感知误差也会影响最终的 reasoning 表现。作者认为如何同时提高 perception 和 reasoning 的表现是一个需要研究的课题。 未来的工作有：\n如何构建更好的 reward 机制，来同时评估结果和中间过程。从而防止模型 reward hacking 如何基于 multimodal training 来提升模型在纯文本任务上的表现。探究 visual reasoning 和 text reasoning 如何互相促进也是一个课题 如何更好评估模型的表现，即如何更好评估模型的 failure modes, 比如幻觉，short reasoning 等 References Arxiv Huggingface Appendix GLM-4.1V-Thinking 的 patch merger 代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Glm4vVisionPatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, hidden_act: str, bias: bool = False) -\u0026gt; None: super().__init__() self.proj = nn.Linear(dim, dim, bias=bias) self.post_projection_norm = LayerNorm(dim) self.gate_proj = nn.Linear(dim, context_dim, bias=bias) self.up_proj = nn.Linear(dim, context_dim, bias=bias) self.down_proj = nn.Linear(context_dim, dim, bias=bias) self.act1 = nn.GELU() self.act_fn = ACT2FN[hidden_act] def forward(self, hidden_state: torch.Tensor) -\u0026gt; torch.Tensor: hidden_state = self.proj(hidden_state) hidden_state = self.act1(self.post_projection_norm(hidden_state)) return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) class Glm4vVisionModel: def __init__(self, config): self.downsample = nn.Conv2d( in_channels=config.hidden_size, out_channels=config.out_hidden_size, kernel_size=config.spatial_merge_size, stride=config.spatial_merge_size, ) self.merger = Glm4vVisionPatchMerger( dim=config.out_hidden_size, context_dim=config.intermediate_size, hidden_act=config.hidden_act ) def forward(self, hidden_states): ... hidden_states = hidden_states.view(-1, self.spatial_merge_size, self.spatial_merge_size, hidden_states.shape[-1]) hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size) hidden_states = self.merger(hidden_states) ","date":"2025-07-14T10:32:04+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-glm-4.1v-thinking/","title":"Notes on GLM-4.1V-Thinking"},{"content":"Qwen 在 2025 年 1 月提出了 Qwen2.5-1M，一个拥有 1M 上下文长度的大语言模型系列。包含 7B，14B 两个开源模型以及 API 模型 Qwen2.5-Turbo. 主要改进方法包括长上下文数据合成，渐进式预训练以及多阶段 post-training 等。作者还对 inference 进行了优化，提高了 inference 的效率。\nMethod 架构上，Qwen2.5-1M 与 Qwen2.5 的架构一致，Qwen2.5-1M 包括 7B，14B 两个 size，还包括一个基于 MoE 的 API 模型 Qwen2.5-Turbo，不同的点在于，Qwen2.5-1M 的上下文长度为 1M，最大生成长度为 8K\nPretraining Data 作者首先从 CC, arxiv, book, code repositories 等 domain 收集了原始数据。但是，作者发现，原始数据的局部相关性强，但是全局相关性弱。因此，作者基于原始数据进行了增广，来提高数据的长上下文依赖关系。具体有三个任务：\nFill in the middle: FIM 是 openAI 提出来的一个做法，核心思想就是将填空类问题转化为 next-token-prediction 的问题。通过这种方式，作者希望提升模型理解长上下文依赖的能力 Keyword-based and Position-based retrieval: 基于 keywords 或者 position 来找到对应的 paragraph，这个任务的目的是提高模型识别并连接相关信息的能力 Paragraph Reordering: 对输入的 paragraphs 进行随机打乱，然后要求模型重新组织段落的关系 Training 作者将训练拆分为了 5 个 stage：\nstage 1 和 stage 2 与 Qwen2.5 的训练过程一致，stage 1 的上下文长度为 4096，stage 2 的上下文长度为 32768, 训练时，作者使用了 ABF 技巧来将 RoPE 的 base frequency 从 10,000 调整到了 1,000,000. stage 3, stage 4 和 stage 5 分别将模型的上下文长度扩展到了 65,536 tokens, 131,072 tokens 以及 262,144 tokens, 对应的 RoPE base frequency 分别为 1M, 5M 和 10M. 训练时，作者使用了 75% 的长文本和 25% 的短文本，这样可以保证模型在短文本任务上的表现 最后，作者在评估了一下每个 stage 的表现，结果如下表所示\nTraining Length RULER Avg. 4K 8K 16K 32K 64K 128K 32,768 Tokens 82.3 96.8 94.7 95.9 92.2 76.4 37.6 65,536 Tokens 86.8 96.5 95.5 93.6 92.5 86.7 56.0 131,072 Tokens 92.5 96.5 95.9 93.0 92.6 93.0 83.8 262,144 Tokens 92.7 95.6 93.8 93.1 94.1 91.9 87.6 可以看到，随着训练的上下文长度的提升，模型在更长上下文下的能力也有提升，说明模型具有一定的泛化性。\nPost-training Post-training 阶段与 Qwen2.5 一样，也分为了 SFT 和 RL 两个阶段。\n在 SFT 阶段，作者从预训练预料中选择了一部分长文档的片段，然后让 Qwen2.5 来生成对应的 query，query 类型包括 summarization, information retrieval, multi-hop QA 等任务。接下来，作者使用 Qwen-Agent 框架基于全文来回答这些问题。最后，作者基于生成的 query，全文，以及模型产生的回答作为训练数据。\nSFT 训练时，作者拆分为了两个 stage。 stage 1 作者在 32768 的上下文上进行训练，来提高模型短文本回答能力。第二个阶段，作者混合了 262,144 和 32768 上下文长度的训练数据。\nRL 训练时，与 Qwen2.5 不一样的是，作者进使用了 offline RL，也就是 DPO。作者仅在 8192 的上下文长度上面进行训练。作者认为，长上下文的 RL 训练是非常耗时的，并且作者发现，短文本上进行 RL 的训练之后，模型在长文本上的表现也能得到提升。结果如下表所示\nModel Before RL After RL Qwen2.5-7B-Instruct-1M 7.32 8.08 (+0.75) Qwen2.5-14B-Instruct-1M 8.56 8.76 (+0.20) Qwen2.5-Turbo 7.60 8.34 (+0.74) Inference 前面是训练部分的优化，主要是提升模型的上下文能力。接下来，作者详细介绍了如何在 Inference 阶段提升整体的推理效率和减少内存占用。\nLength Extrapolation 与 Qwen2.5 一样，Qwen2.5-1M 也是用了 Dual Chunk Attention 和 YARN 来在推理阶段扩展模型的上下文长度，作者做了如下实验，来对比 Qwen2.5, Qwen2.5-1M 加上 DCA 之后的影响\n结果显示，Qwen2.5-1M 的表现比 Qwen2.5 更好，并且加上 DCA 之后，两者的表现都有进一步的提升。\nSparse Attention 为了进一步提高计算效率，作者基于 MInference 来加速 perfilling phase. 并结合了 前面的技巧来防止模型性能下降。\nMInference MInference 的主要思想就是在长上下文中，有一些 critical token 对最终结果的影响是更大的。因此我们可以识别出这些 critical token 并只计算这些 token 对应的 attention score. 这些 critical token 对应的 pattern 被称为 Vertical-Slash pattern.\n为了识别出这个 pattern，作者首先进行离线搜索，来决定最优的 configuration。这个 configuration 决定了 attention 应该如何计算。在 Inference 阶段，MInference 首先计算最后一个 query 和前面所有 key 的 attention，然后基于 configuration 来动态选择 pattern。通过 MInference，我们可以降低 10 倍以上的内存和算力消耗。\nIntegrating with Chunked prefill 但是 MInference 的问题在于，整个 sequence 是并行处理的，这会导致内存占用持续上升。为了解决这个问题，作者提出了 chunked prefilling 的技巧，来降低 VRAM 的消耗。具体做法就是，将整个 sequence 分为若干个 chunk，然后每个 chunk 里，选取最后 64 个 token 作为 query，在每个 chunk 中分别识别出 critical token，这样就降低了 MInference 的内存占用\n接下来，作者在集成 DCA 的时候，发现性能有所下降。作者认为，这是由于 DCA 的 position id 信息不连续所导致的，为了解决这个问题，作者在选择 critical token 的时候，使用了连续版的 position id 信息。在最终推理的时候，还是使用 DCA 本身的位置信息。\nSparsity refinement 前面提到，MInference 需要先进行离线搜索决定最优的 configuration，但是对于 1M token 的上下文，这个过程还是非常耗时的。因此，作者构建了一个加速离线搜索的方法，具体做法就是定义两个 attention score，一个是 full attention, 另一个是 sparse attention， 然后计算两者的差值，如果说相差比较小，则说明 critical token 抓住了全局信息，这个配置是有效的。其公式定义如下：\n$$ \\mathrm{Attention\\_Recall} = \\exp\\left(\\log\\sum_{0\\leq j\\leq i}\\exp \\left(\\frac{q^Tk_j}{\\sqrt {d}}\\right) - \\log\\sum_{j\\in\\mathcal{critical}}\\exp \\left(\\frac{q^Tk_j}{\\sqrt {d}}\\right)\\right) $$Attention Recall 越高，说明选取的 critical token 越好，其 configuration 也就越好。\n作者进一步分析了 sparse attention 对 accuracy 的影响，结果如下\n可以看到，仅使用 MInference 会导致模型性能 下降，但是加入 refinement 之后，模型的表现基本上和 full attention 差不太多。\nInference Engine Kernel Optimization 作者还对 inference engine 进行了优化，作者使用 BladeLLM 作为 Qwen2.5-1M 的推理引擎。\n作者主要做了两点优化，第一是对 sparse attention kernel 进行了优化，提高了 sparse attention 的计算效率，结果发现，在 1M 的上下文下，BladeLLM 比 flash attention 要快 27.8 倍。\n第二是针对 MoE kernel 的优化。作者发现，decoding 的表现是与 memory access speed 相关的。具体来讲，当 batch size 超过 32 之后，获取模型参数成了效率的瓶颈。因此，作者使用了一系列技巧来提高 memory access 的效率\nPipeline parallelism 作者还对 Chunked pipeline parallelism 进行了优化，Chunked pipeline parallelism 的问题在于，在长上下文的场景下，不同长度的 chunk 会对 attention 的计算时间产生很大影响。不同的计算时间会产生 pipeline bubbles.\nBladeLLm 使用了 Dynamic Chunked pipeline parallelism 来解决这个问题，该方法通过计算复杂度来调整每个 chunk 的大小，进而使得最终的处理时间尽可能一致\nScheduling 作者还在 Scheduling 上进行了优化，已有的推理引擎主要分为四个模块：API server, scheduler, model runner 以及 decoder\n已有方法的问题在于，non-GPU 的操作会占用大量时间，导致 GPU 利用率非常低。因此，作者在 BladeLLM 中进行了改进，使用了 Totally Asynchronous Generator (TAG) 的架构，主要有：\nScheduler：动态分配 KV cache，类似于 speculative sampling, 而不必等前面的结果完成 Runner: 基于 Scheduler 分配的任务直接进行处理，处理完之后直接处理下一个任务 Decoder：基于 token id，进行解码，然后发送给前端的 API server Evaluation 作者主要在三个 benchmark 上进行了评测：\nRULER: RULER 是 Needle-in-ahaystack 任务的一个扩展笨笨，其要求模型从不相关的上下文中找到多个 \u0026ldquo;needles\u0026rdquo; 或者回答多个问题，数据最长为 128K tokens. LV-Eval: LV-Eval 要求模型从上文本中同时理解多个 evidence fragments，数据最长为 256K tokens Longbench-Chat: 评估模型在长上下文下与人类偏好对齐的程度，数据最长为 100K tokens Qwen2.5-1M 与 Qwen2.5 的对比表现如下\n可以看到，相比于 Qwen2.5，Qwen2.5 模型的表现有了大幅度的提升。\nConclusion 在本文中，作者提出了 Qwen2.5-1M 系列大语言模型，包括 7B，14B 两个 size，以及一个 MoE 架构的 API 模型 Qwen2.5-Turbo。作者在训练和推理两方面进行了改进，最终将模型的上下文长度扩展到了 1M。从现在的角度来看，不管是 Reasoning model 还是 agent 的训练都依赖 long Context 作为基础能力。\nReferences Qwen2.5-1M Technical Report ","date":"2025-07-12T11:00:47+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen2.5-1m/","title":"Notes on Qwen2.5-1M"},{"content":"2024 年 12 月 Qwen 发布了 Qwen 2.5 系列大语言模型，包括 7 个 dense 模型以及两个 MoE 模型，Qwen2.5 在 pre-training 阶段使用了 18T token 进行。在 post-training 阶段使用了 1M 的样本，还使用了 DPO 以及 GRPO 来进行 RL 的训练\nQwen2.5 主要在以下方面进行了改进\n模型方面，提供了更多的 size，Qwen2 中只有 0.5B, 1.5B, 7B, 72B 四个 size, 在 Qwen2.5 中，加入了 3B, 14B 和 32B 三个 size 的模型 数据方面，pre-training 阶段使用了 18T 的 token， post-training 阶段使用了 1M 的样本 功能方面，Qwen2.5 支持更长的上下文长度（8K），支持结构化输入和输出，拥有更强的工具调用能力。 Method Architecture 模型架构这方面，Qwen2.5 和 Qwen2 的模型架构是一致的，tokenizer 页没有太大变化。为了支持工具调用，作者额外增加了 18 个 control token\nPre-training data Qwen2.5 从以下方面提高了预训练数据的质量\nBetter data filtering: 使用 Qwen2-Instruct 来过滤掉质量的数据，然后从多维度对训练数据进行打分，从而提高数据的质量 Better math and code data: 加入了 Qwen2.5 Math 以及 Qwen2.5 Coder 的训练数据来提高模型的数学和代码能力 Better synthetic data: 作者使用 Qwen2-72B-Instruct 以及 Qwen2-Math-72B-Instruct 来合成 math, code, knowledge domain 的数据，然后通过过滤以及 Qwen2-Math-RM-72B 来提高数据的质量 Better data mixture: 作者使用 Qwen2-Instruct 来分类，然后平衡不同 domain 的数据分布。作者发现 e-commerce, social media 以及 entertainment 的数据重复性高，且大多都是机器生成的。而 technology, science 以及 academic research 等 domain 的数据质量更高。作者对不同 domain 的数据进行了上采样或者下采样。 基于这个过程，作者一共收集了18T tokens\nHyper-parameters 作者构建了针对超参数的 scaling law，即决定最优的训练超参数如 batch size, learning rate 等\n作者通过实验得到了 model size $N$ 以及 pre-training data size $D$ 与 learning rate $\\mu_{opt}$ 和 batch size $B_{opt}$ 之间的关系。\nLong context pre-training 为了提升模型的上下文长度，作者将 pre-training 拆分为两个 stage，第一个 stage 的上下文长度为 4096， 第二个 stage，作者将上下文长度从 4096 扩展到 32768.\n在提升模型上下文过程中，作者使用 ABF 技巧将 Position Encoding 的 base frequency 从 10,000 提升到了 1,000,000.\n对于 Qwen2-5-Turbo，作者实现了渐进式上下文长度扩展策略，模型上下文长度扩展经历四个阶段：32768, 65536, 131072 到最终的 262,144. 此时，RoPE 的 base frequency 为 10,000,000. 在训练的每个阶段，作者都使用了 40% 的长文本以及 60% 的短文本，以保证在扩展模型上下文长度的同时，还能保持模型在不同上下文长度下的表现。\n为了提高模型在 inference 时的长上下文表现，作者使用了 Dual Chunk Attention 和 YARN 两个技巧。通过这两个技巧，作者将 Qwen2.5-Turbo 的上锈阿文扩展到了 1M，将其他模型的上下文长度扩展到了 131072.\nPost-training Qwen2.5 的 post-training 分为两个大的 stage: SFT 和 RL，其中 RL 又分为两个小的 stage，分别是 offline RL 和 online RL\n在 SFT 阶段，作者主要做了以下改进：\nLong-sequence generation: 作者将 Qwen2.5 的输出长度提升到了 8192, 为了扩展模型输出的长度，作者构建了 Long-response 数据集，然后基于 back-translation 来生成对应的 query，最后使用 Qwen2 来过滤低质量的数据 Math: 作者在 SFT 阶段加入了 Qwen2.5-Math 的 CoT 数据，包括公开数据集，K12 问题集一集合成数据等。作者通过 rejection sampling 以及 annotated answers 来生成 CoT 过程 Code: 作者加入了 Qwen2.5-Coder 的 SFT 数据，作者基于多个 agent 来生成多样化高质量的 Instruction, 然后还从 code-related QA website 以及 Github 上获取数据来扩展数据集。对于最终的数据，作者使用了 sandbox 来保证代码的质量 Instruction following: 作者构建了一个基于 code 的验证框架，让 LLM 同时生成 Instruction 和对应的验证代码，验证的单元测试。最后，通过 rejection sampling 来得到最终的数据集 Structured Data Understanding: 作者还构建了针对 tabular QA, fact verification, error correction 以及 structured understanding 等数据集。作者在回答中加入 CoT，作者提高了模型对 structured data 的理解能力 Logical Reasoning: 作者构建了 70,000 个不同 domain 的 query，有多种格式，覆盖了 analogical reasoning, causal reasoning 等 domain Cross-Lingual Transfer: 作者使用了一个翻译模型，来将 Instruction 转换到 low-resource language 上，进而提高模型在对应语种上的表现 Robust System Instruction: 作者构建了不同的 system prompt 用于提升 system prompt 的多样性。作者发现，使用不同的 system prompt 可以减少模型的 variance, 提高模型的 robustness. Response Filtering: 作者使用了多种自动化标注方法来保证最终 response 的质量 最终，作者一共收集到 1M 的 SFT 样本，模型训练了两个 epoch\n在 RL 阶段，作者首先基于 SFT model 来进行采样，然后将高质量的回答作为正样本，低质量的回答作为负样本，通过这个过程，一共采集到了150K的样本。最后，作者使用 DPO 来进行训练。\n然后，作者进行了 online stage 的 RL 训练，这一阶段主要是对齐模型与人类的价值观。这一阶段的数据包括公开数据集，私有数据集。作者使用不同的 checkpoint 来进行采样，然后作者使用 GRPO 来进行训练.\nLong Cotnext Fine-tuning 作者还针对 Qwen2.5-Turbo 做了额外的 post-training, 来进一步提高其在长上下文下的表现。\n在 SFT 阶段，作者使用了一个两阶段方法，第一阶段仅在短文本上进行训练（上下文长度为 32768），这一阶段的训练数据与其他 Qwen2.5 的模型训练数据相同。第二个阶段，作者混合了短文本和长文本（262144）来进行训练，来提高模型在长上下文情景下的指令跟随能力\n在 RL 阶段，作者使用了和其他 Qwen2.5 模型相同的训练策略。作者认为：\n长上下文下训练 RL 代价很大 reward model 更偏向于长文本 RL 尽管只在短文本上进行训练，其还是可以提高模型在长上下文下的表现。 Evaluation 我们仅关注 instruction 版本的 72B,32B 和 7B 模型\n可以看到，Qwen2.5 72B 模型表现和 LLaMA3.1 405B 表现差不多，其他两个 size 的模型基本上达到了 SOTA\n最后，作者评估了一下 DCA+YaRN v.s. Full attention 的表现，结果如下图所示\n可以看到，使用 DCA+YaRN 之后，模型的推理效率比 full attention 要快 3-4 倍。\nConclusion 在本文中，作者提出了 Qwen2.5 系列大语言模型，包括 7 个 dense 模型以及两个 MoE 模型，作者详细介绍了模型的 pre-training 和 post-training. 评测结果发现 Qwen2.5 模型基本上达到了 SOTA.\n作者认为，未来工作有：\n使用更多更多样化的 pre-training 和 post-training 数据 多模态大模型的构建，特别是 omni-modal 提高模型的 Reasoning 能力 References Arxiv ","date":"2025-07-12T10:51:42+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen2.5/","title":"Notes on Qwen2.5"},{"content":"Dual Chunk Attention (DCA) 由阿里巴巴在 2024 年 9 月份提出，DCA 是一个无需训练的，扩展 LLM 上下文长度的方法，后续，DCA 被应用于 Qwen2, Qwen2.5, Qwen2.5-1M 以及 Qwen3 中，与 YARN 一起作为扩展模型上下文的有效手段\nIntroduction 提升 LLM 上下文长度的方法可以分为两类：一类是 training-free 的，包括 LM-infinite 和 StreamingLLM 等，这些方法以损失 long range dependency 为代价来保持较低的 perplexity。另一类为了保留全局信息，则是通过外插来扩展模型的上下文，主要工作我们在 YARN 中已经回顾了。\n第二类方法的问题在于，其依赖训练，在 training-free 的 setting 下，这些方法也会导致 perplexity 的上升\n因此，在本文中，作者就提出了 Dual Chunk Attention (DCA) ，一个无需训练的，扩展 LLM 上下文长度的方法。DCA 的主要做法是将 attention 的计算进行分块，这样就可以提高计算效率。\n通过实验，作者给出了三点关键发现：\nExtrapolation： DCA 可以在无需训练的情况下，将 LLM 的上下文提升到 32K，而不导致 Perplexity 大幅度增加 Orthogonality： DCA 可以和其他方法一起使用，如 YARN, 这一点已经在 Qwen2.5-1M 以及 Qwen3 中得到了应用 Long Context Understanding: DCA 可以在无需训练的情况下，在长上下文设置下，达到已有 SOTA 模型的表现 Preliminary 对于一个长度为 $L$ 的 token 序列，我们首先定义对应的 position id 如下\n$$ P_{\\mathbf{q}} = [0,1,\\dots,L-1],\\quad P_{\\mathbf{k}} = [0,1,\\dots,L-1] $$然后，对于第 $i$ 个位置和第 $j$ 个位置的 token，其 attention score 定义为：\n$$ \\langle f(\\mathbf{q}, i), f(\\mathbf{k}, j)\\rangle =\\langle R_{\\theta,i}\\mathbf{q}, R_{\\theta,j}\\mathbf{k}\\rangle =\\mathbf{q}^TR_{\\theta, i-j}\\mathbf{k} $$具体细节参考 Position Encoding 中的 RoPE 部分介绍。这里面的关键在于，最后的结果只与相对位置 $i-j$ 相关，而与绝对位置 $i$ 和 $j$ 无关。因此，我们可以用一个相对位置矩阵 $M\\in\\mathbb{R}^{L\\times L}$ 来表示这个信息，其中 $M_{ij}=P_{\\mathbf{q},i}- P_{\\mathbf{k},j}$ 代表了第 $i$ 个位置的 query $\\mathbf{q}$ 和第 $j$ 个位置的 key $\\mathbf{k}$ 的相对位置信息，其示意图如下所示\n原始版本的 RoPE 的问题在于，在训练时，模型没有见过更长的上下文，因此其泛化性也最差，这一点在 YARN 已经得到了验证\nMethod DCA 的关键在于将 sequence 分割为若干个 Chunk，然后将 attention 的计算拆分为三个部分：\nintra-chunk：负责计算每个 chunk 内部的 attention inter-chunk ：负责计算 chunk 之间的 attention successive-chunk：负责计算相邻两个 chunk 之间的 attention 为了更好理解 DCA，我们接下来假设 $L=12$, 这时我们有\n$$ P_{\\mathbf{q}} = [0,1,\\dots,11],\\quad P_{\\mathbf{k}} = [0,1,\\dots,11] $$Intra-Chunk Attention 我们首先定义个超参数 chunk size $s\u003e0$, 然后我们将我们的 sequence 分割成 $L/s$ 个 chunk，然后每个 chunk 重新进行编号，就得到了如下的 position id\n$$ P_{\\mathbf{q}}^{Intra} = [0,1,\\dots,L-1]\\mod s,\\quad P_{\\mathbf{k}}^{Intra} = [0,1,\\dots,L-1]\\mod s $$接下来，我们定义 intra-chunk 的相对位置矩阵 $M$， 此时，我们仅在每个 chunk 内部进行计算 attention，即\n$$ M[i][j] = \\begin{cases} P_{\\mathbf{q},i}^{Intra} - P_{\\mathbf{k},j}^{Intra},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor = \\lfloor P_{\\mathbf{k},j} /s\\rfloor ,\\\\ 0,\u0026\\text{otherwise} \\end{cases} $$在上面的例子中，假设 $s=6$, 那么我们新的 position id 就变成了\n$$ \\begin{aligned} P_{\\mathbf{q}}^{Intra} = [0,1,2,3,4,5,0,1,2,3,4,5]\\\\ P_{\\mathbf{k}}^{Intra} = [\\underbrace{0,1,2,3,4,5}_{\\text{Chunk 0}},\\underbrace{0,1,2,3,4,5}_{\\text{Chunk 1}}] \\end{aligned} $$对其进行可视化，我们就得到\nInter-Chunk Attention 接下来，我们来看一下不同 chunk 之间如何计算彼此的 attention。在 Intra-chunk attention 计算中，我们忽略了跨 chunk 的信息，而且，由于现在的 position id 不再是单调递增的了，我们直接使用 $P_{\\mathbf{q}}^{Intra}$ 和 $P_{\\mathbf{k}}^{Intra}$ 给出的位置信息不对，这也是为什么我们在 Intra-chunk attention 中要求 query 和 key 在同一个 chunk 中才能计算的原因。\n为了解决这个问题，作者构建了一个新的 position id。首先我们引入一个新的超参数 $c\u003e\\max_i P_{\\mathbf{q},i}$, $c$ 代表了模型预训练时的上下文长度，如 4096。\n接下来，基于 $c$, 我们定义新的 position id 如下：\n$$ P_{\\mathbf{q}}^{Inter} = [c-1,c-1,\\dots,c-1]\\in\\mathbb{R}^s,\\quad P_{\\mathbf{k}}^{Inter} = P_{\\mathbf{k}}^{Intra} $$ 注：这里的 $P_{\\mathbf{q}}^{Inter}$ 指的是某一个 chunk 中的 position id，每个 chunk 中的 position id 都相同，请参考例子理解，后面不再赘述。\n也就是说，在计算跨 chunk 的 attention 的时候，我们直接把 query 的 position id 设置为最大值，然后 key 的 position id 依然使用 intra-chunk 的位置信息。由于 $\\max_i P_{\\mathbf{k},i}=s-1$, 因此我们有\n$$ M[i][j] = P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j}^{Inter} = c - 1 - P_{\\mathbf{k},j}^{Inter}\\geq c - 1 - (s- 1) \\geq c-s. $$最后，我们对于 inter chunk 的位置矩阵 $M$ 定义如下：\n$$ M[i][j] = \\begin{cases} P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j}^{Inter},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor \\neq \\lfloor P_{\\mathbf{k},j} /s\\rfloor ,\\\\ 0,\u0026\\text{otherwise} \\end{cases} $$在上面的例子中，当 $c=10$ 时，$c-1=9$, 我们有\n$$ P_{\\mathbf{q}}^{Inter}=[\\underbrace{9,9,9,9,9,9}_{\\text{Chunk 0}},\\underbrace{9,9,9,9,9,9}_{\\text{Chunk 1}}] $$对其进行可视化，得到\nSuccessive-Chunk Attention 现在我们既可以计算 intra-chunk，也可以计算 inter-block 的 attention，但是问题是对于相邻的 chunk，其位置信息不对了，从上面的可视化中，我们可以看到，当 $P_{\\mathbf{q},i}=6$ , $P_{\\mathbf{k},j}=5$ 时，我们有\n$$ P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j}^{Inter}=9-5=4\\neq 1 = P_{\\mathbf{q},i}-P_{\\mathbf{k},j} $$也就是说，inter-block 的 attention 会破坏原有的相对位置信息，因此我们就通过 successive chunk attention 来解决这个问题，使得 $P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j}^{Inter}\\approx P_{\\mathbf{q},i}-P_{\\mathbf{k},j}$.\n作者发现，这个问题不是所有的 chunk 都有，而是只存在于相邻的 chunk 中，因此，作者又加入了一个超参数 $w\u003e0$ 代表了 local window size，我们可以直接将其设置为 $c-s$, 通过这个 local window，我们调整对应的 position id 如下：\n$$ P_{\\mathbf{q}}^{Succ} = [\\overbrace{s,s+1,\\dots,s+w-1}^{w \\text{ elements}},c-1, \\dots,c-1]\\in\\mathbb{R}^s,\\quad P_{\\mathbf{k}}^{Succ} = P_{\\mathbf{k}}^{Inter} $$对于 successive chunk 的位置矩阵 $M$ 定义如下：\n$$ M[i][j] = \\begin{cases} P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j}^{Inter},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor=1 ,\\\\ 0,\u0026\\text{otherwise} \\end{cases} $$在上面的例子中，我们设置 $w=4$, 就得到\n$$ P_{\\mathbf{q}}^{Succ}=[\\underbrace{6,7,8,9,9,9}_{\\text{Chunk 0}},\\underbrace{6,7,8,9,9,9}_{\\text{Chunk 1}}] $$对其进行可视化，得到\nComputation 接下来，我们把所有的改进放在一起，就得到\n$$ M[i][j] = \\begin{cases} P_{\\mathbf{q},i}^{Intra} - P_{\\mathbf{k},j},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor=0 ,\\\\ P_{\\mathbf{q},i}^{Succ} - P_{\\mathbf{k},j},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor=1 ,\\\\ P_{\\mathbf{q},i}^{Inter} - P_{\\mathbf{k},j},\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor\u003e1. \\end{cases} $$基于上面的位置矩阵 $M$， 我们再依次计算对应的 attention score\n$$ \\langle f(\\mathbf{q}, i), f(\\mathbf{k}, j)\\rangle = \\begin{cases} \\langle f(\\mathbf{q}, P_{\\mathbf{q},i}^{Intra}) , f(\\mathbf{k}, P_{\\mathbf{k},j})\\rangle,\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor=0 ,\\\\ \\langle f(\\mathbf{q}, P_{\\mathbf{q},i}^{Succ}) , f(\\mathbf{k}, P_{\\mathbf{k},j})\\rangle,\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor=1 ,\\\\ \\langle f(\\mathbf{q}, P_{\\mathbf{q},i}^{Inter}) , f(\\mathbf{k}, P_{\\mathbf{k},j})\\rangle,\u0026\\text{if}\\lfloor P_{\\mathbf{q},i}/s \\rfloor - \\lfloor P_{\\mathbf{k},j} /s\\rfloor\u003e1. \\end{cases} $$Code 首先是 RotaryEmbedding 部分的修改\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 class DCARotaryEmbedding(nn.Module): def __init__(self, max_len, chunk_size, local_window): self.max_len = max_len self.chunk_size = chunk_size self.local_window = local_window self.inv_freq = ... def forward(self, x): q_t = torch.arange(self.chunk_size) qc_t = (q_t + self.chunk_size).clamp(max=self.chunk_size) k_t = torch.arange(seq_len) % self.chunk_size q_freqs = torch.outer(q_t, self.inv_freq) # seq_len x dim/2 qc_freqs = torch.outer(qc_t, self.inv_freq) k_freqs = torch.outer(k_t, self.inv_freq) # seq_len x dim/2 q_emb = torch.cat((q_freqs, q_freqs), dim=-1) # seq_len x dim qc_emb = torch.cat((qc_freqs, qc_freqs), dim=-1) k_emb = torch.cat((k_freqs, k_freqs), dim=-1) # seq_len x dim # compute related sin, cos return q_sin, q_cos, qc_sin, qc_cos, k_sin, k_cos attention 计算时的逻辑\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Attention(nn.Module): def forward(...): key_states = apply_rotary_pos_emb(key_states, k_cos, k_sin, position_ids) q_states_intra = apply_rotary_pos_emb(query_states[:, :, :chunk_len, :], q_cos, q_sin, position_ids[:, :chunk_len]) k_states_prev = key_states[:, :, :chunk_len, :] v_states_prev = value_states[:, :, :chunk_len, :] # first chunk flash_result = do_flash_attn(q_states_intra, k_states_prev, v_states_prev) flash_results.append(flash_result) remain_len = kv_seq_len - chunk_len while remain_len \u0026gt; 0: flash_per_chunk = [] begin = kv_seq_len - remain_len curr_chunk_len = min(chunk_len, remain_len) end = begin + curr_chunk_len # current chunk, intra-chunk attention q_states_intra = apply_rotary_pos_emb(query_states[:, :, begin:end, :], q_cos, q_sin, position_ids[:, begin:end]) k_states_intra = key_states[:, :, begin:end, :] v_states_intra = value_states[:, :, begin:end, :] flash_result = do_flash_attn(q_states_intra, k_states_intra, v_states_intra) flash_per_chunk.append(flash_result) # successive chunk attention q_states_succ = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin, position_ids[:, begin:end]) flash_result = do_flash_attn(q_states_succ, k_states_prev, v_states_prev, False) flash_per_chunk.append(flash_result) # inter chunk attention if begin - (k_states_prev.size(-2)) \u0026gt; 0: prev_len = k_states_prev.size(-2) q_states_inter = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin, position_ids[:, chunk_len - 1][:, None].repeat(1, curr_chunk_len)) k_states_inter = key_states[:, :, :begin - prev_len, :] v_states_inter = value_states[:, :, :begin - prev_len, :] flash_result = do_flash_attn(q_states_inter, k_states_inter, v_states_inter, False) flash_per_chunk.append(flash_result) flash_results.append(flash_per_chunk) k_states_prev = k_states_intra v_states_prev = v_states_intra remain_len = remain_len - chunk_len # merge the final results attn_output = merge_attn_outputs(flash_results) Evaluation 作者还分析了一下 DCA 的效率，结果如下\n可以看到，在 flash attention 的基础上加上 DCA 之后，内存占用和推理时间并没有发生太大变化\n作者还分析了三种 attention 对结果的贡献，如下图所示\n结果显示，intra block 的 perplexity 是最低的，但是其在下游任务上表现是最差的。当三者结合在一起之后，perplexity 和下游任务上的表现都是最好的。\nConclusion 本文中，我们回顾了 Qwen 系列扩展大模型上下文的方法 Dual Chunk Attention （DCA） 通过将 attention 切分成更小的 chunk，然后将 attention 的计算分为 intra-chunk，inter-chunk 和 successive-chunk，分别处理 chunk 内部，chunk 之间以及相邻 chunk 的 attention，通过这种方式，在无需训练的情况下，我们可以有效将模型上下文长度扩展 4 倍以上。\nReferences Arxiv Github ","date":"2025-07-12T10:41:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/dual-chunk-attention/","title":"Dual Chunk Attention"},{"content":"Introduction 2024 年 9 月 Qwen 发布了 Qwen2 系列技术报告，Qwen2 系列包括 4 个 dense 模型（0.5B, 1.5B, 7B, 72B）和一个 MoE 模型（总参数 57B，激活参数 14B），作者主要在架构，数据和长上下文上进行了改进。\nMethod Model 对于 dense 模型，Qwen2 在 Qwen-LLM 的基础上做了如下改动：\n使用 Group Query Attention (GQA) 替换 MHA，来优化 KV cache，提高 throughput 使用 Dual Chunk Attention 和 YARN 来提高模型上下文长度和训练效率 其余与 Qwen 一致，包括 SwiGLU，RoPE，RMSNorm 和 pre-normalization\n对于 MoE 模型，Qwen2-MoE 基于 Qwen1.5 进行了改进，主要是 3 点：\n作者使用了更细粒度的专家个数，作者认为细粒度的专家可以提供更丰富的 combination，这一点与 olmoe 的结论相同 与 DeepSeek-MoE 一样，作者使用了共享专家和路由专家 作者使用了类似 upcycling 的方法来初始化模型。假设一共有 $n$ 个专家，每个专家的维度为 $h_E$, 原始 dense 模型的维度为 $h_{FFN}$, 那么我们会把 dense 模型的参数复制 $[nh_E/h_{FFN}]$ 次，这样就可以扩展到任意个数的 MoE 模型上。作者还对参数进行 shuffle，来提高 diversity。最后，作者还对 50% 的参数进行随机初始化，来提高模型的 capacity。 模型配置如下\nConfiguration 0.5B 1.5B 7B 72B 57B-A14B Hidden Size 896 1,536 3,584 8,192 3,584 # Layers 24 28 28 80 28 # Query Heads 14 12 28 64 28 # KV Heads 2 2 4 8 4 Head Size 64 128 128 128 128 Intermediate Size 4,864 8,960 18,944 29,568 2,560 # Routed Experts - - - - 64 # Activated Experts - - - - 8 # Shared Experts - - - - 8 Embedding Tying True True False False False Vocabulary Size 151,646 151,646 151,646 151,646 151,646 # Trained Tokens 12T 7T 7T 7T 4.5T Pre-training 预训练阶段的数据基于 Qwen 和 Qwen1.5，数据处理策略如下：\n使用基于 heuristic 和 model-based 方法来过滤掉低质量的数据 加入了 code， math 和 multilingual 的数据 平衡了各个类别的数据分布 初始数据包括 12T token，经过过滤得到 7T token。作者发现，使用 12T token 进行训练，模型的表现不如使用 7B token 训练得到的模型效果好。因此除了 0.5B 的模型，其他模型使用的都是 7T 的 token\n对于 MoE 模型，作者使用了额外的 4.5T token 来进行预训练。\n在训练过程中，作者还加入了 multi-task instruction 数据，来提高模型的上下文学习能力和指令跟随能力。\n作者还将 Qwen2 模型系列的上下文长度从 4096 扩展到 32768，扩展过程中作了三个改动：\n加入了更多高质量的长上下文数据 将 RoPE 的 frequency 从 10,000 提升到了 1,000,000 使用了 YARN 来扩展上下文长度 使用了 Dual Chunk Attention 来优化 attention 的计算 Post-training post-training 包括 SFT 和 RLHF 两个阶段\n数据包括 SFT 数据和 RLHF 使用的偏好数据\n数据标注过程有：\n使用 InsTag 对数据进行打标 选取高质量的 instruction 构建了一个 self-evolution 策略，来扩展 instruction 数据 请人类来标注数据 作者还合成了一些数据，合成数据的过程如下：\nrejection sampling：对 LLM 进行多次采样，然后保留结论正确的数据作为 SFT 数据，以正确和错误的数据对作为偏好数据 Execution feedback：对于代码任务，使用 Python 来验证答案的正确性 Data Repurposing：对于写作类任务，以文档为输入，让 LLM 生成对应的 instruction Constitutional Feeback：基于预设的 principle 来生成回答 最终，SFT 数据包括 500, 000 条样本\nRLHF 的训练包括 offline stage 和 online stage，offline stage 就是用收集到的偏好数据。在 online stage，作者使用 reward model 来给输出的回答进行打分，然后再使用 DPO 进行训练。\n与 Qwen 不同，Qwen2 中作者使用了 Online Merging Optimizer 来解决因为 alignment 导致的性能降低\nConclusion 本文提出了 Qwen2 系列，在 Qwen2 中，首次使用了 GQA 代替 MHA，Qwen2 在上下文上做出了初步探索\nReferences Qwen2 tech report ","date":"2025-07-12T10:36:43+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen2/","title":"Notes on Qwen2"},{"content":"Qwen 在 24 年 1 月份发布了 Qwen1.5，包含 0.5B, 1.8B, 4B, 7B, 14B, 32B, 72B, 以及 110B 6 个 size，还有一个 MoE 模型。\n介绍 Qwen1.5 的主要特点：\n支持 12 中语言 统一支持 32768 tokens 上下文长度 。 提供 量化版本 （Int4、Int8、AWQ、GGUF）以适应低资源环境或部署需求。 训练过程使用了 DPO 以及 PPO 来进行对齐\nQwen1.5-MoE Qwen1.5-MoE 的激活参数为 2.7B，一共包含 64 个专家，其中激活 4 个专家，共享 4 个专家\n相比于 Qwen1.5-7B，去训练的 FLOPS 降低了 75%，inference 的速度提高了 174%\nQwen1.5-MoE 采用了改进的 MoE 架构，主要优化包括：\n细粒度专家（Fine-grained experts） ：通过将 FFN 层划分为多个片段，构建更多专家而不增加参数总量。 初始化策略（Upcycling） ：基于 Qwen-1.8B 初始化模型，并引入随机性以加速收敛。 路由机制（Routing Mechanism） ：在每个 MoE 层中使用 64 个专家，其中 4 个共享专家始终激活，60 个路由专家中有 4 个被激活，提高了灵活性和效率。 效率对比 作者对比了 throughput (requests processed per second) 以及 tokens per second (TPS):\nModel Throughput TPS Qwen1.5-7B-Chat 1.15 2298.89 Qwen1.5-MoE-A2.7B-Chat 2.01 4010.27 References Qwen1.5 MoE Qwen 1.5 ","date":"2025-07-03T17:37:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen1.5/","title":"Notes on Qwen1.5"},{"content":"Introduction YaRN (Yet Another RoPE extentionN method) 时23年9月EleutherAI等提出来的一个扩展LLM上下文长度的方法，后来被Qwen系列模型所应用。\nPreliminary 作者首先回顾了一下RoPE, 具体内容请参见上一篇blog。并使用了 $f_{W}(x_m, m, \\theta_d)$ 来表示RoPE：\n$$ f_{W}(x_m, m, \\theta_{d}) = \\Theta_{\\theta, m}^d W x_m $$其中 $\\Theta_{\\theta, m}^d\\in\\mathbb{R}^{d\\times d}$ 是多维旋转矩阵, $\\theta_d=[\\theta_{0,d},\\dots,\\theta_{(d-2)/2,d}]\\in\\mathbb{R}^{d/2}$, $\\theta_{i,d}=\\theta^{2i/d}$, $\\theta\u003e0$ 是一个超参数，RoPE中设置为 $\\theta=10000$, $W\\in\\{W_q,W_k\\}\\subset\\mathbb{R}^{d\\times d}$ 是对应的query/key projection layer的权重矩阵， $x\\in\\mathbb{R}^{d}$ 是输入的hidden states.\n接下来，作者定义了两个新的变量：\nscaling factor 假设预训练的上下文长度为 $L$, 扩展的上下文长度为 $L'\u003eL$, 则我们定义scaling factor 为\n$$ s = \\frac{L'}{L} $$易知 $s\u003e1$.\nwavelength 我们将 $\\lambda_d$ 定义为 $i$-th 维的RoPE embedding对应的 wavelength：\n$$ \\lambda_{i,d} = \\frac{2\\pi}{\\theta_{i,d}} = 2\\pi\\theta^{2i/d}, \\ i=0,\\dots,(d-2)/2 $$wavelength描述了对于第$i$ 个维度，RoPE旋转一周 ($2\\pi$) 所需要的上下文长度。\nUnified Perspective on Related Work 基于 $f_{W}(x_m, m, \\theta_d)$, 作者统一了已有的扩展上下文长度的方法，作者将不同的扩展方法使用一个通用函数 $f_W(x_m,g(m), h(\\theta_d))$ 来表示，这里 $g(m)$ 和 $h(\\theta_d)$ 分别代表了不同的长度外推方法所使用的变换。\nPosition Interpolation Position Interpolation (PI) 的核心思想在于，我们可以通过Interpolation，将超过预训练长度的文本给压缩到当前最大长度，以此来避免RoPE外推产生的问题，其对应的公式为\n$$ f'_W(x_m,m,\\theta_d) = f_W(x_m, \\frac{mL}{L'}, \\theta_d) $$其中 $L'\u003eL$ 为我们扩展之后的上下文长度， $L$ 为我们预训练的上下文长度。使用通用函数表示的话，我们有\n$$ g(m) = \\frac{m}{s},\\quad h(\\theta_d) = \\theta_d $$NTK-aware Interpolation PI的问题是，并没有考虑不同维度的wavelength。 基于NTK理论，DNN当输入维度比较低，且embedding缺少高频内容时，模型就会很难学习到高频信息。对应到RoPE里面，输入的token position id是低位信息（1维），而输出的RoPE是一个 $d$ 维的复杂向量。因此，当输入token非常相似却距离非常近时，RoPE就会丢失高频的细节信息\n因此，为了解决这个问题，作者对不同的维度使用了不同的缩放策略：维度比较小时，其缩放的更多，维度比较大是，其缩放的更少。\n基于这个策略，作者提出了NTK-aware interpolation，其定义如下：\n$$ g(m) = m, \\quad h(\\theta_{i,d}) = \\theta'^{-2i/d}, i=0,\\dots,(d-2)/2 $$其中\n$$ \\theta' = \\theta\\cdot s^{\\frac{d}{d-2}} $$实际中，这种方法会产生out-of-bound的值，因此最终结果会比PI要差一点，为了解决这个问题，一般会使用比 $s$ 更大的scaling factor.\n上式的推导基于一个简单的假设：我们希望最后一个维度的wavelength在scaling之后，是线性变化的，即 $\\theta'_{(d-2)/2,d}=s\\theta_{(d-2)/2,d}$, 求解之后，我们就得到了上面的定义\nPI 和NTK-aware interpolation的问题在于，我们对不同的维度的处理都是一样的。这类不在乎wavelength的方法被称为blind interpolation methods, 接下来要介绍的就是基于wavelength的方法，即target interpolation methods.\nNTK-by-parts Interpolation 与NTK-aware interpolation， NTK-by-parts interpolation基于wavelength来考虑不同维度上所做的变换。\n对于低维度，其 $\\theta_{i,d}$ 非常大，因此旋转一周所需要的上下文长度也非常大。实际上就导致某些维度的embedding并不是均匀分布的，（比如说只有 $0\\sim\\pi$ 这个区间的embedding），这个时候，模型就只能访问到绝对位置信息，而访问不到相对位置信息。 另外，当我们对所有的维度都进行scale的时候，所有的token都会与彼此更加靠近，这损害了模型对于局部信息的获取能力。 基于这些认知，作者基于wavelength，对不同的维度分别进行处理：\n如果wavelength远小于上下文长度 $L$， 则我们不做任何处理 如果wavelength等于或者大于上下文长度 $L$， 则我们使用NTK-aware interpolation 进行处理 对于中间的其他维度，我们进行了一个trade off 作者定义了一个ratio $r$ 来描述原始上下文长度 $L$ 和 wavelength 之间的关系\n$$ r(i,d) = \\frac{L}{\\lambda_{i,d}} = \\frac{L}{2\\pi\\theta'^{2i/d}} $$基于这个ratio，我们可以定义上面的三种处理方式对应的权重\n$$ \\gamma(r) = \\begin{cases} 0, \u0026\\text{if } r \u003c \\alpha\\\\ 1, \u0026\\text{if } r \u003e \\beta\\\\ \\frac{r-\\alpha}{\\beta-\\alpha}, \u0026\\text{otherwise} \\end{cases} $$其中， $\\beta\u003e\\alpha\u003e0$ 是超参数， $r\u003c\\alpha$, $r\u003c\\beta$ 分别代表了上面的第1种，第2种情况。\n最后，NTK-by-parts interpolation的定义如下\n$$ g(m) = m, \\quad h(\\theta_i) = \\left(1-\\gamma(r(i,d)\\right)\\frac{\\theta_i}{s} + \\gamma(r(i,d))\\theta_i $$作者通过实验发现，在LLaMA上，$\\alpha=1$ 和 $\\beta=32$ 是一个比较好的选择\nDynamic NTK Interpolation 在实际中，一个经常遇到的场景就是sequence length会从1逐步上升到最大上下文长度，比如说inference的时候。对于这种情况，我们有两种解决方法：\n在整个inference周期中，RoPE的scaling factor都设置为 $s=K'/L$, 其中$L'$ 是扩展后的上下文长度 在每次foward的过程汇总，都更新sclaing factor $s=\\max(1, \\ell'/L)$, 这里 $\\ell'$ 是当前sequence的长度 作者发现，方案1在sequence长度小于 $L$的时候性能会下降，并且当上下文长度超过 $L'$ 时，性能下降的更快。 但是，方案2可以让模型的性能下降曲线更平缓。因此，作者将这种inference-time方法称为 Dynamic Scaling method, 当其与NTK-aware方法结合时，就得到了 Dynamic NTK interpolation\n作者通过实验发现，Dynamic NTK interpolation在$L'=L$ 时，效果非常好\nYaRN 在YaRN中，作者针对Dynamic NTK interpolation做了进一步改进，也就是在计算attention softmax时，加入了一个温度参数 $t\u003e0$, 这样attention的计算就变成了\n$$ \\mathrm{Attn}(Q, K, V) = \\mathrm{softmax}\\left(\\frac{QK^T}{t\\sqrt{d}}\\right)V $$作者发现，通过这种scaling的方式，YaRN可以在不改变代码的前提下，更改attention的机制。并且，其不增加训练和推理的cost\n作者将YaRN定义为结合了NTK-by-parts interpolation和上述scaling技巧的方法\n对于LLaMA，作者推荐使用如下参数：\n$$ \\sqrt{\\frac{1}{t}} = 0.1\\ln(s) + 1. $$实验结果如下\n作者发现：\n对于合适的 $t$, 扩展上下文之后，perplexity会变的更好 最好的$t$ 对于不同的位置和样本提升都是一样的 Evaluation 实验结果如下\nExtension Method Trained Tokens Context Window Evaluation Context Window Size 2048 4096 6144 8192 10240 PI (s = 2) 1B 8k 3.92 3.51 3.51 3.34 8.07 NTK ($\\theta$ = 20k) 1B 8k 4.20 3.75 3.74 3.59 6.24 YaRN (s = 2) 400M 8k 3.91 3.50 3.51 3.35 6.04 可以看到，YaRN使用的数据更少，并且当模型扩展到10240的时候，其表现下降的最慢，这说明了YaRN在扩展上下文长度时的有效性\n原始RoPE，dynamic-PI和dynamic-YaRN的对比\n可以看到，RoPE的上下文扩展能力很差，Dynamic-YaRN的表现最好。\nCode YaRN的实现在HuggingFace/src/transformers/modeling_rope_utils.py 里的 _compute_yarn_parameters 函数里，其返回 inv_freq 以及 attention_factor 两个量，前者代表了 $\\theta_d$, 后者代表 $t\\sqrt{d}$.\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 def get_mscale(scale): if scale \u0026lt;= 1: return 1.0 return 0.1 * math.log(scale) + 1.0 def find_correction_dim(num_rotations, dim, base, max_position_embeddings): \u0026#34;\u0026#34;\u0026#34;Inverse dimension formula to find the dimension based on the number of rotations\u0026#34;\u0026#34;\u0026#34; return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings): \u0026#34;\u0026#34;\u0026#34;Find dimension range bounds based on rotations\u0026#34;\u0026#34;\u0026#34; low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings)) high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings)) return max(low, 0), min(high, dim - 1) def linear_ramp_factor(min, max, dim): if min == max: max += 0.001 # Prevent singularity linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) ramp_func = torch.clamp(linear_func, 0, 1) return ramp_func class YaRNRotaryEmbedding(nn.Module): def __init__(self, config): beta_fast = config.get(\u0026#34;beta_fast\u0026#34;) or 32 beta_slow = config.get(\u0026#34;beta_slow\u0026#34;) or 1 dim = config[\u0026#34;head_dim\u0026#34;] factor = config.max_position_embeddings / original_max_position_embeddings pos_freqs = base ** (torch.arange(0, dim, 2) / dim) inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (factor * pos_freqs) low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings) # Get n-dimensional rotational scaling corrected for extrapolation inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2) inv_freq = (inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + inv_freq_extrapolation * inv_freq_extrapolation_factor) attention_factor = get_mscale(factor) def forward(self, x, position_ids): ... return cos, sin 实际的transformers代码中，Qwen使用的还是默认的RoPE，在inference时如果我们需要扩展上下文，可以通过修改config的形式：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 from transformers import pipeline model_name_or_path = \u0026#34;Qwen/Qwen3-8B\u0026#34; generator = pipeline( \u0026#34;text-generation\u0026#34;, model_name_or_path, torch_dtype=\u0026#34;auto\u0026#34;, device_map=\u0026#34;auto\u0026#34;, model_kwargs={ \u0026#34;max_position_embeddings\u0026#34;: 131072, \u0026#34;rope_scaling\u0026#34;: { \u0026#34;rope_type\u0026#34;: \u0026#34;yarn\u0026#34;, \u0026#34;factor\u0026#34;: 4.0, \u0026#34;original_max_position_embeddings\u0026#34;: 32768, }, } ) Conclusion 在本文中，作者首先构建了一个统一的表征不同上下文长度扩展的形式，接下来作者分析了不同上下文长度扩展的不足，并提出了YaRN这种上下文长度扩展方式，结果发现，YaRN不仅在短上下文长度下面表现很好，当上下文长度扩展之后，其表现依然非常优秀。\nReferences arxiv YaRN arxiv PI Qwen documentation ","date":"2025-07-03T14:40:49+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-yarn/","title":"Notes on YaRN"},{"content":"Qwen 在 23 年 9 月份发布了 Qwen 系列大语言模型，包括 1.8B， 7B，14B 三个 size，训练过程使用了 3T token. 作者还基于 Qwen，构建了 Code-Qwen-Chat，Math-Qwen-Chat 等系列领域大语言模型。\nPre-training Data 数据一共使用了 3T token，主要是 public web documents, encyclopedia, books, codes, etc，覆盖了中文和英文两种语言\n数据处理：\n语言识别 去重，包括 MinHash 和 LSH 算法 质量过滤，包括基于规则和和基于 ML 的方法 上采样，特定数据会进行上采样 加入指令数据，提高模型的 zero-shot 和 few-shot 表现 Tokenization BPE tokenizer，最终的 tokenizer 大小为 152K\nArchitecture 模型架构基于 LLaMA， 改动：\ntie embdding: input embdding 和 output embdding 使用的权重相同 position encoding:RoPE, inverse frequency 的精度为 FP32 bias: 取消了大部分的 bias，增加了 QKV bias，来提高模型的外推能力 Pre-Norm \u0026amp; RMSNorm Activation function: SwiGLU Training 上下文长度：2048 attention：flash attention optimizer：AdamW， $\\beta_1=0.9$, $\\beta_2=0.95$, $\\epsilon=10^{-8}$. data type: BF16 Context Extention 使用了三个技巧：\nNTK-aware position interpolation log-N scaling window attention 后续前两个统一成了 YARN.\nobservation: lower layer 对上下文长度扩展更敏感, 因此作者动态调整了 window size\nPost-training 包括 SFT 和 RLHF 两个阶段\nSFT data： 使用了 ChatML 格式\nRLHF PPO 算法\nreward model 构建：基于 Qwen-base model\nRL 训练：先更新 value model 50 steps\n发现：top-p 设置为 0.9 比设置为 1.0 更好\nTool-use and Agent 作者使用了 self-instruct 来进行 SFT，基于 ReAct 构建数据，数据包括 2000 条高质量数据\nSpecialization Code-Qwen code-qwen 基于 qwen continue Pretraining 得到，然后基于 code-qwen 进行 sft 得到 code-qwen-chat，包括 7B 和 14B 两个 size\nMath-Qwen 基于 qwen 直接 SFT 得到，包括 7B 和 14B 两个 size\nConclusion 作者在本文中介绍了 Qwen 系列大语言模型，模型使用了 3T token，作者介绍了训练的细节以及如何扩展到领域大语言模型 Code-Qwen 和 Math-Qwen\nReferences Length exploration Arxiv ","date":"2025-07-03T10:47:27+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen-llm/","title":"Notes on Qwen-LLM"},{"content":"Transformer 实现\n我们采用 top-down 的形式构建 transformer 的代码\n架构 我们以 Qwen3 的代码为例子讲解 Assignment1 的代码实现\n我们通过在 transformer 架构上加上一个 linear layer 就可以完成不同的下游任务，比如：\nQwen3ForQuestionAnswering Qwen3ForCausalLM Qwen3ForSequenceClassification 因此，大语言模型是 transformer 的一个附加产物\nCausalLM 编写大语言模型的第一步为定义 Qwen3ForCausalLM\n1 2 3 4 5 6 7 8 9 10 class CausalLM(nn.Module): def __init__(self, config): self.model = Transformer(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) def forward(self, ...): outputs = self.model(***) logits = self.lm_head(outputs) return logits 这里 lm_head 的作用就是构建 embedding space 到 vocabulary 的映射，即 $\\mathbb{R}^d\\to\\mathbb{R}^{|V|}$\nTransformer transformer 部分包括四个部分：\nEmbedding Layer：将 token 映射到 embedding space layers：Transformer 的主体部分，由 $n$ 个 DecodeLayer 组成 Norm：在输出之前，进行一次 Normalization Position Embedding：由于输入的 sequence 长度是固定的，因此我们提前计算好每一层的 position embedding Transformer 部分的代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 class Transformer(nn.Module): def __init__(self, config): self.embedding = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id) self.layers = nn.ModuleList( [DecodeLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = RotaryEmbedding(config) def forward(self, input_ids,...): input_embeds = self.embedding(input_ids) hidden_states = input_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) for decode_layer in self.layers: layer_outputs = decode_layer(hidden_states, position_ids, position_embeddings) hidden_states = layer_outputs[0] hidden_states = self.norm(hidden_states) return logits DecodeLayer DecodeLayer 就是 transformer 的核心部分，里面包含四个模块：\nPre-Normalization：一般是 RMSNorm 或者 LayerNorm Attention：self-attention Post-Normalization：与 Pre-Normalization 一致 MLP：FFN，SwiGLU 或者 MoE DecodeLayer 还会使用 residual connection 来防止梯度消失\nDecodeLayer 部分的代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class DecodeLayer(nn.Module): def __init__(self, config, layer_idx): self.attn = Attention(config, layer_idx) self.mlp = MLP(config) self.pre_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward(self, hidden_states, position_ids, position_embeddings): residual = hidden_states hidden_states = self.pre_norm(hidden_states) hidden_states = self.attn(hidden_states, position_ids, position_embeddings) # residual hidden_states = hidden_states + residual residual = hidden_states hidden_states = self.post_norm(hidden_states) hidden_states = self.mlp(hidden_states) # residual hidden_states = hidden_states + residual return hidden_states 我们接下来按照\nNormalization MLP Attention Position embedding 的顺序来介绍\nRMSNorm RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单\n$$ \\mathrm{LayerNorm}(x) = \\frac{x-\\mathbb{E}[x]}{\\sqrt{\\mathrm{var}[x]+\\epsilon}}\\odot \\beta + \\gamma $$其中 $\\beta,\\gamma\\in\\mathbb{R}^d$ 是可学习的参数\n$$ \\mathrm{RMSNorm}(x) = \\frac{x}{\\sqrt{\\|x\\|_2^2+\\epsilon}}\\odot \\gamma $$其中 $\\gamma\\in\\mathbb{R}^d$ 是可学习的参数\nRMSNorm 代码实现\n1 2 3 4 5 6 7 8 9 10 11 class RMSNorm(nn.Module): def __init__(self, d, eps): self.weight = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x): input_dtype = x.dtype x = x.to(torch.float32) variance = x.pow(2).mean(-1, keepdim=True) x = x * torch.rsqrt(variance + self.eps) return self.weight * x.to(input_dtype) MLP 现在大语言模型的 MLP 使用的激活函数一般都是 SwiGLU, 其定义为\n$$ \\mathrm{SwiGLU}(x) = x\\odot \\sigma(x) $$其中 $\\sigma(\\cdot)$ 是 sigmoid 函数\nMLP 的定义为\n$$ y = W_2(W_3x\\odot \\mathrm{SwiGLU}(W_1x)) $$其中 $W_3,W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$\n一般地，由于 FFN 只有两个权重矩阵，且 $d_{ff}=4d$, 在 SwiGLU 中，为了保证参数量一致，其隐藏层大小设置为 $d_{ff}'=\\frac23d_{ff}=\\frac83 d$.\nMLP 的代码如下所示\n1 2 3 4 5 6 7 8 9 10 11 def SwiGLU(x): return x * torch.sigmoid(x) class MLP(nn.Module): def __init__(self, d, d_ff): self.gate_proj = nn.Linear(d, d_ff, bias=False) self.up_proj = nn.Linear(d, d_ff, bias=False) self.down_proj = nn.Linear(d_ff, d, bias=False) def forward(self, x): return self.down_proj(SwiGLU(self.gate_proj(x)) * self.up_proj(x)) Attention 我们先不考虑 position embedding，直接看 attention，attention 定义为\n$$ \\mathrm{Attention}(X) = \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d}}\\right)V\\in\\mathbb{R}^{m\\times d} $$其中 $X\\in\\mathbb{R}^{m\\times d}$,\n$$ Q = W_QX\\in\\mathbb{R}^{m\\times d},\\quad K =W_KX\\in\\mathbb{R}^{n\\times d},\\quad V = W_VX\\in\\mathbb{R}^{n\\times d} $$在自回归模型里，我们还会加上 mask, 让每个 token 只能看见前面的 token 的信息\n$$ \\mathrm{Attention}(X) = \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d}}\\odot M\\right)V $$其中\n$$ M = [M_{ij}] = \\begin{cases} 1, \u0026\\text{ if } i \u003c j\\\\ 0, \u0026\\text{ otherwise} \\end{cases} $$self-attention 的代码如下：\n1 2 3 4 5 6 7 8 9 10 11 12 13 def scaled_dot_product_attention(Q, K, V, mask) -\u0026gt; torch.Tensor: d_k = Q.shape[-1] # d_k scaled_factor = 1 / d_k**0.5 scores = torch.einsum(\u0026#34;... s_q d_k, ... s_k d_k -\u0026gt; ... s_q s_k\u0026#34;, Q, K) scores *= scaled_factor if mask is not None: scores = scores.masked_fill(mask == 0, float(\u0026#34;-inf\u0026#34;)) scores = scores.softmax(dim=-1) return torch.einsum(\u0026#34;... s_q s_k, ... s_k d_v -\u0026gt; ... s_q d_v\u0026#34;, scores, V) Multi-Head Attention Multi-Head Attention 定义如下\n$$ \\mathrm{MultiHeadAttention}(X) = [\\mathrm{Attention}_1(X),\\dots,\\mathrm{Attention}_h(X)]W_o\\in\\mathbb{R}^{m\\times d} $$其中 $W_o\\in\\mathbb{R}^{d\\times d}$, 且每一个 Attention heads 的维度会从 $d\\to d/h$.\nMulti-Head Attention 的主要作用为：\n让不同的 head 关注不同的信息 并行计算，提高计算效率 MHA 代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int) -\u0026gt; None: self.q_proj = Linear(d_model, d_model) self.k_proj = Linear(d_model, d_model) self.v_proj = Linear(d_model, d_model) self.output_proj = Linear(d_model, d_model,) def forward(self, x, position_embeddings, mask): Q = rearrange(self.q_proj(x), \u0026#34;... seq_len (num_heads head_dim) -\u0026gt; ... num_heads seq_len head_dim\u0026#34;, num_heads=self.num_heads, head_dim=self.head_dim) K = rearrange( self.k_proj(x), \u0026#34;... seq_len (num_heads head_dim) -\u0026gt; ... num_heads seq_len head_dim\u0026#34;, num_heads=self.num_heads, head_dim=self.head_dim) if mask is None: mask = torch.ones(Q.shape[-2], K.shape[-2]) mask = torch.tril(mask) if position_embeddings is not None: sin, cos = position_embeddings Q = apply_rotary_pos_emb(Q, sin, cos) K = apply_rotary_pos_emb(K, sin, cos) V = rearrange(self.v_proj(x), \u0026#34;... seq_len (num_heads head_dim) -\u0026gt; ... num_heads seq_len head_dim\u0026#34;, num_heads=self.num_heads, head_dim=self.head_dim) output = scaled_dot_product_attention(Q, K, V, mask=mask) output = rearrange(output, \u0026#34;... num_heads seq_len head_dim -\u0026gt; ... seq_len (num_heads head_dim)\u0026#34;) return self.output_proj(output) Position Encoding Attention 对于输入的顺序是不敏感的，也就是\n$$ \\mathrm{Attention}(Q, \\Pi K, \\Pi V) = \\mathrm{Attention}(Q, K, V) $$这里 $\\Pi\\in \\{0,1\\}^{d\\times d}$ 是一个置换矩阵 (permutation matrix)\nTransformer 的解决方法是在 query 和 key 上加上位置信息：\n$$ Q' = Q + PE(Q),\\ K'=K + PE(K) $$这样\n$$ \\mathrm{Attention}(Q, K, V) = \\mathrm{softmax}\\left(\\frac{(Q + PE(Q))(K + PE(K))^T}{\\sqrt{d}}\\right)V\\in\\mathbb{R}^{m\\times d} $$就包含了位置信息\n绝对位置编码 Transformer 的使用的位置编码如下所示\n$$ PE(pos, 2i) = \\sin\\left(\\frac{pos}{10000^{2i/d}}\\right) $$$$ PE(pos, 2i+1) = \\cos\\left(\\frac{pos}{10000^{2i/d}}\\right) $$RoPE 苏剑林老师提出了 Position Encoding，现在已经被广泛使用\n$$ q' = R_{\\theta,m}^dq, k' = R_{\\theta,n}^d k $$这样 $\\langle q, k\\rangle$ 就仅包含两者的相对位置信息\n$$ \\langle q_m, k_n\\rangle = x^TW_qR_{\\theta, n-m}^d W_kx_n $$RoPE 的矩阵定义如下\n$$ R_{\\theta,m}^d = \\mathrm{diag}(M_1,\\dots,M_{d/2}) $$其中\n$$ M_i = \\begin{bmatrix} \\cos m\\theta_i \u0026 -\\sin m\\theta_i\\\\ \\sin m\\theta_i \u0026 \\cos m\\theta_i \\end{bmatrix} $$这里\n$$ \\theta_i = \\frac{1}{10000^{2(i-1)/d}}, i\\in\\{1,2,\\dots,d/2\\} $$简化后得到\n$$ R_{\\theta,m}^dq = \\begin{bmatrix} \\cos m\\theta_0\\\\ \\cos m\\theta_0\\\\ \\vdots\\\\ \\cos m\\theta_{d/2}\\\\ \\cos m\\theta_{d/2}\\\\ \\end{bmatrix}\\odot \\begin{bmatrix} x1\\\\ x2\\\\ \\vdots\\\\ x_{d-1}\\\\ x_d\\\\ \\end{bmatrix} + \\begin{bmatrix} \\sin m\\theta_0\\\\ \\sin m\\theta_0\\\\ \\vdots\\\\ \\sin m\\theta_{d/2}\\\\ \\sin m\\theta_{d/2}\\\\ \\end{bmatrix}\\odot \\begin{bmatrix} -\\ x_2\\\\ x_1\\\\ \\vdots\\\\ -x_d\\\\ x_{d-1}\\\\ \\end{bmatrix} $$RoPE 代码 Naive 实现 RotaryEmbedding 代码\n1 2 3 4 5 6 7 8 9 10 11 12 class RotaryEmbedding(nn.Module): def __init__(self, dim): self.inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) def forward(self, x, position_ids): freqs = einsum(torch.arange(self.max_seq_len), self.inv_freq, \u0026#34;seq_len, d_k_half -\u0026gt; seq_len d_k_half\u0026#34;)[token_positions] sin = torch.sin(freqs) cos = torch.cos(freqs) return sin, cos 计算部分代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -\u0026gt; torch.Tensor: x_even = x[..., ::2] # (seq_len, d_k_half) x_odd = x[..., 1::2] # (seq_len, d_k_half) odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half) evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half) stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half) stacked_trans = rearrange( stacked, \u0026#34;... seq_len double d_k_half -\u0026gt; ... seq_len d_k_half double\u0026#34; ) # (...,seq_len, d_k_half, 2) out = rearrange( stacked_trans, \u0026#34;... seq_len d_k_half double -\u0026gt; ... seq_len (d_k_half double)\u0026#34; ) # (..., seq_len, d_k) return out RoPE 标准实现 RotaryEmbedding 代码 (LLaMA)\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 class LlamaRotaryEmbedding(nn.Module): def __init__(self, config: LlamaConfig, device=None): inv_freq = inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim)) # d_k_half def forward(self, x, position_ids): # (bsz, d_k_half, 1) inv_freq_expanded = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1) # (bsz, 1, seq_len) position_ids_expanded = position_ids[:, None, :] # (bsz, seq_len, d_k_half) freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) # (..., seq_len, d_k) cos = emb.cos() # (..., seq_len, d_k) sin = emb.sin() # (..., seq_len, d_k) return cos, sin 计算部分代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed References Qwen3 transformer source code position encoding blog ","date":"2025-06-29T11:40:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/hands-on-llm2-transformer/","title":"Hands on LLM(2) Transformer"},{"content":"Introduction 在上一篇blog里，我们介绍了MLE和KL minimization的等价性。在这篇blog里，我们将要基于这个等价性，推导masked diffusion LLM, autoregressive LLM, any-order diffusion LLM之间的等价性。最终，我们发现，这几种建模方式本质上都是一致的。\nPreliminary 对于 $\\bm{x}=(x_1,\\dots,x_L)\\in\\mathbb{R}^L$, 基于概率的链式法则，我们有\n$$ p(\\bm{x}) = \\prod_{i=1}^Lp(x_i\\mid x_{","date":"2025-06-28T15:02:09+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/unified-perspective-on-dllm-and-llm/","title":"Unified perspective on dLLM and LLM"},{"content":"MLE 最大似然估计，即MLE (maximum likelihood estimation), 是一个估计参数分布的方法，其核心思想是：模型的参数，应该让观察样本出现的概率最大。\n假设我们有一个参数分布 $p(x\\mid \\theta)$, 其中 $\\theta$ 是参数，如正态分布中的均值和方差。我们从$p(x\\mid \\theta)$进行采样得到 $i.i.d.$ 的数据 $X=\\{x_1,\\dots,x_n\\}$.\n似然函数 (likelihood function) 定义为给定数据 $X$ 的联合分布，即：\n$$ \\mathcal{L}(\\theta\\mid X) = P(X\\mid \\theta) $$由于 $X=\\{x_1,\\dots,x_n\\}$ 是 $i.i.d.$, 因此，我们可以将上式改写为：\n$$ \\mathcal{L}(\\theta\\mid X) = \\prod_{i=1}^n p(x_i\\mid \\theta) $$这样我们的优化目标就是\n$$ \\begin{aligned} \\theta_{MLE}^* \u0026= \\arg\\max_{\\theta} \\mathcal{L}(\\theta\\mid X)\\\\ \u0026= \\arg\\max_{\\theta} \\prod_{i=1}^n p(x_i\\mid \\theta)\\\\ \u0026= \\arg\\max_{\\theta} \\log\\prod_{i=1}^n p(x_i\\mid \\theta)\\\\ \u0026=\\arg\\max_{\\theta} \\sum_{i=1}^n \\log p(x_i\\mid \\theta)\\\\ \\end{aligned} $$即\n$$ \\theta_{MLE}^* = \\arg\\max_{\\theta} \\sum_{i=1}^n \\log p(x_i\\mid \\theta) $$KL divergence KL divergence 用于衡量概率分布 $Q(x)$ 到概率分布 $P(x)$ 的不同程度，我们可以将其理解为：如果我们用 $Q(x)$来替换 $P(x)$, 会有多大的信息损失？\n连续概率分布的KL divergence的定义如下\n$$ D_{KL}(P\\mid\\mid Q) =\\int P(x)\\log\\left(\\frac{P(x)}{Q(x)}\\right)dx $$ 离散概率分布的KL divergence定义如下\n$$ D_{KL}(P\\mid\\mid Q) = \\sum_{x} P(x)\\log\\left(\\frac{P(x)}{Q(x)}\\right) $$KL divergence有两个关键性质：\n非负性：$D_{KL}(P\\mid\\mid Q)\\geq0$, 且 $D_{KL}(P\\mid\\mid Q)=0$ 当且仅当 $P(x)=Q(x)$ 对任意 $x$成立 非对称性： 一般情况下，$D_{KL}(P\\mid\\mid Q)\\neq D_{KL}(Q\\mid\\mid P)$. MLE和KL Divergence的等价性 我们假设 $p_{data}(x)$ 是数据$X$的真实分布， 我们现在需要找到合适的参数 $\\theta$ 以及其对应的分布 $p(x\\mid \\theta)$ 来近似 $p_{data}(x)$, 此时我们可以用KL Divergence作为我们的目标函数，即\n$$ \\theta_{KL} = \\arg\\min_{\\theta}D_{KL}(p_{data}(x)\\mid\\mid p(x\\mid \\theta)) $$我们将上面的式子进行展开得到\n$$ \\begin{aligned} \\theta_{KL}^* \u0026= \\arg\\min_{\\theta}D_{KL}(p_{data}(x)\\mid\\mid p(x\\mid \\theta))\\\\ \u0026= \\arg\\min_{\\theta} \\int p_{data}(x)\\frac{p_{data}(x)}{p(x\\mid \\theta)} dx\\\\ \u0026= \\arg\\min_{\\theta}\\int p_{data}(x)\\log p_{data}(x) dx - \\int p_{data}(x)\\log p(x\\mid \\theta)dx \\\\ \u0026= \\arg\\min_{\\theta} - \\int p_{data}(x)\\log p(x\\mid \\theta)dx \\\\ \u0026= \\arg\\max_{\\theta} \\int p_{data}(x)\\log p(x\\mid \\theta)dx \\end{aligned} $$实际上，真实的数据分布 $p_{data}(x)$ 是未知的，我们只有从 $p_{data}(x)$ 采样得到的一批数据 $X=\\{x_1,\\dots,x_n\\}\\sim p_{data}(x)$.\n基于大数定律，我们有\n$$ \\frac{1}{n}\\sum_{i=1}^n\\log p(\\theta_i\\mid \\theta)=\\mathbb{E}_{x\\sim p_{data}}[\\log p(x\\mid \\theta)] = \\int p_{data}(x)\\log p(x\\mid \\theta)dx, n\\to \\infty $$这样，最大似然估计就与最小化KL divergence构建起了联系：\n$$ \\begin{aligned} \\theta_{MLE}^*\u0026=\\arg\\max_{\\theta} \\sum_{i=1}^n \\log p(x_i\\mid \\theta)\\\\ \u0026= \\arg\\max_{\\theta} \\int p_{data}(x)\\log p(x\\mid \\theta)dx\\\\ \u0026= \\theta_{KL}^*, n\\to\\infty. \\end{aligned} $$也就是说，当采样样本足够多的时候，最大似然估计和最小KL divergence是等价的。\n","date":"2025-06-27T11:35:33+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/relationship-between-mle-and-kl-divergence/","title":"Relationship between MLE and KL divergence"},{"content":"介绍 小米在5月30号发布了MiMo-VL 7B多模态推理大模型，模型包括SFT和RL两个版本，作者在技术报告里讲解了MiMo-VL的架构，预训练和后训练过程，最后对MiMo-VL的性能进行了评测\n架构 MiMo-VL 7B的架构是一个标准的 \u0026ldquo;ViT-MLP-LLM\u0026quot;的架构，其中：\nViT使用了Qwen2.5-VL的Qwen2.5-ViT MLP是一个两层的SwiGLU MLP Layer LLM是之前发布的MiMo-7B 预训练 数据 预训练阶段一共使用了2.4T的token\nImage caption data: 使用了rewrite caption等方法提升和过滤数据，作者还基于MetaCLIP来降低数据的不平衡性 Interleaved data: 主要从webpage, books, papers中提取，最后基于relevance, complementarity, balance of information density来保证数据的质量 OCR and grounding data: OCR数据包括文档，表格，数学公式等。grounding data包括单物体和多物体 video data: rewrite caption, caption有对应的时间戳，加入了video analysis数据 GUI data: 覆盖mobile, web, desktop三种场景, 加入了GUI grounding和GUI action数据 Synthetic reasoning data:使用一个reasoning model来生成带有思考过程的reasoning数据 训练 训练有四个阶段，基本上和之前的模型一致，先训练MLP，然后解冻ViT，再训练全部参数，最后扩展模型的上下文。训练过程如下表所示\n后训练 作者提出了Mixed On-policy Reinforcement Learning (MORL)框架来把RLVR和RLHF结合起来\n数据 RLVR的数据包括:\nvisual reasoning, 80K数据，使用rule-based math-verify library来评估 text reasoning, 使用了MiMo的数学推理数据，评估方式与visual reasoning一致 Image grounding, 包括general and GUI grounding任务，使用GIoU来计算奖励 Visual Counting, 使用准确率来打分 Temporal Video Grounding, temporal video grounding任务，使用IoU来计算奖励 RLHF的数据包括中英文的query，作者使用MiMo-VL-7B和其他模型来采样，然后使用一个advanced VLM来排序. RLHF的reward model使用Breadley-Terry model作为训练目标，text-only reward由MiMo-7B初始化得到，多模态reward model由MiMo-VL-7B初始化得到\n训练 训练过程就是将RLVR和RLHF放在一起进行训练，作者使用了MiMo的seamless rollout engine.\n训练目标与GRPO一致，但是本文中作者使用了on-policy的训练方式。\n评估 作者评估了MiMo-VL-7B的通用能力，reasoning能力，GUI交互理解能力，最后还给出了MiMo-VL-7B在ChatbotArena上的排名\nAblation study 作者首先探究了在预训练阶段加入Long CoT数据对模型表现的影响，结果发现模型正在MMMU, OlympiadBench等benchmark上都有了提升\n作者还比较了原始的GRPO和本文中使用的on-policy版本的GRPO，结果发现on-policy版本的GRPO效果更好\n最后，作者探讨了一下将不同任务放在一起训练导致的互相干扰问题，这个在MiMo里也提到过，核心问题就是有的任务是short CoT的，有的任务是Long CoT的，如何让模型根据任务来选择不同的reasoning length是一个需要探讨的问题，这个问题在VC-PPO里进行了初步的探讨。\n结论 与MiMo一样，MiMo-VL-7B通过在预训练阶段加入reasoning data，来提高模型的reasoning能力，并指出模型可以通过这个方式来媲美其他模型的表现\n参考文献 Github ","date":"2025-06-05T10:51:43+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-mimo-vl/","title":"Notes on MiMo-VL"},{"content":"Introduction 在自然语言处理中，tokenizer的作用是将一个文本序列通过一个字典转化为一个token id的序列。我们回顾图片分类任务，模型在预测的时候，实际上预测的是类别对应的id，而不是类别本身。tokenizer做的事情就是提供一个类似于从类别到对应id的字典。\n一般来说，一个tokenizer处理文本序列的过程有两步：\npre-tokenize，也就是预处理，我们需要将文本序列分割成合适大小的chunks (words) tokenize，构建chunks (words)到token id的映射 实际上, huggingface的tokenizer包括四个步骤, 其中第二第三个步骤与上述一致. 在pre-tokenize之前, 我们有一个normalization过程, 该过程会对文本序列进行处理, 如将文本序列变为小写, 删掉声调符号等, 如下面例子所示：\n1 2 3 4 5 6 from tokenizers import normalizers from tokenizers.normalizers import NFD, StripAccents normalizer = normalizers.Sequence([NFD(), StripAccents()]) normalizer.normalize_str(\u0026#34;Héllò hôw are ü?\u0026#34;) # \u0026#34;Hello how are u?\u0026#34; 在tokenize之后, 我们会有一个post-processing过程, 比如BERT会在生成的token系列前后加入 [CLS] token 和 [SEP] token, 例子如下:\n1 2 3 4 5 6 7 from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(\u0026#34;bert-base-cased\u0026#34;) token_ids = tokenizer.encode(\u0026#34;I love NLP.\u0026#34;) print(token_ids) # [101, 146, 1567, 21239, 2101, 119, 102] # represents [[CLS], \u0026#34;I\u0026#34;, \u0026#34;love\u0026#34;, \u0026#34;NL\u0026#34;, \u0026#34;##P\u0026#34;, \u0026#34;.\u0026#34;, [SEP]] 其完整流程如下图所示 (图源: huggingface llm-course)\n构建好tokenizer之后, 我们还要保证tokenizer提供两个接口：\nencoding, 给定文本序列, 将其映射到字典中去得到token id序列 decoding, 给定token id序列, 将其解码成文本序列 接下来, 我们将简单介绍一下word tokenizer, character tokenizer以及byte tokenizer, 并分析它们各自的不足。 然后, 我们介绍现代大语言模型中使用最多的BPE tokenizer。最后, 我们介绍一些sub-word tokenizer。\nTraining-free tokenizer 本节我们将要介绍word tokenizer, character tokenizer以及byte tokenizer, 它们的特点就是简单易懂, 不需要额外的规则和学习.\nWord tokenizer 给定一个文本序列, 我们现在需要将其转化为一个token序列。一个比较自然的想法是, 我们按照空格将序列拆分成若干个单词, 这样每个单词的语义都能比较好地保留。下面是一个例子：\n1 2 3 4 5 6 import tiktoken tokenizer = tiktoken.get_encoding(\u0026#34;gpt2\u0026#34;) indices = tokenizer.encode(\u0026#34;hello world\u0026#34;) # indices = [31373, 995] # decode = [\u0026#34;hello\u0026#34;, \u0026#34; world\u0026#34;] 接下来我们基于一个预定义好的词典, 将其转化为一个token id的序列。\nword tokenizer的优点是能够保留语义信息，且压缩率比较高（每个token包含的bytes数），其问题是不能处理预定义好的词典之外的词 (out of vocabulary, OOV)。现有的处理方法是使用 \u0026lt;UNK\u0026gt; token来表示这些OOV的词。 但这样显然会丢失语义信息, 因为我们编码成 \u0026lt;UNK\u0026gt; token之后, 就没办法再解码回原有的语义信息了。\nword tokenizer的缺点为：\n单词数量很大, 很多罕见单词的出现频率很低, 降低了tokenizer的利用率 对于不在词典内的单词只能用\u0026lt;UNK\u0026gt; token表示, 损害了语义信息 既然基于word的tokenizer有OOV的问题，我们能否想办法解决这个问题呢？答案是可以的, 我们可以使用 character tokenizer。\nCharacter tokenizer Character tokenizer的基本思想是使用字符而不是单词来编码文本序列。其实现方式如下：\n1 2 3 4 5 6 class CharacterTokenizer: def encode(self, s: str) -\u0026gt; list[int]: return list(ord(c) for c in s) def decode(self, token_ids: list[int]) -\u0026gt; str: return \u0026#34;\u0026#34;.join(chr(token_id) for token_id in token_ids) character tokenizer的词表大小取决于我们的编码方式，UTF-8的编码大概有110K code points。character tokenizer的缺点总结如下：\ncharacter tokenizer会导致我们的词表非常大 和word tokenizer一样, 很多character非常罕见, 会降低词表的利用率 token序列的上下文语义信息较差 Byte tokenizer 我们发现, character tokenizer和word tokenizer的词表都很大, 我们能否想办法降低词表大小, 提升每个token的利用率呢？答案是使用Byte tokenizer.\nByte tokenizer的基本思想是, 所有的字符(character)都是由byte组成的, 比如对于UTF-8编码来说, 每个字符由1-4个byte组成。 因此, 所有满足UTF-8编码的文本, 我们都可以将它们转换为基于byte的token序列。 由于现在的大部分文本都是基于UTF-8的, 因此, 我们只讨论UTF-8编码的文本。\nByte tokenizer的实现如下：\n1 2 3 4 5 6 class ByteTokenizer: def encode(self, s: str) -\u0026gt; list[int]: return list(s.encode(\u0026#34;utf-8\u0026#34;)) def decode(self, token_ids: list[int]) -\u0026gt; str: return bytes(token_ids).decode(\u0026#34;utf-8\u0026#34;) byte tokenizer的词表很小, 其词表大小为 256, 这是因为一个byte可以有256中可能的值.\n尽管byte tokenizer实现简单，并且词表也很小，可以说byte tokenizer解决了character tokenizer和word tokenizer的问题。 但是，byte tokenizer的问题在于，其encode的到的token序列可能会非常长！我们知道，transformer计算量与token序列的长度是平方级关系的，也就是说token序列长度增加10倍，整体的计算量就会增加100倍，因此我们势必需要考虑token序列的长度。\n总之，byte tokenizer的问题为：\n产生的token序列过长, 增加了transformer的计算量 没有上下文语义信息 总结 我们总结一下word tokenizer, character tokenizer以及byte tokenizer三者各自的特点:\nFeature Word Tokenizer Character Tokenizer Byte Tokenizer Granularity Coarse Medium Fine Vocabulary Yes No No Support OOV Bad Good Best #Tokens Small Large Very Large Chinese Yes Yes Yes Support Spell Error Bad Yes Yes Context Good Bad Worst 因此, 这三种tokenizer尽管实现起来很简单, 但是其都有各自的问题. 为了解决这些问题, 我们的做法就是折衷, 使用sub-word tokenizer, 也就是介于word tokenizer和byte tokenizer之间的方法.\nBPE 基本原理与实现 实际生活中，对于出现频率比较高的词，我们会有一个简写的方式，也就是我们使用一个新的单词来表示这个词。比如在英语中，我们会使用plz 来代替 please 以及使用how r u 来代替how are you。 BPE，即byte pair tokenizer的原理也是类似的，对于出现频率比较高的byte pair或者character pair, 我们会使用一个新的token来表示这个pair，这样就压缩了sequence的长度。\nBPE算法包括以下几个步骤:\n对文本序列进行pre-tokenize, 分割成不同的单词 当len(vocab)\u0026lt;vocab_size时, 重复以下步骤: 对所有单词, 统计其相邻character或者byte pair的频率 计算出现频率最高的pair, 使用一个新的token来表示这个pair 将新的token和其对应的token_id加入到vocab中, 并更新单词的分割表示 算法如下图所示 (参考文献2)\n注意：实际上，我们实现的是BBPE (byte BPE算法)，BBPE与BPE的区别在于我们的最小单元是character还是bytes。本质上原理是一致的\n实现代码见 Github naive BPE\n高效实现 BPE的原理很简单, 我们也实现了其naive版本, 但是naive版本的问题是太慢了。因此我们将要优化naive版本的效率。\n首先我们发现, 我们不需要遍历所有的word, 只有含有best_pair的word我们才会进行处理, 因此, 我们的第一个改进就是使用 pair_to_word 来记录每个pair的来源, 比如：\n1 2 3 4 pair_to_word = { (b\u0026#39; \u0026#39;, b\u0026#39;t\u0026#39;): [b\u0026#39; the\u0026#39;, b\u0026#39; it\u0026#39;], (b\u0026#39;t\u0026#39;, b\u0026#39;h\u0026#39;): [b\u0026#39;the\u0026#39;] } 这样, 我们在merge的时候, 直接使用 pair_to_word[best_pair] 来获取需要被更新的token序列就可以了。\n其次, 注意到每次merge之后, 我们都需要重新计算一次 pair_freq, 而实际上, 只有被merge的token序列才需要被重新计数, 其他大部分token序列都是不需要重新计数的。 因此, 一个改进点就是我们在merge的过程中就更新 pair_freq, 而不是重新计算。为了达到这个目标, 我们其实只需要两个操作。 我们用(b'x', b'a', b'b', b'y') 和 best_pair=(b'a', b'b')来说明, merge之前, 这个序列贡献的pair_freq为：\n1 2 3 4 5 { (b\u0026#39;x\u0026#39;, b\u0026#39;a\u0026#39;): 1, (b\u0026#39;a\u0026#39;, b\u0026#39;b\u0026#39;): 1, (b\u0026#39;b\u0026#39;, b\u0026#39;y\u0026#39;): 1, } merge之后, token序列变成了(b'x', b'z', b'y') (假设best_pair对应的新的token为b'z'), 这时候的计数为:\n1 2 3 4 5 6 7 { (b\u0026#39;x\u0026#39;, b\u0026#39;a\u0026#39;): 0, (b\u0026#39;a\u0026#39;, b\u0026#39;b\u0026#39;): 0, (b\u0026#39;b\u0026#39;, b\u0026#39;y\u0026#39;): 0, (b\u0026#39;x\u0026#39;, b\u0026#39;z\u0026#39;): 1, (b\u0026#39;z\u0026#39;, b\u0026#39;y\u0026#39;): 1, } 也就是说, merge之后, 三个pair的计数减少了1, 分别是(token_seq[i-1], merge_pair[0]),merge_pair 和 (merge_pair[1], token_seq[i+2])。两个pair的个数增加了1, 分别是 (token_seq[i-1], new_token)和(new_token, token_seq[i+2]) (这里我们假设merge_pair=(token_seq[i], token_seq[i+1]))。\n基于这个结论，我们就可以优化BPE算法了，具体逻辑就是：\npretokenize, 将 text 切分为若干个 word 计算word_count, pair_freq, pair_to_word, 使用splits记录每个word对应的token分布 重复以下过程： 挑选频率最高的pair将其merge为一个新的token, 基于pair_to_words更新对应的pair_freq 对每个split, 按照上述方式更新pair_freq和split 其具体实现见 Github\nOther subword tokenizers WordPiece WordPiece是Google在预训练BERT时采用的tokenizer，WordPiece的基本思想和BPE差不多，都是从一个较小的vocab开始的。\n首先，WordPiece会通过加上prefix ##来把单词进行切分，比如 word 会被拆分为：\n1 [\u0026#39;w\u0026#39;, \u0026#39;##o\u0026#39;, \u0026#39;##r\u0026#39;, \u0026#39;##d\u0026#39;] 接下来，对于pair $(a, b)$, WordPiece定义了merge pair的规则如下：\n$$ \\mathrm{score}((a, b)) = \\frac{\\#(a, b)}{\\#a \\times \\#b} $$其中 $\\#(a, b)$, $\\#a$, $\\#b$ 分别代表 pair $(a, b)$, 元素 $a$ 和元素 $b$ 的频率。 通过这个方式，我们会给予包含元素出现频率较低的pair更高的优先级。通过这个方式，我们选取score最高的pair，然后将其用一个新的token表示，然后和BPE算法一样，继续这一过程直到我们的vocab size达到指定大小。\n在tokenize的时候，WordPiece会找出现在vocab中的最长的subword, 比如对于'hugs', 假设从左向右在词典中的最长subword是'hug', 那么'hugs' 就会被拆分为 ['hug', '##s']。如果我们在词表中找不到对应的subword，这个时候我们就会使用'[UNK]'来表示。\n具体实现见 Github wordpiece (基于huggingface llm course)。 代码实现除了选择最优pair的方式不同之外，和BPE基本一致。\nUnigram Unigram也是由Google提出来的tokenizer，与BPE和wordpiece不同，unigram从一个非常大的vocab开始，然后merge token来降低vocab的size，直到达到指定大小。初始的vocab可以基于BPE算法或者使用prefix subword来构建。并且，初始vocab还包含所有的base characters来保证所有的word都可以被tokenize。\n算法的描述如下:\n我们来看一下算法的细节, 首先对于一个word, 我们有多种切割方式, 比如'bug'可以被切分为如下三种形式:\n1 [[\u0026#39;b\u0026#39;, \u0026#39;u\u0026#39;, \u0026#39;g\u0026#39;], [\u0026#39;b\u0026#39;, \u0026#39;ug\u0026#39;], [\u0026#39;bu\u0026#39;, \u0026#39;g\u0026#39;]] unigram 假设每个 word 出现的概率是其 subword 出现概率的乘积, 即对于包含 $n$个subword的单词 $\\bm{x}=(x_1,\\dots,x_n)$, 我们有:\n$$ p(\\bm{x}) = \\prod_{i=1}^n p(x_i) $$其中，对于给定的vocab $\\mathcal{V}$, 我们有：\n$$\\sum_{v\\in\\mathcal{V}} p(x)=1$$unigram的目的就是选择合适的切分 $\\bm{x}\\in S(\\bf{x})$ (这里我们用 $\\bf{x}$ 表示单词本身, 用 $\\bm{x}$ 表示 $\\bf{x}$ 的一个切分), 使得 $p(\\bm{x})$的概率最大. 这样我们就可以写出unigram的损失函数了:\n$$ \\mathcal{L} = \\sum_{i=1}^{N} \\log\\left(\\sum_{\\bm{x}\\in S(\\bf{x})}p(\\bm{x})\\right) $$其本质就是: 我们希望对每个单词找到一种合适的切分, 切分得到的subword的概率分布满足其求和为1, 并且使得每个单词的概率最大.\n但是直接对上面概率最大化的问题就是我们每个subword的概率是未知的, unigram的做法是使用EM算法求解这个问题.\n当我们求解完成之后, 对每个subword, 我们都尝试将其从 $\\mathcal{V}$中移除, 然后计算移除后的损失 $loss_i$, 我们依照$loss_i$对subword进行排序, 然后我们去掉 $\\eta \\%$ 比例的subword.\nunigram的伪代码逻辑如下:\n1 2 3 4 5 6 7 8 9 while len(model) \u0026gt; vocab_size: scores = compute_scores(model) sorted_scores = sorted(scores.items(), key=lambda x: x[1]) # Remove percent_to_remove tokens with the lowest scores. for i in range(int(len(model) * percent_to_remove)): _ = token_freqs.pop(sorted_scores[i][0]) total_sum = sum([freq for token, freq in token_freqs.items()]) model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()} 其中 compute_scores 用于计算最优分割以及从model中去掉每个token之后的loss.\n具体实现见 Github wordpiece (基于huggingface llm course)。代码实现的关键在于为每个word选取最优分割，huggingface是采取了动态规划的方法，也就是我们使用 dp[i] 来表示 word[:i] 的最优score，这样我们有：\n$$ dp[i] = \\max_{0 \\leq j \u003c i} dp[j]* p(word[j:i]),\\quad \\mathrm{s.t.}\\ word[j:i]\\in \\mathcal{V} $$这里的乘法代表 $p(\\bm{x}) = \\prod_{i=1}^n p(x_i)$, 在实现的时候我们会取log变成加法，然后概率会由频率来代替。\nSubword tokenizer总结 sub-word tokenizer的对比 (来自huggingface llm course)：\nMethod BPE WordPiece Unigram Start Vocabulary Small Small Large Train Merge tokens Merge tokens Remove tokens Training Step Merge with most frequent pair Merge with best score Remove all tokens minimized the loss Learns Merge rules and a vocab A vocab A vocab with a score for each token Encoding Splits into words and applies merge rules Find the longest subword from the beginning that is in the vocab Finds the most likely split into tokens with learned scores Model GPT BERT T5 实践 tiktoken tiktoken是openAI提出来的一个BPE tokenizer, openAI的模型都基于这个tokenizer, 其主要用于调用GPT系列模型是对token进行计数, 我们可以在tokenizer 这个网站查看其分词情况.\nSentencePiece SentencePiece是google开源的一个无监督的text tokenizer，其实现了BPE和unigram两种算法，SentencePiece还是一个语言无关的tokenizer，使其更适合多语种大语言模型的开发。\nTokenizer tokenizer 是huggingface推出的为基于transformer服务的tokenizer库, 其支持BPE, wordpiece和unigram等分词算法, 使用简便. 并且, huggingface的tokenizer包括两种:\nfast tokenizer, 即Tokenizer库, 这个库是基于Rust开发的 slow tokenizer, 这个是transformer库里模型自带的, 比如ChatGLM就有自己开发的tokenizer huggingface比较了并行处理时两者的区别:\nSetting Fast Tokenizer Slow Tokenizer batched=True 10.8s 4min41s batched=False 59.2s 5min3s huggingface提供的tokenizer库已经非常齐全了, 如果我们要训练新的基于transformer的模型的话，建议直接使用Huggingface的AutoTokenizer。\n总结 特性 SentencePiece Tokenizer tiktoken 是否适合中文 √ √ × 是否适合英文 √ √ √ 是否适合训练 √ √ × 是否快速 √ √ fast 是否用于 GPT 系列 × × √ 是否可解码 √ √ √ 是否支持多语言 √ √ × 结论 本文中, 我们介绍了大语言模型中的tokenizer, 我们从byte level, word level到sub-word level, 再到现代大语言模型最常使用的BPE tokenizer, 并给出了其（高效版本）实现。最后, 我们介绍了一下tokenizer-free的大语言模型和huggingface的tokenizer库。在未来, 我们将继续深入了解大语言模型的基本原理和实现细节。\n参考文献 cs336 Lecture1 Neural Machine Translation of Rare Words with Subword Units SentencePiece Unigram WordPiece Huggingface LLM Course ","date":"2025-05-24T19:56:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/hands-on-llm1-tokenizer/","title":"Hands on LLM(1) Tokenizer"},{"content":"Introduction 我们知道，transformer使用position encoding的一个原因就是，attention layer具有置换不变性，也就是说，我们随机打乱输入token的顺序，并不影响其最终结果 (我们后面会证明，实际上只对key和value具有置换不变性，对query具有置换等变性，也就是改变query的顺序之后，结果的顺序也相应改变)。因此为了让模型学习到正确的上下文知识，我们需要加上position encoding。\n已有的工作大部分都在讨论如何构建更好的position encoding，但是鲜有工作探究为什么attention layer具有置换不变性. 因此，本文将从这一点出发，抽丝剥茧探究其内在原因，最后通过数学公式证明原始transformer是如何具有置换不变性的。\nattention layer介绍 原始transformer layer的架构比较简单，其结构具有attention-LayerNorm-FFN-LayerNorm的形式。给定输入 $X\\in\\mathbb{R}^{d\\times m}$ 和上下文 $Y\\in\\mathbb{R}^{d\\times n}$. 其中，attention的定义为\n$$ \\mathrm{Attn}(X, Y, Y) = V\\mathrm{softmax}\\left(\\frac{K^TQ}{\\sqrt{d}}\\right)\\in\\mathbb{R}^{d\\times m} $$ 其中 $d$是模型的hidden_size, $Q=W_QX\\in\\mathbb{R}^{d\\times m}$, $K=W_KY\\in\\mathbb{R}^{d\\times n}$, $V=W_VY\\in\\mathbb{R}^{d\\times n}$, $W_Q, W_K, W_V\\in\\mathbb{R}^{d\\times d}$ 分别是QKV projection layer的参数.\nLayerNorm的定义为：\n$$ \\mathrm{LayerNorm}(x) = \\frac{x-\\mathbb{E}[x]}{\\sqrt{\\mathrm{var}[x]+\\epsilon}}\\odot \\gamma + \\beta $$ 其中 $\\epsilon\u003e0$是一个超参数， $\\gamma, \\beta\\in\\mathbb{R}^d$ 是可学习的参数.\nFFN的定义为：\n$$ \\mathrm{FFN}(x) = W_2\\max(0, W_1x+b_1)+b_2 $$ 其中 $W_1\\in\\mathbb{R}^{d_{ff}\\times d}$, $W_2\\in\\mathbb{R}^{d\\times d_{ff}}$, $b_1\\in\\mathbb{R}^{d_{ff}}$, $b_2\\in\\mathbb{R}^d$ 是可学习的参数。\n最后，一个attention layer的的结构可以表达为：\n$$ X = X + \\mathrm{LayerNorm}(\\mathrm{Attn}(X, Y, Y))\\\\ X = X + \\mathrm{LayerNorm}(\\mathrm{FFN}(X))\\\\ $$置换不变性的定义 置换不变性(permutation invariant)的定义：假设 $f:\\mathbb{R}^n\\to\\mathbb{R}^n$，如果\n$$ f(\\sigma(\\bm{x})) = (f(\\bm{x})) $$则我们说 $f$是置换不变的. 这里 $\\sigma:\\mathbb{R}^n\\to\\mathbb{R}^n$ 是一个置换函数 (permutation function). 当输入的是一个矩阵时，我们默认置换其列，即对 $X=[X_1,\\dots,X_n]\\in\\mathbb{R}^{d\\times n}$, 我们有 $\\sigma(X)=[X_{\\sigma_1},\\dots, X_{\\sigma_n}]=Y\\Pi $, 其中 $\\Pi\\in\\mathbb{R}^{n\\times n}\\in \\{0,1\\}^{n\\times n}$ 是一个置换矩阵 (permutation matrix)。\n置换等变性 (permutation equivariant)的定义：假设 $f:\\mathbb{R}^n\\to\\mathbb{R}^n$，如果\n$$ f(\\sigma(\\bm{x})) = \\sigma(f(\\bm{x})) $$ 则我们说 $f$是置换等变的.\nattention的置换不变性与置换等变性 我们首先证明attention 对于key和value是置换不变的，即\n$$ \\boxed{\\mathrm{Attn}(X, \\sigma(Y),\\sigma(Y)) = \\mathrm{Attn}(X, Y, Y)} $$证明: 我们直接计算即可得到：\n$$ \\begin{aligned} \\mathrm{Attn}(X, \\sigma(Y),\\sigma(Y)) \u0026= V\\Pi\\mathrm{softmax}\\left(\\frac{(K\\Pi)^TQ}{\\sqrt{d}}\\right)\\\\ \u0026=V\\Pi\\mathrm{softmax}\\left(\\frac{\\Pi^TK^TQ}{\\sqrt{d}}\\right)\\\\ \\end{aligned} $$ 由于softmax是按列计算的，置换只是改变了元素的顺序，因此我们自然有\n$$ \\mathrm{Attn}(X, \\sigma(Y),\\sigma(Y)) = V\\Pi\\Pi^T\\mathrm{softmax}\\left(\\frac{K^TQ}{\\sqrt{d}}\\right)=V\\mathrm{softmax}\\left(\\frac{K^TQ}{\\sqrt{d}}\\right)=\\mathrm{Attn}(X, Y, Y) $$ 这里我们使用了性质 $\\Pi\\Pi^T=\\mathbf{I}$.\n接下来我们证明，attention对于query是置换等变的，即\n$$ \\boxed{\\mathrm{Attn}(\\sigma(X), Y, Y) = \\sigma(\\mathrm{Attn}(X,Y,Y))} $$证明: $$ \\begin{aligned} \\mathrm{Attn}(\\sigma(X), Y, Y) \u0026= V\\mathrm{softmax}\\left(\\frac{K^TQ\\Pi}{\\sqrt{d}}\\right)\\\\ \u0026= V\\mathrm{softmax}\\left(\\frac{K^TQ\\Pi}{\\sqrt{d}}\\right)\\\\ \u0026= V\\mathrm{softmax}\\left(\\frac{K^TQ}{\\sqrt{d}}\\right)\\Pi\\\\ \u0026= \\mathrm{Attn}(X,Y,Y)\\Pi\\\\ \u0026= \\sigma(\\mathrm{Attn}(X,Y,Y)) \\end{aligned} $$从以上的证明可以看到，attention layer对于key和value具有置换不变性，也就是说，我们改变文字顺序不影响最终的输出结果。 但是，我们发现，尽管我们证明了attention具有置换不变性，我们却忽略了一件事：那就是我们计算query, key和value的时候，没有加上bias! 为什么bias如此重要呢？这是因为，$W\\sigma(x) = \\sigma(Wx)$, 但是 $W\\sigma(X)\\neq \\sigma(Wx+b)$. 因此，我们就会思考，难道是transformer实际上可以通过增加bias的方式来让模型学习到上下文知识？事实上并非如此，我们将要通过分析表明，我们计算query, key和value时，增加的query bias和key bias会被softmax操作给消除掉，而key bias则会被LayerNorm消除掉。因此，我们加与加bias，对attention的置换不变性没有任何影响。\nBias对attention layer的影响 接下来，我们考虑在计算query, key和value时加入bias。为了简化，我们只考虑query为一个向量的情况，即 $X=\\bm{x}\\in\\mathbb{R}^d$, 我们计算query, key和value如下：\n$$ \\bm{q} = W_Q\\bm{x}+\\bm{b}_Q\\in\\mathbb{R}^{d}\\\\ K = W_KY + \\bm{b}_K\\mathbf{1}^T\\in\\mathbb{R}^{d\\times n}\\\\ V = W_VY + \\bm{b}_V\\mathbf{1}^T\\in\\mathbb{R}^{d\\times n} $$这里 $\\mathbf{1}^T\\in\\mathbb{R}^{n}$. 我们这里简化了scaling的操作，因为其不对结果产生影响。\n注：以下证明参考了【参考文献2】\n我们首先展开attention中的 $V$:\n$$ \\begin{aligned} \\mathrm{Attn}(\\bm{x}, Y, Y) \u0026= V\\mathrm{softmax}\\left(K^T\\bm{q}\\right)\\\\ \u0026= \\left(W_VY + \\bm{b}_V\\mathbb{1}^T\\right)\\mathrm{softmax}\\left(K^T\\bm{q}\\right)\\\\ \u0026= W_VY\\mathrm{softmax}\\left(K^T\\bm{q}\\right) + \\bm{b}_V\\mathbb{1}^T \\mathrm{softmax}\\left(K^T\\bm{q}\\right) \\end{aligned} $$ 由于 $\\mathrm{softmax}\\left(K^T\\bm{q}\\right)\\in\\mathbb{R}^{n}$的列求和为$1$, 因此，$\\mathbb{1}^T\\mathrm{softmax}\\left(K^T\\bm{q}\\right)=1$, 我们有\n$$ \\mathrm{Attn}(\\bm{x}, Y, Y) = W_VY\\mathrm{softmax}\\left(K^T\\bm{q}\\right) + \\bm{b}_V $$接下来，我们展开 $K$:\n$$ \\begin{aligned} \\mathrm{Attn}(\\bm{x}, Y, Y) \u0026= W_VY\\mathrm{softmax}\\left(K^T\\bm{q}\\right) + \\bm{b}_V\\\\ \u0026= W_VY\\mathrm{softmax}\\left((W_KY + \\bm{b}_K\\mathbf{1}^T)^T\\bm{q}\\right) + \\bm{b}_V\\\\ \u0026= W_VY\\mathrm{softmax}\\left(Y^TW_K^Tq + \\mathbf{1}\\bm{b}_K^T\\bm{q}\\right) + \\bm{b}_V\\\\ \\end{aligned} $$这里，我们需要用到softmax函数的平移不变性，即 $\\mathrm{softmax}(\\bm{x}+\\delta\\mathbf{1})=\\mathrm{softmax}(\\bm{x})$, 这里 $\\delta\\in\\mathbb{R}$ 是一个常数，证明起来很简单： $$ \\mathrm{softmax}(\\bm{x}+\\delta)_i = \\frac{e^{x_i+\\delta}}{\\sum_{j}e^{x_j+\\delta}} = \\frac{e^{x_i} * e^{\\delta}}{\\sum_{j}e^{x_j} * e^{\\delta}} = \\mathrm{softmax}(\\bm{x})_i $$ 而这里 $\\bm{b}_K^T\\bm{q}\\in\\mathbb{R}$，因此我们可以将这一项给去掉，我们得到：\n$$ \\mathrm{Attn}(\\bm{x}, Y, Y) = W_VY\\mathrm{softmax}\\left(Y^TW_K^T\\bm{q}\\right) + \\bm{b}_V $$接下来，我们展开 $\\bm{q}$, $$ \\boxed{ \\begin{aligned} \\mathrm{Attn}(\\bm{x}, Y, Y) \u0026= W_VY\\mathrm{softmax}\\left(Y^TW_K^T\\bm{q}\\right) + \\bm{b}_V\\\\ \u0026= W_VY\\mathrm{softmax}\\left(Y^TW_K^T(W_Q\\bm{x}+\\bm{b}_Q)\\right) + \\bm{b}_V\\\\ \\end{aligned}} $$因此，我们最终的结论为： key bias对attention输出没有任何贡献，query bias和key bias会影响结果。\n到这里，看了参考文献3，我本以为可以进一步简化。但实际上并不行。参考文献3关于“transformer block is equivariant\u0026quot;的结果是错的，因为在attention layer之后还有一个LayerNorm，而LayerNorm不是置换不变的，这也是LayerNorm和BatchNorm之间的区别。也就是如果我们在nn.Linear后加一个BatchNorm，那么nn.Linear的bias是无效的，反之如果是LayerNorm的话，则bias是有效的.\n为什么没有bias 实际上这个问题并没有定论。特别是加入position encoding之后，就更难探究bias对最终结果的影响了。但是，我认为一个原因就是bias其实就是某种先验知识，假设输入满足高斯分布，那么我们有\n$$ \\mathbb{E}[W\\bm{x}+b] = b $$加上先验知识后，当训练数据出现distribution shift之后，模型在训练过程中可能就会不稳定(PaLM). 而后来将LayerNorm替换为RMSNorm，使用RoPE而不是其他的additive position encoding, 我认为也是避免模型学习到先验知识，从而影响其泛化性。在未来，我认为transformer里应该是没有bias的，尽管这样效果可能会差一些，但是其稳定性更好，泛化性应该也会更好。\n结论 在本文中，我们分析了attention的性质，我们发现，在原始transformer架构中，attention对于key和value有置换不变性，对于query有置换等变性。然后，我们给出了一些猜测，也就是bias会让模型产生先验知识，而这种先验知识很可能会影响训练的稳定性和模型的泛化性。\n参考文献 Attention is All you Need Role of Bias Terms in Dot-Product Attention Are Transformers universal approximators of sequence-to-sequence functions? 附录 下面是测试上面结论的python代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 import torch import torch.nn as nn import torch.nn.functional as F # 设置随机种子，确保可复现性 torch.manual_seed(42) # 输入参数 batch_size = 1 seq_len = 16 embed_dim = 1024 # 嵌入维度 num_heads = 32 # 多头注意力头数 head_dim = embed_dim // num_heads # 输入张量 (batch_size, seq_len, embed_dim) x = torch.randn(batch_size, seq_len, embed_dim) # 有 bias 的 QKV 线性层 class Attention(nn.Module): def __init__(self, embed_dim, q_bias=False, k_bias=False, v_bias=False): super().__init__() self.q = nn.Linear(embed_dim, embed_dim, bias=q_bias) self.k = nn.Linear(embed_dim, embed_dim, bias=k_bias) self.v = nn.Linear(embed_dim, embed_dim, bias=v_bias) def forward(self, x): B, N, C = x.shape q = self.q(x).reshape(B, N, num_heads, head_dim).transpose(1, 2) k = self.k(x).reshape(B, N, num_heads, head_dim).transpose(1, 2) v = self.v(x).reshape(B, N, num_heads, head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) * (1.0 / (head_dim**0.5)) attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1, 2).reshape(B, N, C) return x, attn # 初始化模型 model_no_bias = Attention(embed_dim, q_bias=False, k_bias=False, v_bias=False) model_with_bias = Attention(embed_dim, q_bias=False, k_bias=True, v_bias=False) model_with_bias.q.weight.data = model_no_bias.q.weight.data model_with_bias.k.weight.data = model_no_bias.k.weight.data model_with_bias.v.weight.data = model_no_bias.v.weight.data # 推理 out_no_bias, attn_no_bias = model_no_bias(x) out_with_bias, attn_with_bias = model_with_bias(x) # 比较差异 diff_output = torch.abs(out_no_bias - out_with_bias).mean() diff_variance = torch.abs(out_no_bias - out_with_bias).var() diff_attn = torch.abs(attn_no_bias - attn_with_bias).mean() print(\u0026#34;\\nMean difference in output:\u0026#34;, diff_output.item()) print(\u0026#34;Mean difference in variance:\u0026#34;, diff_variance.item()) print(\u0026#34;Mean difference in attention weights:\u0026#34;, diff_attn.item()) # Mean difference in output: 1.2734082233123445e-08 # Mean difference in variance: 1.7173628739783402e-16 # Mean difference in attention weights: 3.949708116124384e-09 ","date":"2025-05-22T15:25:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-attention-bias/","title":"Notes on attention bias"},{"content":" 本文前半部分参考 参考文献1，推荐大家看博客原文。\nPosition encoding总结 在 上一篇blog 中, 我们介绍了 Attention 的两个性质，也就是在不加 position encoding 的情况下，Attention 对于 query 是 permutation equivariant 的，对于 key 和 value 是 permutation invariant 的。\n但是“我爱你”和“你爱我”这两句话所表示的含义应该是不一样的，我们将这两句话作为 key 和 value 的时候，我们发现模型的输出是一致的，这显然是不能接受的。因此，我们就需要加入 position encoding，让模型学习到语序信息，从而明白不同的语序有不同的含义。\n下面是测试代码 （来自 参考文献1）\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 import torch import torch.nn as nn from transformers import AutoTokenizer, AutoModel model_id = \u0026#34;meta-llama/Llama-3.2-1B\u0026#34; tok = AutoTokenizer.from_pretrained(model_id) model = AutoModel.from_pretrained(model_id) text = \u0026#34;The dog chased another dog\u0026#34; tokens = tok(text, return_tensors=\u0026#34;pt\u0026#34;)[\u0026#34;input_ids\u0026#34;] embeddings = model.embed_tokens(tokens) hdim = embeddings.shape[-1] W_q = nn.Linear(hdim, hdim, bias=False) W_k = nn.Linear(hdim, hdim, bias=False) W_v = nn.Linear(hdim, hdim, bias=False) mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True) with torch.no_grad(): for param in mha.parameters(): nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings)) dog1_out = output[0, 2] dog2_out = output[0, 5] print(f\u0026#34;Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}\u0026#34;) #True Position encoding 可以分为绝对位置编码 (absolute position encoding, APE)，相对位置编码 (relative position encoding, RPE) 以及可学习的位置编码。可学习位置编码主要是 BERT 类的模型在使用，其训练成本比较高，本文不做讨论。绝对位置编码是原始 transformer 里提出的编码模式，现在的大多数基于 transformer 模型使用的都是相对位置编码。\n本文中，我们先介绍位置编码应该具有的性质，然后我们分别介绍绝对位置编码和相对位置编码，我们将着重关注苏剑林老师提出来的 RoPE。\n位置编码 在介绍位置编码之前，我们首先应该关注位置编码的性质，位置编码的目标是为输入的 token embedding 增加位置信息，那么理想的位置编码应该是怎么样的呢？\n我们这里直接引用 参考文献1 中给定的性质：\n性质 1: token sequence 中每个位置的位置编码都是唯一的。这个很好理解，如果不唯一的话，那么根据前面推导的性质，这两个位置的 attention 输出就完全一致了 性质 2: 线性相关性。也就是说，如果我们知道了位置 $p$ 处的位置编码，那么理想情况下，我们应该能比较简单地得到 $p+k$ 处的位置编码，理想情况下，我们应该有 $PE(p+k)=W_kPE(p)$. 性质 3: 泛化到长上下文中去。我们希望位置编码不仅在 8K 的上下文起作用，还希望位置编码能够泛化到 32K 的上下文 性质 4: 生成模式是固定的。固定的模式有助于模型更好地学习位置相关的信息 性质 5: 可以扩展到多维。我们希望位置编码可以从文本扩展到图片再到视频，也就是从 $1D$ 到 $nD$. 绝对位置编码 绝对位置编码依照其名称，其思想就是为每个位置的 token 分配一个固定的位置信息，也就是对于输入的 hidden states $\\bm{x}=[\\bm{x}_1,\\dots,\\ \\bm{x}_m]\\in\\mathbb{R}^{m\\times d}$, 我们有\n$$ \\bm{x}_i' = \\bm{x}_i + p_i, i=1,\\dots, m $$这里，$p_i\\in\\mathbb{R}^d$. 我们的 attention 就变成了\n$$ \\mathrm{Attn}(X) = \\mathrm{softmax}\\left(\\frac{(Q+P)(K+P)^T}{\\sqrt{d}}\\right)V\\in\\mathbb{R}^{m\\times d} $$这里\n$$ P = [p_1,\\dots,p_m]\\in\\mathbb{R}^{m\\times d}， Q= W_QX\\in\\mathbb{R}^{m\\times d}, K=W_KX, V=W_VX\\in\\mathbb{R}^{n\\times d} $$整数位置编码 一个最简单的想法就是我们使用正整数来标记 token 所在的位置，也就是\n$$ PE(i) = [i, \\dots, i]=i\\mathbf{1}_{d\\times 1}\\in\\mathbb{R}^d,\\ i=1,\\dots,m $$可以看到，这个简单的设计满足性质 1，性质 2，性质 3，性质 4.\n但是，注意到 attention 的输入 $X$ 通常是经过 Layer Normalization 处理过后的，因此其按列符合正态分布，并且均值和方差一般较小。当我们加上整数位置编码之后，其 token 本身的信息就会被污染，也就是信噪比非常低。一个解决方法就是我们对 $PE(i)$ 进行 normalization，即\n$$ PE(i)' = \\frac{1}{m}PE(i) = \\frac{i}{m}\\mathbf{1}_{d\\times 1} $$现在所有的位置编码的值都比较小，但是我们发现新的位置编码不满足性质 2 了，这是因为现在位置编码还和 sequence 长度有关，我们从位置 $p$ 到位置 $p+k$ 不仅取决于 $k$ 还取决于 sequence 长度 $m$\n二进制位置编码 既然整数位置编码的主要问题是对输入影响太大，我们能否找一个不影响输入的整数位置编码方式呢？ 参考文献1 提出了二进制位置编码，因为每个 token 是 $d$ 维的，因此我们可以使用 $d$ 位二进制来表示 $i$. 比如说，当 $d=3$, $m=4$ 时，我们的位置编码分别为\n$$ PE(0) =p_{(000)_2} = [0, 0, 0],\\ PE(1) =p_{(001)_2}= [0, 0, 1],\\ PE(2) =p_{(010)_2} = [0, 1, 0],\\ PE(3) =p_{(011)_2} = [0, 1, 1] $$现在，我们二进制位置编码满足性质 1，性质 2. 对于性质 3，由于 $d$ 位二进制的表示范围为 $[0, 2^d-1]$，因此其泛化性受到 $d$ 的影响。\n参考文献1 画出了不同位置的值的变化情况。我们这里也模仿绘制出类似的曲线图\n我们发现，二进制位置编码高位，也就是 $PE(i)_{d}$ 的变化很慢，而低位，也就是 $PE(i)_{0}$ 变化很快，\n二进制位置编码解决了整数位置编码的信噪比过低和线性相关性。但是其问题是其对不同位置的 token embedding 产生的影响是不一样的。比如位置 1 和位置 2 的相同的 token embedding 之间的区别是：\n$$ (\\bm{x}_2 + PE(2)) - (\\bm{x}_1 + PE(1)) = (\\bm{x}_2-\\bm{x}_1)+ [0, 1, -1] $$一般来说, $\\bm{x}_2-\\bm{x}_1$ 比较小，因此使用二进制位置编码的问题是输入位置的微小变化（增加一个 token 或减少一个 token）都会对最终结果产生巨大影响。因此，我们需要想办法解决这个问题。\nSinusoidal 前面提到二进制位置编码的问题是相邻 token 之间变化太大，不够光滑。因此我们想要增加一个光滑性质，也就是说我们希望：\n位置编码值在 $[-1, 1]$ 之间，防止对 token embedding 产生影响 相邻 token 的位置编码尽可能相近，即 $|PE(k+p)-PE(p)| \\leq \\delta |k|$, 其中 $\\delta\u003e0$ 是一个比较小的数。 与二进制一样，高位的变化比较慢，低位的变化比较快 一个想法就是利用三角函数 $\\sin$ 或者 $\\cos$，三角函数满足前两个性质， 对于第三个性质，我们可以通过控制频率来满足。这样我们得到的位置编码就具有如下形式：\n$$ PE(p, i) = \\sin\\left(\\frac{p}{\\theta^{i/d}}\\right) $$其中 $\\theta$ 是我们的超参数。\n我们现在来推导一下上面位置编码的线性相关性：\n$$ PE(p+k) = \\sin\\left(\\frac{p+k}{\\theta^{i/d}}\\right)=PE(p)\\cos\\left(\\frac{k}{\\theta^{i/d}}\\right) + \\cos\\left(\\frac{p}{\\theta^{i/d}}\\right)\\sin\\left(\\frac{k}{\\theta^{i/d}}\\right) $$我们发现，$\\sin$ 位置编码不满足线性相关性。但是出现的 $\\cos$ 给了我们启发，也就是我们可以同时使用 $\\sin$ 和 $\\cos$ 来完成位置编码，这也是原始 transformer 里提出来的 Sinusoidal 位置编码，其形式为：\n$$ \\begin{aligned} PE(p, 2i) \u0026= \\sin\\left(\\frac{p}{\\theta^{2i/d}}\\right)\\\\ PE(p, 2i+1) \u0026= \\cos\\left(\\frac{p}{\\theta^{2i/d}}\\right) \\end{aligned} $$现在，记 $\\omega_i=1/\\theta^{2i/d}$, 我们再推导一下线性相关性，就得到：\n$$ \\begin{aligned} \\begin{bmatrix} PE(p+k, 2i)\\\\ PE(p+k, 2i+1)\\\\ \\end{bmatrix}\u0026=\\begin{bmatrix} \\sin \\omega_i(p+k)\\\\ \\cos \\omega_i(p+k) \\end{bmatrix}\\\\ \u0026=\\begin{bmatrix} \\sin \\omega_i(\\omega_ip)\\cos(\\omega_ik)+\\cos w_i(\\omega_ip)\\sin(\\omega_ik)\\\\ \\cos \\omega_i(\\omega_ip)\\cos(\\omega_ik)-\\sin w_i(\\omega_ip)\\sin(\\omega_ik) \\end{bmatrix}\\\\ \u0026= \\begin{bmatrix} \\cos(\\omega_ik)\u0026 \\sin(\\omega_ik)\\\\ -\\sin(\\omega_ik)\u0026 \\cos(\\omega_ik) \\end{bmatrix}\\begin{bmatrix} PE(p, 2i)\\\\ PE(p, 2i+1)\\\\ \\end{bmatrix} \\end{aligned} $$也就是说，Sinusoidal 位置编码满足线性相关性。对于 Sinusoidal 位置编码我们也可以进行可视化：\n相对位置编码 前面介绍了绝对位置编码，每个位置的位置编码是固定的。但是绝对位置编码的问题是，模型比较难以学习相对位置关系。\n举个例子，我们提到上下文时，通常会使用“上一节”，“上一章”这些表示相对位置关系的词。\n因此，我们希望让模型学习相对位置关系而不是绝对位置关系，因为相对关系更符合我们的认知。\nRoPE RoPE 由苏剑林老师提出，最早应用于 LLaMA 架构（没有确认），后续被大多数模型所采用。\n之前的 PE 大多数关注于加性位置编码，也就是假设位置编码的形式为 $\\bm{x}+\\bm{p}$, 基于这种假设，已有的工作基本都集中于优化下面的 Q 和 K 的内积\n$$ \\langle f_q(\\bm{x}_q, m), f_k(\\bm{x}_k, n) \\rangle $$这里 $f_q(\\bm{x}_q, m)=W_q(\\bm{x}_q+\\bm{p}_m)$, $f_k(\\bm{x}_k, n)=W_k(\\bm{x}_k+ \\bm{p}_n)$.\n而 RoPE 里面，作者使用了一个不同的假设： 假设内积应该仅包含两者的相对信息，也就是\n$$ \\langle f_q(\\bm{x}_q, m), f_q(\\bm{x}_k, n)\\rangle := g(\\bm{x}_q,\\bm{x}_k, m-n) $$这里的 $f$ 和 $g$ 都是未知函数。我们的目标就是从这个公式中推导出一个合适的位置编码出来。\n不失一般性，我们可以假设\n$$ f_q(\\bm{x}_m,0) = \\bm{x}_q,\\quad f_q(\\bm{x}_n, 0) = \\bm{x}_k $$这个假设代表初始条件下，我们不对输入做任何改变，也就是不增加位置信息。\n2D 推导 与 RoPE 一样，我们直接使用复平面来进行推导。\n我们假设 $d=2$, 注意到二维平面上的每个点都可以表示为如下形式\n$$ \\bm{z} = (x,y) = re^{i\\theta} $$其中 ($\\mathrm{atan2}$ 定义参考 维基百科)\n$$ r = \\|\\bm{z}\\|_2 = \\sqrt{x^2+y^2}\\in\\mathbb{R},\\quad \\theta = \\mathrm{atan2}(y, x)\\in\\mathbb{R}, $$现在，对于三个向量 $f_q(\\bm{x}_q, m)$, $f_q(\\bm{x}_k, n)$, $g(\\bm{x}_q,\\bm{x}_k, m-n)$ 我们可以写出其极坐标形式：\n$$ \\begin{aligned} f_q(\\bm{x}_q,m) \u0026:= r_q(\\bm{x}_q,m)e^{i\\theta_q(\\bm{x}_q,m)}\\\\ f_k(\\bm{x}_k, n) \u0026:= r_k(\\bm{x}_k, n)e^{i\\theta_k(\\bm{x}_k, n)}\\\\ g(\\bm{x}_q,\\bm{x}_k, m-n) \u0026:= r_g(\\bm{x}_q,\\bm{x}_k, m-n)e^{i\\theta_g(\\bm{x}_q,\\bm{x}_k, m-n)} \\end{aligned} $$我们计算内积并比较同类项得到：\n$$ \\begin{aligned} r_g(\\bm{x}_q,\\bm{x}_k, m-n) \u0026:= r_q(\\bm{x}_q,m)r_k(\\bm{x}_k, n)\\\\ \\theta_g(\\bm{x}_q,\\bm{x}_k, m-n) \u0026:= \\theta_q(\\bm{x}_q,m)-\\theta_k(\\bm{x}_k, n) \\end{aligned}\\tag{3} $$我们接下来分别推导 $r_g(\\bm{x}_q,\\bm{x}_k, m-n)$ 和 $\\theta_g(\\bm{x}_q,\\bm{x}_k, m-n)$ 的形式\n$r_g(\\bm{x}_q,\\bm{x}_k, m-n)$ 我们令 $m=n=0$ 可以得到初始条件\n$$ r_g(\\bm{x}_q,\\bm{x}_k, 0) = r_q(\\bm{x}_q,0)r_k(\\bm{x}_k, 0)=\\|\\bm{q}\\|_2\\|\\bm{k}\\|_2 $$我们再令 $n=0$,得到\n$$ r_g(\\bm{x}_q,\\bm{x}_k, m) = r_q(\\bm{x}_q,m)r_k(\\bm{x}_k, 0)=r_q(\\bm{x}_q,m)\\|\\bm{k}\\|_2=\\frac{r_g(\\bm{x}_q,\\bm{x}_k, m-n)}{r_k(\\bm{x}_k, n)}\\|\\bm{k}\\|_2 $$这里最后一个等式带入了原始等式 (3)，注意到左侧与 $n$ 无关，因此右侧我们选取 $n=1$, 得到\n$$ r_g(\\bm{x}_q,\\bm{x}_k, m) = \\frac{r_g(\\bm{x}_q,\\bm{x}_k, m-1)}{r_k(\\bm{x}_k, 1)}\\|\\bm{k}\\|_2 =\\cdots= r_g(\\bm{x}_q,\\bm{x}_k, 0)\\left(\\frac{\\|\\bm{k}\\|_2 }{r_k(\\bm{x}_k, 1)}\\right)^{m+1} $$令 $m=0$ 我们有\n$$ r_k(\\bm{x}_k, 1) = \\|\\bm{k}\\|_2. $$因此我们最终的表达式为：\n$$ r_g(\\bm{x}_q,\\bm{x}_k, m) = r_g(\\bm{x}_q,\\bm{x}_k, 0) = \\|\\bm{q}\\|_2\\|\\bm{k}\\|_2. $$并且，通过分别设置 $m=0$ 以及 $n=0$ 我们还可以得到\n$$ r_q(\\bm{x}_q,m) = \\|\\bm{q}\\|_2,\\quad r_k(\\bm{x}_k, n) = \\|\\bm{k}\\|_2 $$$\\theta_g(\\bm{x}_q,\\bm{x}_k, m-n)$ 令 $m=n=0$, 我们得到初始条件\n$$ \\theta_g(\\bm{x}_q,\\bm{x}_k, 0) = \\theta_q(\\bm{x}_q,0)-\\theta_k(\\bm{x}_k, 0)=\\theta_q-\\theta_k $$令 $n=1$, 我们有\n$$ \\begin{aligned} \\theta_g(\\bm{x}_q,\\bm{x}_k, m-1) \u0026= \\theta_q(\\bm{x}_q,m)-\\theta_k(\\bm{x}_k, 1)\\\\ \u0026=\\theta_g(\\bm{x}_q,\\bm{x}_k, m-n) + \\theta_k(\\bm{x}_k, n)-\\theta_k(\\bm{x}_k, 1) \\end{aligned} $$这里我们带入了公式 (3)，注意到公式左边与 $n$ 无关，因此在公式右侧我们令 $n=0$, 得到\n$$ \\theta_g(\\bm{x}_q,\\bm{x}_k, m-1) = \\theta_g(\\bm{x}_q,\\bm{x}_k, m)+ \\theta_k(\\bm{x}_k, 0)-\\theta_k(\\bm{x}_k, 1) $$分别令 $m=1,2,\\dots$ 并相加这些等式，我们得到\n$$ \\theta_g(\\bm{x}_q,\\bm{x}_k, 0) = \\theta_g(\\bm{x}_q,\\bm{x}_k, m) + m(\\theta_k(\\bm{x}_k, 0)-\\theta_k(\\bm{x}_k, 1)) $$即\n$$ \\theta_g(\\bm{x}_q,\\bm{x}_k, m) = m(\\theta_k(\\bm{x}_k, 1)-\\theta_k(\\bm{x}_k, 0))+(\\theta_q-\\theta_k)\\tag{4} $$注意到\n$$ \\theta_g(\\bm{x}_q,\\bm{x}_k, m) = \\theta_q(\\bm{x}_q,m)-\\theta_k(\\bm{x}_k, 0)=\\theta_q(\\bm{x}_q,m)-\\theta_k $$带入上式我们就得到\n$$ \\theta_q(\\bm{x}_q,m) = m(\\theta_k(\\bm{x}_k, 1)-\\theta_k(\\bm{x}_k, 0))+\\theta_q $$在 (4) 式中再令 $m=m-n$，并带入 $\\theta_q(\\bm{x}_q,m)$ 就有\n$$ \\theta_k(\\bm{x}_k,n) = n(\\theta_k(\\bm{x}_k, 1)-\\theta_k(\\bm{x}_k, 0))+\\theta_k $$汇总 最后，我们将以上结果放在一起，就得到\n$$ f_q(\\bm{x}_q,m) = \\bm{q}e^{im\\theta}, f_v(\\bm{x}_k,m) = \\bm{k}e^{in\\theta} $$这里 $\\theta=\\theta_k(\\bm{x}_k, 1)-\\theta_k(\\bm{x}_k, 0)$ 是一个超参数，用于控制频率。\n我们记\n$$ R_{\\theta,m} = \\begin{bmatrix} \\cos m\\theta \u0026 -\\sin m\\theta\\\\ \\sin m\\theta \u0026 \\cos m\\theta \\end{bmatrix} $$则我们有：\n$$ f_q(\\bm{x}_q,m) = R_{\\theta,m}\\bm{q}, f_v(\\bm{x}_k,m) = R_{\\theta,n}\\bm{k}. $$并且\n$$ \\langle f_q(\\bm{x}_q,m), f_v(\\bm{x}_k,m)\\rangle = \\bm{q}^TR_{\\theta,m-n}\\bm{k} \\tag{5} $$多维扩展 上面是 2D 的情况，对于多维情况，苏剑林老师通过将两个元素组对，然后分别进行处理，得到了多维的情形：\n$$ R_{\\theta,m}^d = \\begin{bmatrix} R_{\\theta_1,m} \u0026 \u0026 \u0026\u0026 \u0026 \\\\ \u0026 \u0026 R_{\\theta_2,m} \u0026 \u0026 \u0026 \\\\ \u0026\u0026\u0026\u0026 \\ddots \u0026 \\\\ \u0026\u0026\u0026\u0026 \u0026 R_{\\theta_{d/2},m} \\end{bmatrix}\\in\\mathbb{R}^{d\\times d} $$我们可以验证公式 (5) 仍然是成立的。\nRoPE 的远程衰减性质 我们接下来看一下结果与相对距离 $m-n$ 之间的关系， 注意到\n$$ \\langle f_q(\\bm{x}_q,m), f_v(\\bm{x}_k,m)\\rangle = \\bm{q}^TR_{\\theta,m-n}\\bm{k} = \\sum_{i=1}^{d/2} \\bm{q}_i^TR_{\\theta, m-n}\\bm{k}_i $$这里 $\\bm{q}_i=[q_{2i},q_{2i+1}]^T$, $\\bm{k}_i=[k_{2i},k_{2i+1}]^T$ 分别是对应的 pair，我们考虑其中一个分量，不妨假设 $\\|\\bm{q}\\|_2=\\|\\bm{k}\\|_2=1$, 我们有\n$$ \\begin{aligned} \\bm{q}_i^TR_{\\theta, m-n}\\bm{k}_i \u0026\\leq \\bm{q}_i^TR_{\\theta, m-n}\\bm{q}_i\\\\ \u0026= \\bm{q}_i^T\\left(\\frac{R_{\\theta, m-n}+R_{\\theta, m-n}^T}{2}\\right)\\bm{q}_i\\\\ \u0026\\leq \\lambda_{\\max}\\left(\\frac{R_{\\theta, m-n}+R_{\\theta, m-n}^T}{2}\\right) \\\\ \u0026= \\cos (m-n)\\theta_i \\end{aligned} $$其中，第一个不等式是因为 两个向量相等时其内积最大，第二个不等式是由于二次型最大值为矩阵的特征值。\n这样我们就有\n$$ \\langle f_q(\\bm{x}_q,m), f_v(\\bm{x}_k,m)\\rangle \\leq \\sum_{i=1}^{d/2}\\cos (m-n)\\theta_i. $$我们可以简单画出对应的曲线：\n这里针对不同的配置，结果也会略微不同，具体分析可以参加知乎回答 RoPE的远距离衰减\nRoPE 代码实现与理解 Naive 实现 我们接下来看一下如何实现 RoPE\n$$ \\Theta_m\\bm{x}=\\begin{bmatrix} \\cos m\\theta_0 \u0026 -\\sin m\\theta_0 \u0026 \u0026\u0026\\cdots \u0026\\cdots \u0026\\cdots \\\\ \\sin m\\theta_0 \u0026 \\cos m\\theta_0 \u0026 \u0026\u0026\\cdots \u0026\\cdots \u0026\\cdots\\\\ \u0026 \u0026 \\cos m\\theta_1 \u0026 -\\sin m\\theta_1 \u0026 \\cdots \u0026\\cdots \u0026\\cdots\\\\ \u0026 \u0026 \\sin m\\theta_1 \u0026 \\cos m\\theta_1 \u0026 \\cdots \u0026\\cdots \u0026\\cdots \\\\ \u0026\u0026\u0026\u0026 \\ddots \u0026\\vdots \u0026\\vdots\\\\ \u0026\u0026\u0026\u0026 \u0026 \\cos m\\theta_{d/2} \u0026 -\\sin m\\theta_{d/2}\\\\ \u0026\u0026\u0026\u0026 \u0026 \\sin m\\theta_{d/2} \u0026 \\cos m\\theta_{d/2} \\end{bmatrix}\\begin{bmatrix} x_1\\\\ x_2\\\\ \\vdots \\\\ x_d \\end{bmatrix} $$在实现的时候，我们一般根据 $\\sin$ 和 $\\cos$ 进行分组，也就是\n$$ \\Theta_m\\bm{x}=\\begin{bmatrix} \\cos m\\theta_0\\\\ \\cos m\\theta_0\\\\ \\vdots\\\\ \\cos m\\theta_{d/2}\\\\ \\cos m\\theta_{d/2}\\\\ \\end{bmatrix}\\odot \\begin{bmatrix} x1\\\\ x2\\\\ \\vdots\\\\ x_{d-1}\\\\ x_d\\\\ \\end{bmatrix} + \\begin{bmatrix} \\sin m\\theta_0\\\\ \\sin m\\theta_0\\\\ \\vdots\\\\ \\sin m\\theta_{d/2}\\\\ \\sin m\\theta_{d/2}\\\\ \\end{bmatrix}\\odot \\begin{bmatrix} -\\ x_2\\\\ x_1\\\\ \\vdots\\\\ -x_d\\\\ x_{d-1}\\\\ \\end{bmatrix} $$我们通常按照奇偶 index 来分别计算，然后通过重排序来得到最终的结果，实现代码如下：\n1 2 3 4 5 6 7 8 9 10 11 12 13 def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -\u0026gt; torch.Tensor: x_even = x[..., ::2] # (seq_len, d_k_half) x_odd = x[..., 1::2] # (seq_len, d_k_half) odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half) evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half) stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half) stacked_trans = rearrange( stacked, \u0026#34;... seq_len double d_k_half -\u0026gt; ... seq_len d_k_half double\u0026#34; ) # (...,seq_len, d_k_half, 2) out = rearrange( stacked_trans, \u0026#34;... seq_len d_k_half double -\u0026gt; ... seq_len (d_k_half double)\u0026#34; ) # (..., seq_len, d_k) return out LLaMA 实现 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(end, device=freqs.device) # type: ignore freqs = torch.outer(t, freqs).float() # type: ignore freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): ndim = x.ndim assert 0 \u0026lt;= 1 \u0026lt; ndim assert freqs_cis.shape == (x.shape[1], x.shape[-1]) shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] return freqs_cis.view(*shape) def apply_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, ) -\u0026gt; Tuple[torch.Tensor, torch.Tensor]: xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) freqs_cis = reshape_for_broadcast(freqs_cis, xq_) xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) 在 LLaMA 中，我们首先还是计算 $\\theta_i$, 然后在计算的过程中，我们将 $(x_i,x_{i+1})$ 视作一个复数，然后 乘以 $\\exp(im\\theta)$, 最后再取实部得到最终的结果\n通用实现 实际上，naive 版本的实现与现在大语言模型所采用的实现并不一致，我们先看一下现有的大语言模型的 RoPE 实现，这里我们将 LLaMA的transformer代码 放在下面，\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 def rotate_half(x): \u0026#34;\u0026#34;\u0026#34;Rotates half the hidden dims of the input.\u0026#34;\u0026#34;\u0026#34; x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class LlamaRotaryEmbedding(torch.nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer(\u0026#34;inv_freq\u0026#34;, inv_freq) self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum(\u0026#34;i,j-\u0026gt;ij\u0026#34;, t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\u0026#34;cos_cached\u0026#34;, emb.cos()[None, None, :, :].to(dtype), persistent=False) self.register_buffer(\u0026#34;sin_cached\u0026#34;, emb.sin()[None, None, :, :].to(dtype), persistent=False) def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len \u0026gt; self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), ) 我们将上述代码翻译成公式，现在我们的 $\\Theta$ 变成了 (对应 emb = torch.cat((freqs, freqs), dim=-1))\n$$ \\Theta = [\\theta_0,\\dots,\\theta_{d/2},\\theta_0,\\dots,\\theta_{d/2}]^T $$实际上 $\\sin$ 部分对应的向量现在变成了\n$$ [-x_{d/2+1}, -x_{d/2+2}, \\dots, -x_{d}, x_1, \\dots, x_{d/2}]^T $$我们带回到原始公式，可以得到对应的 RoPE 操作变成了\n$$ R_{\\theta,m}^d=\\begin{bmatrix} \\cos m\\theta_0 \u0026 \u0026 \u0026 -\\sin m\\theta_0 \u0026 \\cdots \u0026\\cdots \u0026\\cdots \\\\ \u0026 \u0026 \\cos m\\theta_1 \u0026 \u0026-\\sin m\\theta_1 \u0026 \\cdots \u0026\\cdots \\\\ \u0026 \u0026 \u0026 \\cos m\\theta_2 \u0026 \u0026-\\sin m\\theta_2 \u0026 \\cdots \\\\ \\vdots\u0026\\vdots\u0026\\vdots\u0026\\vdots\u0026 \\vdots \u0026\\vdots \u0026\\vdots\\\\ \u0026 \u0026 \u0026 \u0026\\cos m\\theta_{d/2 - 1} \u0026 \u0026 -\\sin m\\theta_{d/2 - 1} \\\\ \\sin m\\theta_0 \u0026\u0026 \u0026 \\cos m\\theta_0 \u0026\u0026\\cdots \u0026\\cdots \\\\ \u0026 \u0026 \\sin m\\theta_1 \u0026 \u0026\\cos m\\theta_1 \u0026 \\cdots \u0026\\cdots \\\\ \u0026 \u0026 \u0026 \\sin m\\theta_2 \u0026 \u0026\\cos m\\theta_2 \u0026 \\cdots \\\\ \\vdots\u0026\\vdots\u0026\\vdots\u0026\\vdots\u0026 \\vdots \u0026\\vdots \u0026\\vdots\\\\ \u0026\u0026\u0026\u0026 \\sin m\\theta_{d/2 - 1}\u0026 \u0026 \\cos m\\theta_{d/2 - 1} \\end{bmatrix} $$这列每一行的 $\\cos$ 和 $\\sin$ 都相差了 $d/2$ 列.\n因此，这里的区别在于，原始 RoPE 计算的 pair 为 $(x_{i}, x_{i+1})$, 而 LLaMA 里的 RoPE 计算的 pair 为 $(x_{i}, x_{i+d/2})$. transformers library 使用这种方式，可以减少计算量，提高整体的计算效率。\n为了适应使用 LLaMA 中实现的 RoPE 的，Huggingface 对权重进行了转换，使得基于原始 RoPE 实现的模型也可以获得加速.\n假设 $d=8$，原始 RoPE 的 pair 为 [(q_0, q_1), (q_2, q_3), (q_4, q_5), (q_6, q_7)], 新的 pair 为 [(q_0, q_4), (q_1, q_5), (q_2, q_6), (q_3, q_7)]. 我们希望对 index 进行 remap，我们发现一个满足条件的 permutation 为 [0, 2, 4, 6, 1, 3, 5, 7], 也就是 q_0-\u0026gt;q_0, q_2-\u0026gt;q_1, \u0026hellip;, q_7-\u0026gt;q_7.\n但是，如果我们在推理时这样做，就会降低整体速度，因此 Huggingface 的做法是改变 $W_Q$ 和 $W_K$ 的权重，具体来说，就是 $\\Pi q=(\\Pi W_Q)x$， 左边是在线转换，右侧离线转换。转换好 $W_Q$ 之后，正常计算就可以了。具体代码 为\n1 2 3 4 5 6 7 8 9 10 11 12 13 # permute for sliced rotary def permute(w, n_heads=n_heads, dim1=dim, dim2=dim): return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2) state_dict = { f\u0026#34;model.layers.{layer_i}.self_attn.q_proj.weight\u0026#34;: permute( loaded[f\u0026#34;layers.{layer_i}.attention.wq.weight\u0026#34;] ), f\u0026#34;model.layers.{layer_i}.self_attn.k_proj.weight\u0026#34;: permute( loaded[f\u0026#34;layers.{layer_i}.attention.wk.weight\u0026#34;] ), ... } 结论 本文中，我们回顾了位置编码，包括绝对位置编码和相对位置编码，我们着重介绍了 RoPE 的原理，推导以及代码实现。\n参考文献 You could have designed state of the art positional encoding Is LLaMA rotary embedding implementation correct? [LLaMA] Rotary positional embedding differs with official implementation RoPE blog RoFormer 位置编码之路 RoPE的远距离衰减 ","date":"2025-05-19T10:46:39+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-position-encoding/","title":"Notes on Position encoding"},{"content":"Qwen 在 2025 年 5 月发布了 Qwen3 系列大语言模型，Qwen3 包括 6 个 dense 模型和 2 个 MoE 模型，主要亮点为多语种能力，自适应快慢思考能力以及支持用户设置 thinking budget.\nQwen3 包括 6 个 dense 模型和 2 个 MoE 模型，其旗舰模型是一个 235B 的 MoE 模型，激活参数为 22B. Qwen3 系列的主要亮点如下:\n快慢思考融合，模型原生支持在 reasoning/non-reasoning 模式之间切换 Reasoning budget, 用户可以指定思考需要的 budget，来平衡 latency 和 performance Distillation, 使用蒸馏的方法训练小模型，大幅度提高模型的表现 多语种支持，相比于 Qwen2.5，Qwen3 支持 119 中语言和方言 Method Model Qwen3 的 dense 模型的架构与 Qwen2.5 基本一致，包括使用 GQA , SwiGLU, RoPE, RMSNorm 和 pre-normalization. Qwen3 进一步移除了 QKV bias, 然 后加入了 QK-Norm 来提高训练的稳定性。\nQwen3 的 MoE 架构使用了 128 个专家，激活专家个数为 8 个。与 Qwen2.5-MoE 不同，Qwen3 里没有使用 shard experts。并且，Qwen3 加入了 global-batch load balancing loss,来提高 expert 的特化程度。\n在 tokenizer 方面，Qwen 系列的 tokenizer 一直都是一样的，这也是 Qwen 系列领先的一点。\n模型的具体参数如下两张表所示。\nMoE 架构：上下文长度为 128K，128 个专家，每个 token 由 8 个专家负责处理\nQwen3-235B-A22B, 总参数 235B，激活参数 22B Qwen3-30B-A3B, 总参数 30B，激活参数 3B Models Layers Heads (Q / KV) # Experts (Total / Activated) Context Length Qwen3-30B-A3B 48 32 / 4 128 / 8 128K Qwen3-235B-A22B 94 64 / 4 128 / 8 128K dense 架构: Qwen3-32B, Qwen3-14B, Qwen3-8B, Qwen3-4B, Qwen3-1.7B, and Qwen3-0.6B\nModels Layers Heads (Q / KV) Tie Embedding Context Length Qwen3-0.6B 28 16 / 8 Yes 32K Qwen3-1.7B 28 16 / 8 Yes 32K Qwen3-4B 36 32 / 8 Yes 32K Qwen3-8B 36 32 / 8 No 128K Qwen3-14B 40 40 / 8 No 128K Qwen3-32B 64 64 / 8 No 128K Pre-training 预训练数据一共包括 36T token，覆盖了 119 种语言。数据包括 coding, STEM, reasoning, books, multilingual texts 以及合成数据。\n为了扩展训练数据，作者微调了 Qwen2.5-VL 来从 PDF 文档中提取文字，然后使用 Qwen2.5 来进行修正。最终收集到了几 T 的 token。另外，作者还使用 Qwen2.5, Qwen2.5-Math, Qwen2.5-Coder 来合成不同格式的数据，包括教科书，QA，指令以及代码片段等。最后，作者加入了更多的多语种数据。\n作者从 educational value, fields, domains 以及 safety 对数据进行了标注。在数据混合时，Qwen3 在 instance 层面进行操作。\n预训练阶段包括 3 个 stage：\nGeneral Stage (S1): 这一阶段的目的是让模型掌握世界知识，使用了 30T 的 token，模型上下文长度为 4096 Reasoning Stage (S2): 这一阶段的目的是提高模型的推理能力，使用了 5T 的高质量 token，模型上下文长度为 4096，数据包括 STEM, coding, reasoning 以及合成数据 Long Context Stage (S3): 这一阶段的目的是提升模型的长上下文能力，使用了几百 B的 token，模型上下文长度为 32768.训练时数据混合 75% 的长文档数据，25% 的短文本数据。作者将 RoPE 的 frequency 从 10000 提升到了 1,000,000. 作者还是用 YARN 以及 Dual Chunk Attention 来提高 inference 效率 对 pre-training 的 base model 进行评测之后，作者发现：\nQwen3-235B-A22B-Base 超过了其他 base 模型的表现，包括 DeepSeek-V3 Base, Llama-4-Maverick Base, Qwen2.5-72B Base Qwen3-MoE 模型与相同大小的 Qwen3-Dense 模型参数相比，其只需要 1/5 的参数就可以达到相同的表现 Qwen3-MoE 模型与 2 倍参数量的 Qwen2.5-MoE 模型表现差不多 Qwen3-Dense 模型与大一个量级的 Qwen2.5-Dense 模型表现差不多 Post-training Qwen3 的 post-training 如下图所示：\n对于旗舰模型 (Qwen3-235B-A22B, Qwen3-32B) 的训练，Qwen3 使用了一个四阶段的训练 pipeline。对于轻量化模型（其他模型）的训练，Qwen3 使用了知识蒸馏。\n旗舰模型的训练包括四个阶段，前两个阶段用于提升模型的 reasoning 能力，后两个阶段用于将 reasoning 和 non-reasoning 能力结合起来。\nFlashship Model Stage 1 (Long CoT Cold Start) 这个阶段的目的是让模型掌握 reasoning 的基础。这个阶段使用了数学，代码，逻辑推理和通用的 STEM 相关问题。每个问题都有参考答案或者 test-cases. 作者使用了 Qwen2.5-72B 来过滤数据，包括 non-verifiable prompts 以及太简单的 prompt. 作者认为，这一阶段应该减少训练使用的样本和训练步数。\nStage 2 (Reasoning RL) 这个阶段的目的是提升模型的 reasoning 能力。该阶段使用了 3,995 条过滤得到的样本，算法为 GRPO. 作者发现提高 batch size 和每个 query 的 rollouts 可以提高模型的表现。作者通过调整模型的 entropy 来控制 exploration 和 exploitation 的平衡\nStage 3 (Thinking Mode Fusion) 这一阶段的目的是将 non-reasoning 能力加入到之前的 reasoning 模型中。作者在第二阶段的 model 上进行了 continual SFT，然后构建了一个 chat template 用于融合两种模式。\nreasoning 数据来源于 stage1 的 rejection sampling 和 stage 2 的模型. non-reasoning 数据来源于各种任务，如 coding, math, multilingual 等。为了保证模型的多语种能力，作者还加入了一些翻译相关的数据。\n作者还构建了一个 chat template, 用于统一数据格式。chat template 如下图所示\n作者使用 /think 和 /no_think 来标记两种模式，对于 non-reasoning mode, 其 \u0026lt;think\u0026gt;\u0026lt;/think\u0026gt; 会被置空。模型在默认情况下处于 reasoning mode, 因此作者加入了一些不包含 /think 的 reasoning 数据。\n作者发现，通过这种 Think mode fusion, 模型可以学会在 reasoning mode 和 non-reasoning mode 下进行回答，因此，模型也可以基于中间结果来给出最终的答案。 当超出 budget 之后，作者使用以下 Instruction\nConsidering the limit time by the user. I have to give the solution based on the thinking directly now. \\n\u0026lt;/think\u0026gt;.\\n\\n\n来让模型直接终止思考二给出最终的答案。\nStage 4 (General RL) 这个阶段的目的是提升模型在不同场景下的能力。作者构建了一个 reward system 来覆盖 20 多种不同的任务。这些任务包括：instruction following, format following, preference alignment, agent ability 以及 abilities for specialized scenarios.\n作者构建了三种不同的 rewards:\nRule-based rewards: 覆盖的任务包括 instruction following 和 format following Model-based rewards: 作者使用 Qwen2.5-72B 来判别答案的正确性 Model-based Reward without reference answer: 作者训练一个 reward model 来给模型的回答进行打分 Lightweight Model 对于轻量化的模型，作者发现直接通过蒸馏可以有效提高学生模型的表现，并且训练效率也更高。蒸馏训练包括两个阶段：\nOff-policy Distillation: 这个阶段的目的是让模型拥有基本的 reasoning 能力并且可以在不同的模式中进行切换。作者使用了教师模型的 reasoning 输出和 non-reasoning 输出来蒸馏学生模型 On-policy Distillation: 在这个阶段，学生模型生成回答，然后基于教师模型的输出，使用 KL-divergence 来更新学生模型的参数 Evaluation Thinking budget. 作者发现当我们提高 Thinking budget 之后，模型的表现是可以持续提升的。结果如下图\nEfficiency of distillation. 作者发现使用 distillation 可以大幅度提高模型的表现和训练效率。下面是结果如下图所示\nEffects of Thinking mode fusion and RL 作者进一步探究了三个 stage 对模型表现的影响，为此，作者构建了 in-house benchmarks 来评估模型的表现，这些 benchmarks 包括：\nCounterFactrQA. 问题是不符合事实的，用于评估模型的幻觉 LengthCtrl. 有长度要求的写作任务，评估生成内容长度和给定长度之间的差别 ThinkFollow. 多轮对话，每轮对话随机插入 /think 和 /no_think flag，评估模型是否能在两种模式之间切换 Tooluse. 评估模型的工具调用能力 结果如下\n结论如下：\nStage3 可以提高模型在两种 reasoning mode 切换的能力，并且 stage3 还可以提高模型的通用以及 instruction following 能力 Stage4 进一步提高模型在两种模式下的通用，instruction following 和 agent 能力 Stage3 和 stage4 并没有显著提高模型在 knowledge, STEM, math 和 coding 相关任务上的表现。甚至在一些竞赛如 AIME24 上模型的表现还有所下降，作者认为这是由于我们提升了模型的通用能力而导致其特化能力下降导致的，作者认为作为一个通用模型，这是可以接受的。 Conclusion 在本文中，作者提出了 Qwen3 系列大语言模型，包括 6 个 Dense 模型和 2 个 MoE 模型。Qwen3 模型标志了一个新的 SOTA，其特点主要是快慢思考结合，thinking budget，以及多语种。\n作者认为后续工作有以下几点：\n使用更高质量的数据来进行预训练 优化模型架构和训练方式，提升模型的上下文 提高针对 RL 的计算资源，来进一步提高模型的 agent 能力 References Arxiv Github ","date":"2025-05-15T14:48:11+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen3/","title":"Notes on Qwen3"},{"content":"字节Seed在5月11号发布了Seed1.5-VL技术报告。技术报告详细介绍了Seed1.5-VL的架构，训练和评估细节\n简介 Seed1.5-VL是多模态大模型，它是一个标准的 Vision Encoder - MLP - LLM 的架构，Vision Encoder是自研的Seed-ViT，参数量为532M，大语言模型的激活参数为20B（虽然论文未提及总参数量，但是从seed-Thinking v1.5来看，总参数量应该是200B）。Seed1.5-VL强调了agent能力，比如GUI control和gameplay。\n介绍 现有的LVLM还不能与人类的通用能力相比，特别是在3D空间理解，物体技术，交互游戏等方面。并且，LVLM训练预料的多样性也不如LLM。最后，多模态数据的异质性也对模型的训练和推理提出了挑战。\n基于这些问题，作者提出了Seed1.5-VL，一个视觉多模态大模型。通过构建高质量的训练数据，作者大幅度提高了模型的表现。作者还进一步探究了模型的scaling law以及pre-training和post training算法的设计。为了提高训练和推理效率，作者在infra上也进行了改进\n架构 Seed1.5-VL的架构如下图所示\nSeed1.5-VL的架构包含三个部分：\nVision Encoder: vision encoder是seed自研的Seed-ViT，拥有532M参数 Adapter：一个两层的MLP LLM：Seed1.5-LLM，一个MoE架构的LLM Seed1.5-VL支持文本，图片和视频输入，输入为文本。其可以进行理解和推理(reasoning)。\nSeed-ViT Seed-ViT主要是在两个方面进行了改进：\n在pre-training阶段使用了视频数据，来提高模型的时空间感知能力 使用了2D-RoPE来处理动态分辨率的图片 Seed-ViT的超参数如下表所示\nPatch size Pos Embedding Head dim #heads Embedding dim MLP ratio #layers 14 2D RoPE 64 20 1280 4.0 27 对于输入的图片，Seed-ViT首先将图片resize到28*28的倍数，接下来图片会被分割为14*14的patch. 与NaViT一样，Seed-ViT也是用了token-packing的操作，来将多个图片packing到一个sequence里进行训练。最后，对于Seed-ViT输出的图片特征，作者还是用了 $2\\times 2$的 average pooling来降低token个数。\nSeed-ViT的预训练包括三个阶段，其设置和超参数如下表所示：\nCategories Unlabeled Image Image-text pairs Video-audio-text tuples Training samples 2.2B 4.8B 65M Token percentages 4.0% 91.2% 4.8% Batch sizes 55,296 32,768 1,024 LR warm up steps 1,692 2,000 12,800 Maximum LR $7.06\\times 10^{-3}$ $1.0\\times 10^{-4}$ $5.0\\times 10^{-5}$ Minimum LR $1.05\\times 10^{-5}$ $1.2\\times 10^{-6}$ $2.02\\times 10^{-7}$ 在预训练时，作者考虑了三点：\n训练效率：尽可能提高模型的训练效率 原生图片精度输入处理：让模型可以尽早处理不同分辨率图片输入 数据利用率：可以使用不同类型的数据 基于这三点考虑，Seed-ViT的训练包括三个阶段：\nMasked Image modeling with 2D RoPE：这一阶段的目的是提高模型对与几何和结构的感知能力。作者使用 EVA02-CLIP-E作为教师模型，然后随机将75%的patches进行mask，最后让Seed-ViT输出的特征尽可能匹配教师模型输出的特征。这里的损失函数是cosine similarity. 作者发现，scale up MIM之后，VLM在OCR和文档理解任务上的表现得到了大幅度的提升。 Native Resolution Contrastive Learning：这一阶段和CLIP,SigLIP的训练差不多，text encoder与EVA02-CLIP-E相同。损失函数为SigLIP loss和SuperClass loss Omni-modal pret-raining：这一阶段使用了MiCo框架，来将video frames, audio, visual captions以及audio captions进行对齐，text encoder编码caption，ViT编码其他部分。通过这一阶段的训练，模型可以学习omni-model representations,并且还可以提高模型的图片和视频理解人物上的表现。 视频编码 对于视频输入，Seed1.5-VL采用了 Dynamic Frame-Resolution Sampling技巧，核心思想在于根据budget和处理的任务来动态调整采样率(frame)和每一帧的图片精度(resolution)。\n具体来讲，在时序上，默认的FPS为1，如果要处理有关时序信息任务的话，FPS调整为2，如果要处理motion tracking等任务的话，FPS提高到5.为了进一步提高模型的时间感知能力，Seed 1.5-VL加入了timestamp tokens (e.g., [1.5 second])。\n在空间上，每个frame的可选精度为 $\\{640, 512, 384, 256, 160, 128\\}$，选择的精度根据我们的budget进行决定。如果说我们使用最低配置还是不能处理视频输入的话，我们就会先进行均匀采样，然后再进行处理，这样可以让模型处理所有的视频信息。\nPre-training pre-training阶段一共包括了3T token. 作者详细介绍了每一个类别。这里我们就简单总结一下。数据一共包括以下7个类别\nGeneric Image-text pairs \u0026amp; knowledge data. 这部分数据主要是图文交错以及image text-pairs数据。作者介绍了一下针对长尾数据（稀有物种识别）的数据构建方式 OCR. 这部分数据有1B左右的sample，包括documents, scene texts, tables, charts以及flowcharts等。其中，作者通过SynthDog和Latex合成了200M左右的文本密集图片。作者还对合成的数据进行了增广;Chart数据包括公开数据集和合成的数据集，最终一共有100M的样本; Table数据集也是通过提取得到的，一共包含50M左右的图片。作者还基于OCR结果构建了VQA数据集，来提高模型理解图片文字的能力 Visual Grounding \u0026amp; Counting: 提高模型识别和定位物体的能力。数据格式包括bounding box和center points. 其中bounding box主要是公开数据集；作者还是用Grounding DINO来标记了一部分数据，最后一共收集了200M的样本; point数据主要是PixMo, 作者还使用了Molmo以及CountGD标注了一部分数据，最后一共包括170M的指令以及110B的token; Counting数据从bounding box 和point数据集中筛选得到，包括了8M的样本和13B的token。最后在训练时，与InternVL一样，坐标会被normalize到 [1,1000] 3D spatial understanding. 主要包括相对深度，绝对深度和3D Grounding数据；对于相对深度，作者使用DepthAnything V2来标注，得到了3.2B的token; 对于绝对深度，作者使用了公开数据集，最终包含18M的指令以及28B的token;对于3D Grounding，作者使用了公开数据集，最终包含770的指令以及1.3B的token Video. 主要包含general video understanding数据, video temporal grounding and moment retrieval以及video streaming data. STEM. 主要包含image comprehension data和problem-solving data. 前者包括收集的3.2M教育相关的数据，以及10M结构化表格，4.5M化学结构表达式，1.5M坐标系数据。后者包括100M K12的习题 GUI. 作者使用UI-TARS来合成数据。数据主要包括perception, grounding 以及 reasoning 预训练包括三个stage，其实验设置如下表所示：\nStages Stage 0 Stage 1 Stage 2 Training budget (tokens) 16B 3T 240B Sequence length 32,768 32,768 131,072 Trainable components MLP adaptor all all Batch sizes (tokens) 8.4M 71M 71M LR warmup steps 100 500 0 Maximum LR $2.52 \\times 10^{-4}$ $5.22 \\times 10^{-5}$ $5.22 \\times 10^{-6}$ Minimum LR $4.50 \\times 10^{-5}$ $5.22 \\times 10^{-6}$ $5.22 \\times 10^{-6}$ Stage 0: 仅训练MLP，用于对齐vision encoder和LLM,使用了16B token Stage 1：训练全部参数，用于获取知识以及掌握grounding和OCR能力，使用了3T token Stage 2：训练全部参数，提升模型的长上下文能力以及对下游任务的适应性，包括视频理解，coding和3D空间理解，上下文长度从32,768提升到131,072 作者还进行了消融实验，采取了和Qwen2-VL一样的训练方式，结果发现，本文的训练方式效果更好。作者认为，解冻vision encoder，可能让vision encoder弥补LLM的不足，进而损害其自身的感知能力。\n作者还探究了模型的scaling law，使用的建模函数为：\n$$ \\log(\\hat{L}) \\sim \\log(B) - \\beta \\log(D) = -a \\log(D) + b $$ 这里 $L$ 是NLL loss, D是训练的token数，$a$和$b$是待拟合的参数。实验结果如下图所示。 结论是在特定领域数据训练的loss可以作为模型在下游任务上表现的一个估计\nPost-training post-training分为两个阶段：SFT和RL，其中RL包括RLHF, RLVR. 总流程如下图所示\nSFT SFT数据主要包括General Instruction data以及Long Chain-of-Thought (LongCoT) data，劝着途胜模拟徐保国的指令跟随能力，后者提升模型的reasoning能力。\n通用指令微调数据包括13,000条收集的数据和40,0000高质量的公开数据集。\nRLHF RLHF的数据包括人类标注数据和合成数据。对于reward model，作者使用了LVLM作为reward model，来给不同的response进行打分。最后，使用变体的PPO算法进行训练。\nRLVR RLVR的训练数据主要包含Visual STEM和Visual Perception and Reasoning. 前者包含1M左右的问题，主要是STEM和K12的相关问题。后者包括grounding, visual instruction following以及visual puzzles \u0026amp; games等\nHybrid RL RL的训练过程包含了RLHF和RLVR两个RL算法。作者介绍了以下细节：\nformat reward. 要求模型输出格式为 \u0026lt;think\u0026gt;{thought}\u0026lt;/think\u0026gt;{solution} hybrid reward. reward包括RM和verifier两部分 Shared critic. 使用一个critic model来估计value function KL coefficients. 对于RLHF和RLVE，使用不同的KL divergence稀疏 Iterative Update by Rejection Sampling Fine-tuning 在训练的过程中，作者使用了一个迭代训练策略。一个开始模型在低质量的Long-CoT数据上进行SFT，在RL过后，作者会重新评估问题的难度，然后过程掉一部分数据，再进行SFT，这样反复进行迭代训练以提高模型的表现。\nInfra部分好像都是一些细节，这里就跳过了。\nEvaluation 作者评估了\nSeed-ViT的表现 Seed1.5-VL的通用视觉理解和推理能力 Seed1.5-VL的视频相关能力 Seed1.5-VL的agent能力 Seed1.5-VL在内部benchmark上的表现 由于Seed1.5-VL是一个API模型，并且其内部benchmark不公开，因此Seed-ViT以及其在内部benchmark上的表现我们就不列出来了。我们主要看一下其和领先多模态大模型的对比\nConclusion 在本文中，作者介绍了Seed1.5-VL，一个领先的多模态大模型，可以处理图片和视频输入。作者介绍了Seed-ViT以及Seed1.5-VL的数据，训练和评估。\n作者还分析了以下Seed1.5-VL的不足以及未来的改进方向\nSeed1.5-VL在复杂视觉感知任务中，其计数能力会下降。 Seed1.5-VL的高阶推理能力还尚有不足，作者拿Klotski举了一个例子 Seed1.5-VL的3D spatial reasoning能力不足 Seed1.5-VL的temporal reasoning能力不足 Seed1.5-VL与其他模型一样，存在幻觉问题 未来，有以下值得探索的方向：\n涌现现象(emergent abilities) scaling law，进一步探索模型参数量，算力以及数据对模型表现的影响 提高模型的3D spatial reasoning能力，降低模型的幻觉，提高模型分组合搜索能力 参考文献 arxiv ","date":"2025-05-14T09:28:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-seed1.5-vl/","title":"Notes on Seed1.5-VL"},{"content":"在本文中，我们将要分析与大语言模型相关的参数量和计算量。在计算之前，我们会首先回顾一下大语言模型的架构\n大语言模型架构 大语言模型参数计算 计算量估计 checkpointing KV cache 参考文献 回旋托马斯x 文章 ","date":"2025-05-13T11:26:36+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/%E5%88%86%E5%B8%83%E5%BC%8F%E8%AE%AD%E7%BB%83%E5%8F%82%E6%95%B0%E9%87%8F%E4%B8%8E%E8%AE%A1%E7%AE%97%E9%87%8F%E5%88%86%E6%9E%90/","title":"分布式训练：参数量与计算量分析"},{"content":"本节中，我们将介绍模型训练的基本数学原理，以及在分布式训练中我们需要考虑的精度，优化器等问题。\n训练的数学原理 在最优化里面，我们需要解决的问题一般有如下形式：\n$$ \\min_x\\ f(x) $$ 这里 $f$是我们的目标函数, $x$是我们的变量。一个比较简单的例子就是求一个给定函数的最小值。\n如果说，我们想要基于数据来训练一个模型，这个时候，我们目标函数的输入就包括两部分，一部分是模型参数，另一部分是数据，为了方便起见，我们使用$\\theta$来代表模型的参数，用 $\\{x_i,y_i\\}_{i=1}^N$ 来表示模型的训练集。上述的优化问题改写如下：\n$$ \\min_{\\theta}\\ \\frac{1}{N}\\sum_{i=1}^Nf(x_i,y_i\\;\\theta) $$比如我们训练 resnet 作为分类器，那么 resnet 的模型参数就是这里的 $\\theta$, 训练集就是我们的图片和对应的标签，比如ImageNet等，对应的$f$可以设定为 cross entropy loss.\n现在有了优化问题之后，我们就需要设计算法求解这个优化问题。一个最简单的优化算法就是梯度下降算法：\n$$ \\theta^{k+1} = \\theta^k - \\alpha_k\\frac{1}{N}\\sum_{i=1}^N\\nabla_{\\theta}f(x_i,y_i;\\theta^k) $$ 这里 $\\nabla_{\\theta}f(x;\\theta^k)$ 是 $f$ 相对于 $\\theta$ 在 $\\theta^k$ 处的梯度。\n但是，当我们模型过于复杂的时候，梯度往往计算起来非常复杂。为了简化模型的训练，现在的框架如tensorflow和pytorch都支持自动微分。因此，我们只需要定义如何从输入 $(x_i,y_i)$计算得到 $f(x_i,y_i;\\theta)$ 就可以了，框架会帮我们计算参数的梯度。\n自动微分 自动微分的目的是将求导的过程交给框架，从而让用户专注于模型的开发（也就是设计forward函数）。\n自动微分的核心思想就是链式法则.\n假设我们有一个复合函数 $y= f(g(h(x)))$, 一般来说，我们求 $y$ 关于 $x$的导数过程为 $$ \\frac{dy}{dx} = \\frac{dy}{df}\\frac{df}{dg}\\frac{dg}{dh}\\frac{dh}{dx} $$ 如果我们的中间函数 $g$, $h$非常复杂的话，那么整个求导过程就会非常复杂。而链式法则则是将这样一个全局过程给分解成了若干个局部过程。我们将 $y$ 表示为：\n$$ \\begin{aligned} y \u0026= f(y_1)\\\\ y_1\u0026=g(y_2)\\\\ y_2\u0026=h(y_3)\\\\ y_3\u0026=x \\end{aligned} $$接下来，\n总结 在本文中，我们简单介绍了一下如何训练一个模型，我们使用pytorch作为例子展示了现在训练框架的工作方式。在下一篇博客中，我们将会探究训练精度和优化器。训练精度和优化器是模型在训练过程中需要考虑的重点之一。\n训练 优化器 SGD Adam AdamW 精度 ","date":"2025-05-13T11:26:36+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/%E5%88%86%E5%B8%83%E5%BC%8F%E8%AE%AD%E7%BB%83%E5%A6%82%E4%BD%95%E8%AE%AD%E7%BB%83%E4%B8%80%E4%B8%AA%E6%A8%A1%E5%9E%8B/","title":"分布式训练：如何训练一个模型"},{"content":" 说明：本文参考了 nanotron/ultrascale-playbook 和 Colossal-AI Concepts\n什么是分布式系统 分布式系统允许一个软件的多个组件运行在不同的机器上。与传统集中式系统不一样，分布式系统可以有效提高系统的稳健性。 一个比较比较经典的分布式就是Git，Git允许我们把代码保存在多个remote上。这样当一个remote宕机时，其他remote也能提供服务。\n评估一个分布式系统的重要标准就是规模效益(scalablity)，也就是说，我们希望使用8台设备应该要比4台设备快2倍。但是，由于通信带宽等原因，实际上加速比并不是和设备数量成线性关系。因此，我们需要设计分布式算法，来有效提高分布式系统的效率。\n为什么需要分布式训练 我们需要分布式训练的原因主要是以下几点：\n模型越来越大。当下（2025）领先模型如Qwen，LLaMA系列的最大模型都超过了100B [2][3]。LLaMA系列最大的模型甚至超过了1000B。Scaling law告诉我们模型表现与参数量，算力，数据量成正相关关系。 数据集越来越大。现在领先的模型需要的数据量基本都需要100M以上，而大语言模型训练需要的token数量也都超过了10T的量级 [2][3]. 算力越来越强。现有最强的GPU H100其显存为80GB，拥有3.35TB/s 的带宽 (PcIe)，这让训练大规模模型成为可能。 超大的模型使得我们很难在一张GPU上进行训练，甚至我们都很难使用单张GPU进行部署。而10T级的数据也也需要几个月的时间才能训练完毕。因此，如何高效利用多张GPU在大规模数据上训练超大模型就是我们需要解决的问题。\n基本概念 我们先来熟悉一下分布式训练中的一些基本概念：\nHost: host (master address)是分布式训练中通信网络的主设备(main device). 一般我们需要对其进行初始化 Node: 一个物理或虚拟的计算单元，可以是一台机器，一个容器或者一个虚拟机 Port: port (master port)主要是用于通信的master port Rank: rank是通信网络中每个设备唯一的ID world size: world size是通信网络中设备的数量 process group: 一个process group是通信网络中所有设备集合的一个子集。通过process group, 我们可以限制device只在group内部进行通信 我们以下图为例： 上图中一共包含2个node (2台机器)，每台机器包含4个GPU (device)，当我们初始化分布式环境时，我们一共启动了8个进程（每台机器4个进程），每个进程绑定一个GPU。\n在初始化分布式环境之间，我们需要指定host和port。假设我们指定host为node 0和port为 29500，接下来，所有的进程都会基于这个host和port来与其他进程连接。默认的process group（包含所有device）的world size 为8. 其细节展示如下\nprocess ID rank Node index GPU index 0 0 0 0 1 1 0 1 2 2 0 2 3 3 0 3 4 4 1 0 5 5 1 1 6 6 1 2 7 7 1 3 我们可以创建一个新的process group，使其仅包含ID为偶数的process：\nprocess ID rank Node index GPU index 0 0 0 0 2 1 0 2 4 2 1 0 6 3 1 2 Remark: 注意，rank与process group相关，一个process在不同的process group里可能会有不同的rank.\n通信方式 接下来，我们需要介绍一下设备间的通信方式，这是我们后面分布式训练算法的基础。根据设备数量的不同，我们可以将设备间通信分为：\none-to-one: 两个device之间互相进行通信 one-to-many: 一个device与多个device进行通信 many-to-one: 多个device与一个device之间进行通信 many-to-many: 多个device之间互相进行通信 One-to-one One-to-one的情况很简单，一个process与另一个process进行通信，通信通过 send 和 recv 完成。还有对应的 immediate版本，即 isend 和 irecv，示意图如下所示\n测试代码如下：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 # send_recv.py import os import torch import torch.distributed as dist def init_process(): dist.init_process_group(backend=\u0026#39;nccl\u0026#39;) torch.cuda.set_device(dist.get_rank()) def example_send(): if dist.get_rank() == 0: tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda() dist.send(tensor, dst=1) elif dist.get_rank() == 1: tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f\u0026#34;Before send on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.recv(tensor, src=0) print(f\u0026#34;After send on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_send() # run with # torchrun --nproc_per_node=2 send_recv.py 结果输出\n1 2 Before send on rank 1: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:1\u0026#39;) After send on rank 1: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:1\u0026#39;) 注：为了方便，后续代码仅定义函数和运行方式，init_process()和import部分省略\nsend/recv的特点是在完成通信之前，两个process是锁住的。与之相反，isend/irecv 则不会加锁，代码会继续执行然后返回Work对象，为了让通信顺利进行，我们可以在返回之前加入wait()\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 # isend_irecv.py def example_isend(): req = None if dist.get_rank() == 0: tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda() req = dist.isend(tensor, dst=1) print(\u0026#34;Rank 0 is sending\u0026#34;) elif dist.get_rank() == 1: tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f\u0026#34;Before irecv on rank {dist.get_rank()}: {tensor}\u0026#34;) req = dist.irecv(tensor, src=0) print(\u0026#34;Rank 1 is receiving\u0026#34;) req.wait() print(f\u0026#34;After isend on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_isend() # run with # torchrun --nproc_per_node=2 isend_irecv.py 结果输出\n1 2 3 4 Before irecv on rank 1: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:1\u0026#39;) Rank 0 is sending Rank 1 is receiving After isend on rank 1: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:1\u0026#39;) 由于isend/irecv这种不锁的特性，我们不应该\n在dist.isend()之前修改发送的内容tensor 在dist.irecv()之后读取接受的内容tensor req.wait() 可以保证这次通信顺利完成，因此我们可以在req.wait()之后再进行修改和读取。\nOne-to-many One-to-many 情形下，可以分为两种：scatter 和 broadcast\nscatter的作用是将一个process的数据均分并散布到其他process。broadcast的作用是将一个process的数据广播到其他process。两者不同的地方在于其他process获取到的是全量数据(copy)还是部分数据(slice)，其示意图如下所示\nscatter 测试代码：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 # scatter.py def example_scatter(): if dist.get_rank() == 0: scatter_list = [ torch.tensor([i+1] * 5, dtype=torch.float32).cuda() for i in range(dist.get_world_size()) ] print(f\u0026#34;Rank 0 scatter list: {scatter_list}\u0026#34;) else: scatter_list = None tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f\u0026#34;Before scatter on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.scatter(tensor, scatter_list, src=0) print(f\u0026#34;After scatter on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_broadcast() # run with # torchrun --nproc_per_node=4 broadcast.py 结果输出以下内容（输出内容有优化，后续不再说明）：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Rank 0 scatter list: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:0\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:0\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:0\u0026#39;) ] Before scatter on rank 2: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:2\u0026#39;) Before scatter on rank 1: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:1\u0026#39;) Before scatter on rank 3: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:3\u0026#39;) Before scatter on rank 0: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:0\u0026#39;) After scatter on rank 0: tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;) After scatter on rank 2: tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;) After scatter on rank 3: tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) After scatter on rank 1: tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;) broadcast 测试代码:\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 # broadcast.py def example_broadcast(): if dist.get_rank() == 0: tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda() else: tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f\u0026#34;Before broadcast on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.broadcast(tensor, src=0) print(f\u0026#34;After broadcast on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_broadcast() # run with # torchrun --nproc_per_node=3 broadcast.py 结果输出：\n1 2 3 4 5 6 7 Before broadcast on rank 1: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:1\u0026#39;) Before broadcast on rank 2: tensor([0., 0., 0., 0., 0.], device=\u0026#39;cuda:2\u0026#39;) Before broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:0\u0026#39;) After broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:0\u0026#39;) After broadcast on rank 1: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:1\u0026#39;) After broadcast on rank 2: tensor([1., 2., 3., 4., 5.], device=\u0026#39;cuda:2\u0026#39;) Many-to-one Many-to-one 情形下，也可以分为两种：gather 和 reduce, Gather对应one-to-many的scatter操作，负责将多个process的内容汇聚到一起，形成一个完整的向量。而reduce的操作则是通过一个函数 $f(\\cdot)$ 来把数据进行汇总，常见的函数有求和以及求平均，示意图如下所示\ngather 测试代码：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 # gather.py def example_gather(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() if dist.get_rank() == 0: gather_list = [ torch.zeros(5, dtype=torch.float32).cuda() for _ in range(dist.get_world_size()) ] print(f\u0026#34;Rank 0 gather list: {gather_list}\u0026#34;) else: gather_list = None print(f\u0026#34;Before gather on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.gather(tensor, gather_list, dst=0) if dist.get_rank() == 0: print(f\u0026#34;After gather on rank {dist.get_rank()}: {gather_list}\u0026#34;) init_process() example_gather() # run with # torchrun --nproc_per_node=4 gather.py 结果输出：\n1 2 3 4 5 6 7 8 9 10 11 Before gather on rank 3: tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) Before gather on rank 2: tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;) Before gather on rank 0: tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;) Before gather on rank 1: tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;) After gather on rank 0: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:0\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:0\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:0\u0026#39;) ] reduce 测试代码:\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 # example_reduce.py def example_reduce(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() print(f\u0026#34;Before reduce on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM) if dist.get_rank() == 0: print(f\u0026#34;After reduce on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_reduce() # run with # torchrun --nproc_per_node=3 example_reduce.py 这里我们使用求和dist.ReduceOp.SUM作为我们的汇总操作，Pytorch还支持其他的reduce operations. 结果输出以下内容：\n1 2 3 4 5 6 Before reduce on rank 2: tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;) Before reduce on rank 3: tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) Before reduce on rank 0: tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;) Before reduce on rank 1: tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;) After reduce on rank 0: tensor([10., 10., 10., 10., 10.], device=\u0026#39;cuda:0\u0026#39;) Many-to-many Many-to-many 情形下的两种通信方式为：All-Reduce 和 All-Gather，分别是reduce和gather的升级版，all-reduce对所有process都执行一次reduce操作，而all-gather则对所有process执行一次gather操作，其示意图如下所示\nall-gather 测试代码：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 # example_all_gather.py def example_all_gather(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() gather_list = [ torch.zeros(5, dtype=torch.float32).cuda() for _ in range(dist.get_world_size()) ] print(f\u0026#34;Before all gather on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.all_gather(gather_list, tensor) print(f\u0026#34;After all gather on rank {dist.get_rank()}: {gather_list}\u0026#34;) init_process() example_all_gather() # run with # torchrun --nproc_per_node=3 example_all_gather.py 测试输出结果：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 Before all gather on rank 2: tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;) Before all gather on rank 3: tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) Before all gather on rank 0: tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;) Before all gather on rank 1: tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;) After all gather on rank 0: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:0\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:0\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:0\u0026#39;)] After all gather on rank 2: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:2\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:2\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:2\u0026#39;) ] After all gather on rank 3: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:3\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:3\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:3\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) ] After all gather on rank 1: [ tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:1\u0026#39;), tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;), tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:1\u0026#39;), tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:1\u0026#39;) ] all-reduce 测试代码：\n1 2 3 4 5 6 7 8 9 10 11 12 # example_all_reduce.py def example_all_reduce(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() print(f\u0026#34;Before all reduce on rank {dist.get_rank()}: {tensor}\u0026#34;) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) print(f\u0026#34;After all reduce on rank {dist.get_rank()}: {tensor}\u0026#34;) init_process() example_all_reduce() # run with # torchrun --nproc_per_node=3 example_all_reduce.py 测试输出结果\n1 2 3 4 5 6 7 8 Before all reduce on rank 1: tensor([2., 2., 2., 2., 2.], device=\u0026#39;cuda:1\u0026#39;) Before all reduce on rank 0: tensor([1., 1., 1., 1., 1.], device=\u0026#39;cuda:0\u0026#39;) Before all reduce on rank 2: tensor([3., 3., 3., 3., 3.], device=\u0026#39;cuda:2\u0026#39;) Before all reduce on rank 3: tensor([4., 4., 4., 4., 4.], device=\u0026#39;cuda:3\u0026#39;) After all reduce on rank 0: tensor([10., 10., 10., 10., 10.], device=\u0026#39;cuda:0\u0026#39;) After all reduce on rank 2: tensor([10., 10., 10., 10., 10.], device=\u0026#39;cuda:2\u0026#39;) After all reduce on rank 3: tensor([10., 10., 10., 10., 10.], device=\u0026#39;cuda:3\u0026#39;) After all reduce on rank 1: tensor([10., 10., 10., 10., 10.], device=\u0026#39;cuda:1\u0026#39;) Barrier 除了之前这些传输数据的方式之外，我们还有Barrier，用于在所有process之间进行同步。Barrier会确保所有的process在同一时间点完成某些操作。其流程为，先让每个process完成各自的任务，然后当process到达barrier时，process会通知系统自己已到达。最后当所有process都到达barrier之后，阻塞会解除，所有process继续执行下一步操作。\nbarrier 测试代码\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 # example_barrier.py def example_barrier(): import time rank = dist.get_rank() t_start = time.time() print(f\u0026#34;Rank {rank} sleeps {rank} seconds\u0026#34;) time.sleep(rank) dist.barrier() print(f\u0026#34;Rank {rank} is done at {time.time() - t_start:.4f} seconds\u0026#34;) init_process() example_barrier() # run with # torchrun --nproc_per_node=3 example_barrier.py 结果输出\n1 2 3 4 5 6 7 8 9 Rank 2 sleeps 2 seconds Rank 0 sleeps 0 seconds Rank 1 sleeps 1 seconds Rank 3 sleeps 3 seconds Rank 3 is done at 3.3046 seconds Rank 1 is done at 3.3229 seconds Rank 2 is done at 3.8437 seconds Rank 0 is done at 3.6613 seconds 可以看到，四个process的到达时间都在3s左右，这是因为rank 3需要3s才能完成当前任务\nAdvanced 除了前面的通信方式之外，还有 Reduce-Scatter和Ring All-Reduce，这两个通信方式等我们学习ZeRO的时候再一并讲解。\nReference Colossal-AI LLaMA 4 blog Qwen3 blog Pytorch tutorial nanotron/ultrascale-playbook ","date":"2025-05-12T10:15:17+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/distributed-training--basic/","title":"Distributed training--Basic"},{"content":"介绍 Meta在2025年4月10号发布了LLaMA4系列，包含三个模型：Llama 4 Scout, Llama 4 Maverick 以及Llama 4 Behemoth, 三个模型都基于MoE架构，且支持多模态\nModel Layers Heads (Q / KV) Context Length #Parameters (activated/total) #Experts (activated/total) #Tokens LLaMA 4 Behemoth - - - 288B / 2T (1 shared + 1 routed) / 16 - LLaMA 4 Maverick 48 40/8 1M 17B / 109B (1 shared + 1 routed) / 128 ~22T LLaMA 4 Scout 48 40/8 10M 17B / 400B (1 shared + 1 routed) / 16 ~40T 训练数据截止到2024年8月。LLaMa4支持200多种语言，其中100多种语言的训练token数超过了1B\n亮点 原生多模态 LLaMA 4是一个原生多模态架构 超长上下文 LLaMA 4的上下文超过了1M iRoPE 通过交替dense和MoE MLP来提高整体推理效率 基于MetaCLIP的vision encoder MetaP 使用MetaP来调整超参数 FP8精度训练 Pre-training LLaMA 4 仍然是一个基于transformer的架构，但是引入了MoE，其示意图如下所示\nMoE架构中包含1个shared expert以及1个routed expert. 并且，与其他LLM不同，LLaMA 4使用了一个交替MLP个MoE的架构，即iRoPE，即特定的transformer layer是MoE架构，其余的是MLP架构，其核心代码如下：\n1 2 3 4 5 self.is_moe_layer = layer_idx in config.moe_layers if self.is_moe_layer: # the 128E model interleaves dense / sparse self.feed_forward = Llama4TextMoe(config) else: self.feed_forward = Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp) early fusion. LLaMA 4称其一个原生多模态大模型，但是其架构仍然是 Vision Encoder-MLP-LLM 的形式，其不同点在于patch embedding没有使用convolution, 而是使用 nn.Unfold直接进行展平，然后使用一个线性层与vision encoder进行对齐。代码如下\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 class Llama4UnfoldConvolution(nn.Module): def __init__(self, config): super().__init__() kernel_size = config.patch_size if isinstance(kernel_size, int): kernel_size = (kernel_size, kernel_size) self.unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=config.patch_size) self.linear = nn.Linear( config.num_channels * kernel_size[0] * kernel_size[1], config.hidden_size, bias=False, ) def forward(self, hidden_states: torch.Tensor) -\u0026gt; torch.Tensor: hidden_states = self.unfold(hidden_states) hidden_states = hidden_states.permute(0, 2, 1) hidden_states = self.linear(hidden_states) return hidden_states 其他训练优化技巧如下：\nMetaP：用于选择超参数 FP8 precision：与 DeepSeek-V3 一样，使用FP8精度进行训练 mid-training：在预训练阶段之后，额外增加了一个训练阶段，来提高模型的长上下文等关键能力 Post-training post-training包括三个阶段：\nSFT online RL DPO 作者发现SFT和DPO会限制模型的探索能力，特别是在math, coding等domain。为了解决这个问题，作者使用LlaMA对问题进行难度分级，然后移除了50%的简单数据。\n在online RL阶段，作者设计了一个continuous online RL策略，让模型在训练和筛选问题两种模式之间进行切换，以平衡效率和准确率。\nDPO的目的是为了提升模型输出的质量\n评测 benchmark LLaMA 4 Maverick LLaMA 4 Maverick LLaMA 4 Scout Gemeni 2.0 Flash GPT-4o MMMU 76.1 73.4 69.4 71.7 69.1 Math Vista 73.7 70.7 73.1 63.8 ChartQA - 90.0 88.8 88.3 85.7 DocVQA - 94.4 94.4 - 92.8 LiveCodeBench 49.4 43.4 32.8 34.5 32.3 MMLU Pro 82.2 80.5 74.3 77.6 - GPQA Diamond 73.7 69.8 57.2 60.1 53.6 结论 LLaMA 4 采用了MoE架构，是一个原生的多模态大模型系列。在架构上，与DeepSeek-MoE, aria和OLMoE不同，LLaMA4并没有增加expert granularity，OLMoE分析认为，增加granularity可以提高模型的flexibility， 下面总结了一下相关模型的参数\nModel Layers Heads (Q / KV) Context Length #Parameters (activated/total) #Experts (activated/total) #Tokens DeepSeek-MoE(144.6B) 62 32/32 2048 22.2B/144.6B (1 shared + 7 routed)/64 245B DeepSeek-V3 61 128(MLA) 128K 37B/671B (1 shared + 8 routed)/257 14.8T Aria 28 20/20 64K 3.5B/24.9B (2 shared+ 6 routed)/66 6.4T(text) OLMoE 16 16/16 4096 1.3B/6.9B 8/64 5T LLaMA 4 Maverick 48 40/8 1M 17B / 109B (1 shared + 1 routed) / 128 ~22T LLaMA 4 Scout 48 40/8 10M 17B / 400B (1 shared + 1 routed) / 16 ~40T Reference LLaMA 4 Model Card LLaMA 4 Blog ","date":"2025-04-30T10:44:19+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-llama4-blog/","title":"Notes on LLaMA4 blog"},{"content":"介绍 Qwen3发布，包含两种架构的模型，每种架构均包含对应的base model和post-trained model.\nMoE架构：上下文长度为128K，128个专家，每个token由8个专家负责处理\nQwen3-235B-A22B, 总参数235B，激活参数22B， Qwen3-30B-A3B, 总参数30B，激活参数3B Models Layers Heads (Q / KV) # Experts (Total / Activated) Context Length Qwen3-30B-A3B 48 32 / 4 128 / 8 128K Qwen3-235B-A22B 94 64 / 4 128 / 8 128K dense架构: Qwen3-32B, Qwen3-14B, Qwen3-8B, Qwen3-4B, Qwen3-1.7B, and Qwen3-0.6B\nModels Layers Heads (Q / KV) Tie Embedding Context Length Qwen3-0.6B 28 16 / 8 Yes 32K Qwen3-1.7B 28 16 / 8 Yes 32K Qwen3-4B 36 32 / 8 Yes 32K Qwen3-8B 36 32 / 8 No 128K Qwen3-14B 40 40 / 8 No 128K Qwen3-32B 64 64 / 8 No 128K 亮点 Hybrid Thinking modes qwen3支持两种思考模式： thinking mode 和 non-thinking mode，前者用于解决复杂的问题，后者用于解决简单的问题 multilingual support qwen3支持119中语言和方言 Improved agentic capabilities 提升了qwen3的coding和agentic能力，并支持MCP 训练 Pre-training Qwen3使用了36T token进行训练 （与之相比，Qwen2.5使用方的token数量为18T），数据来源于互联网和PDF，作者使用Qwen2.5-VL来提取内容，然后使用Qwen2.5来提升内容质量。作者还基于Qwen2.5-Math和Qwen2.5-Coder来合成数学以及代码数据\n训练包含三个stage：\n上下文长度为4K tokens，训练数据为30T tokens, 目标是让模型掌握初步的语言能力和知识 上下文长度为4K tokens，训练数据为5T tokens,这部分数据主要是knowledge intensive的数据，比如STEM, coding和reasoning等 上下文长度扩展到32K tokens，训练数据主要是高质量长上下文的数据 Post-training Post-training包含四个阶段，如下图所示 Stage 1, Long CoT code start: 作者基于math, coding, logical reasoning, STEM等domain的Long CoT数据来微调模型，让模型拥有初步的推理能力，这和Kimi-VL是类似的 Stage 2, reasoning-based RL: 作者使用rule-based rewards来提供奖励，然后使用RL来训练模型，经一部提高模型的exploration和exploitation能力 Stage 3, thinking mode fusion: 作者混合了一部分instruction-following和Long CoT数据来提升模型的non-thinking能力，这样可以让模型在两种思考模式之间切换 Stage 4,general RL： 作者使用RL在20多个general-domain任务上进一步提高模型的通用能力，包括instruction following, format following以及agent capability等 Future work 作者希望在未来能够在模型架构和训练方式上进行提升，包括：scaling data, increasing model size, extending context length, broadening modalities, advancing RL with environmental feedback for long-horizon reasoning.\n结论 与Gemini2.5 pro，Kimi-VL等reaosning model不同，qwen3可以在快思考和慢思考之间进行转换。感觉未来有两个趋势，一个是如何在快思考和慢思考之间进行切换，切换的逻辑是什么？第二个就是qwen3以及qwen2.5-vl都在强调的agent能力，也就是我们不仅仅是在做一个LLM，而是逐步延伸到了agent这个层面。\n参考文献 blog ","date":"2025-04-29T11:23:04+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen3-blog/","title":"Notes on Qwen3 blog"},{"content":"介绍 简单总结一下已有的介绍了训练数据配比的多模态大模型，方便后续使用。 根据之前看的一些论文进行总结，如有缺漏，欢迎批评指正。\n相关工作 LLaVA LLaVA 1.5 SFT数据配比如下：\nLLaVA-OneVision LLaVA OneVision 的数据集统计在表16里，这里总结一下数据配比(总量为3.15M)\nCategory Ratio general vqa 36.1 Doc/Chart/Screen 20.6 Math/Reasoning 20.1 general OCR 8.9 text only 14.3 Apollo Apollo 是Meta发布的一个视频理解多模态大模型，其数据集没有开源，论文中给出了其训练数据配比 Cambiran-1 Cambiran-1 是纽约大学提出了一个多模态大模型系列，论文发表在了 NeurIPS 2024(Oral) 上，作者给出了 Cambrain-10M 和 Cambrain-7M 两个数据集，Cambrain-7M 的数据分布如下\nIdefics Idefics系列(1/2/3)是huggingface提出的视觉多模态大模型系列，在 Idefics2 中，作者构建了The Cauldron数据集，其数据配比在表14里面。总结如下：\nCategory Ratio general vqa 11.02 captioning 5.14 OCR 17.47 chart/figures 14.05 table 11.3 reasoning 10.32 textbook 1.58 difference 2.38 screenshot2code 0.31 text only 26.41 在 Idefics3 中，作者基于The Cauldron进行了扩展，最终数据集的比例如下：\nMolmo Molmo 是Allen AI发布的一个多模态大模型，其SFT数据配比如下\nEagle 2/2.5 Eagle (1/2/2.5) 是NVLab提出了系列多模态大模型，Eagle 2 给出了 stage 1.5 和 stage 2的数据配比\nEagle 2.5 在 Eagle 2的基础上加入了一些 long context 数据，其数据列表在表11里\nSmolVLM SmolVLM 是huggingface开发的一款轻量化视觉多模态大模型，论文中的数据配比如下：\nMM1/1.5 MM1 通过实验确定了训练数据的配比：\nCategory Ratio interleaved image-text 45 imagetext 45 text only 10 MM1.5 的SFT数据配比如下： InternVL InternVL2.5 在论文里总结了其使用的pretraining数据和SFT数据，但是没有具体的数据配比，请参考论文的表4和表5\nSFT数据配比\nCategory Ratio single-image 45.92% multi-image 9.37% video 39.79% pure-text 4.92% MiniCPM V MiniCPM V 是 OpenBMB发布的轻量化多模态大模型，其在表1和表2列出了ptraining和SFT数据的具体量级和类别。\nFlash-VL Category Ratio Special Enhancement 4% Text 21% Caption 4% Chart 16% Math 11% OCR 3% Code 8% General 33% 具体数据参见原论文\n参考文献 LLaVA 1.5 LLaVA OneVision Apollo Cambiran-1 Idefics2 Idefics3 Molmo Eagle 2 Eagle 2.5 SmolVLM MM1 MM1.5 InternVL2.5 MiniCPM V ","date":"2025-04-25T10:25:48+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/data-mixture-in-mllm/","title":"Data mixture in MLLM"},{"content":"上大学以及读研的时候，从没生过病，女朋友还说我身体素质强大。工作之后，先是新冠，然后轻度脂肪肝，脂肪肝，今年体检又有颈椎曲度变直，腰椎间盘突出一系列问题。说实话，这些毛病说大也不算大，工作久了都有类似的问题。但是前两天因为睡觉着凉感冒了两天，就忽然觉得年少时以为身体强健，不过是身体在替你硬抗。现在工作之后，久坐，不健康饮食，熬夜等坏习惯就开始了反击。都说人类理想的寿命不过25岁，25岁你应该已经完成结婚生子，可以安享晚年了。而也是在这个年龄段，身体开始摆烂，让你自己硬抗。现在来看，果然如此。\n有时候，还在那里想，人要完成一个什么样的目标才算是成功呢，是发很多顶会成为学术大佬，还是赚很多钱，实现人生自由。现在来看，身体健康，家庭美满，有一份说得过去的工作就已经很好了，不要好高骛远。人可以向前看，向前走，但是不能不看自己的身体状况。\n想起王立群老师的话，“人的一生就像一场马拉松，只要你在路上，就还有机会”。但是现在来看，看着身边那么多优秀的人，很容易就把马拉松跑成百米赛跑，没跑过人家不说，百米过了，还浑身酸痛。\n又想起高职务的那句话，“有时候想想这官当多大才叫大啊”。现在看，人也是这样啊，赚多少的钱才算人生自由，发多少的论文才算成功。向上看，总是会给自己无形的压力，让自己陷入无尽的焦虑之中。\n我觉得我已经算是比较豁达的人了，只会跟自己比，就算这样也还是会有一些焦虑。我的焦虑主要是来自于自己啥都没做，就已经周三了。周日晚上着凉，周一周二脑子就完全宕机，一周的计划直接泡汤一半。这种计划被打乱的挫败感对于T人来说确实难以接受。而细细一想，学生时代对于时间真的没什么概念，每天就瞎搞，没有目标，反而时间就过得慢。现在则是，一低头一抬头可能两个小时就过去了。难说这是年龄大了之后，对时间的迟钝性，还是工作之后，磨平了人心中的激情。\n最近开始了健身，希望能花半年时间降低体脂，我越来越感觉到身体健康的重要性，很多东西失去了才真正懂得拥有是多么难得，感冒鼻塞了才知道呼吸到空气不是那么一件简单的事情。摔伤之后才知道洗澡也可能是奢望。\n在明朝那些事的结尾，当年明月没有提到书中的大人物，而是以徐霞客收尾，“所谓百年功名、千秋霸业、万古流芳，与一件事情相比，其实算不了什么。这件事情就是——用你喜欢的方式度过一生。”希望大家都有健康的体魄，毕竟身体是革命的本钱。与诸君共勉。\n","date":"2025-04-23T13:24:02+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/%E9%9A%8F%E7%AC%94-%E8%BA%AB%E4%BD%93%E5%81%A5%E5%BA%B7/","title":"随笔-身体健康"},{"content":"Abstract 字节Seed团队提出了 Value-based Augmented Proximal Policy Optimization (VAPO), 用于提高 reasoning model 的表现。 VAPO 通过集成 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。\nIntroduction 现有的RL 训练方法可以分为 value-free 和 value-based 两大类。 其中 value-free 方法不需要使用 value model, 比如 GRPO 和 GRPO 的变体 DAPO, 这类方法通过多次采样，然后使用 leave-one-out estimate 来代替 value model. 这类方法的优点是不需要训练value model, 但是缺点是在复杂的任务中表现不是很稳定。\n另一方面，value-based 方法需要训练一个 value model, 比如 VC-PPO, 这类方法的优点是：\n提供更细粒度的奖励信号 提供lower-varaince value estimation, 从而提高训练的稳定性 拥有更好的泛化能力 但是，value-based 方法在训练过程中存在一些问题：\n训练一个low-bias 的 value model 比较困难， 尤其是在long trajectory 上，因为 bias 会随着 trajectory 的长度增加而增加 在heterogeneous sequence lengths during training 中表现不佳，对于短文本和长文本，我们需要考虑 bias-variance 的trade-off 在sparse reward signal 中表现不佳 为了解决这些问题，字节Seed团队提出了 VAPO, 一个基于value-based 的 RL 训练方法，VAPO 通过结合 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。\nPreliminary Preliminary包括 token-level MDP, RLHF, PPO 三个部分，这部分请参考 VC-PPO. 这里不做重复。\nVAPO 作者针对value-based 方法在训练过程中存在的三个问题，在VAPO中分别进行了解决。\nMitigating Value Model Bias over Long Sequences 在 VC-PPO 中，作者提到了 value model 和 reward model 的不一致性，这种不一致性会导致 bias, 尤其是在long sequences上。 在VAPO中，作者就直接使用了 VC-PPO的做法，包括value pretraining 和 decoupled GAE 来解决这个问题。\nManaging Heterogeneous Sequence Lengths during Training 针对heterogeneous sequence的问题，作者提出了 Length-Adaptive GAE. 在VC-PPO中， $\\lambda_{\\mathrm{policy}}$ 被设置为 $0.95$. 但是当 sequence 非常长时， TD-error会迅速下降，导致GAE被一部分TD-error所主导，从而不利于模型的训练。\n为了解决这个问题，作者将 $\\lambda_{\\mathrm{policy}}$ 与 sequence 的长度 $\\ell$ 联系起来，具体来说， 两者的关系如下：\n$$ \\sum_{t=0}^{\\infty}\\lambda_{\\mathrm{policy}}^t = \\frac{1}{1-\\lambda_{\\mathrm{policy}}} := \\alpha\\ell $$其中 $\\alpha$ 是一个超参数，用来控制 bias-variance 的trade-off. 给定 $\\ell$, $\\lambda_{\\mathrm{policy}}$ 可以被计算为：\n$$ \\lambda_{\\mathrm{policy}} = 1 - \\frac{1}{\\alpha\\ell} $$同时，为了平衡短文本和长文本的贡献，基于 DAPO, 作者构建了 token-level policy gradient loss， 其具体形式如下：\n$$ \\mathcal{L}_{\\mathrm{PPO}}(\\theta) = \\frac{1}{\\sum_{t=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\min\\left(r_{i,t}(\\theta)\\hat{A}_{i,t},\\mathrm{clip}\\left(r_{i,t}(\\theta), 1-\\epsilon, 1+\\epsilon\\right)\\hat{A}_{i,t}\\right) $$Dealing with Sparsity of Reward Signal in Verifier-based Tasks 与DAPO一致，为了解决reward signal的稀疏性问题，作者提出了Clip-Higher, 来让更小概率的输出也能获得较大的更新概率。其更新公式如下：\n$$ \\mathcal{L}_{\\mathrm{PPO}}(\\theta) = \\frac{1}{\\sum_{t=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\min\\left(r_{i,t}(\\theta)\\hat{A}_{i,t},\\mathrm{clip}\\left(r_{i,t}(\\theta), 1-\\epsilon_{low}, 1+\\epsilon_{high}\\right)\\hat{A}_{i,t}\\right) $$Clip-Higher的介绍见DAPO\n然后，作者还将next-token prediction loss 和 PPO loss 结合起来，来降低奖励的稀疏程度。\n$$ \\mathcal{L}_{\\mathrm{NTP}}(\\theta) = -\\frac{1}{N}\\sum_{o_i\\in\\mathcal{T}}\\sum_{t=1}^{|o_i|}\\log \\pi_{\\theta}(a_{i,t}|s_{i,t}) $$其中 $\\mathcal{T}$ 是正确答案的集合。 最终的loss为：\n$$ \\mathcal{L}_{\\mathrm{VAPO}}(\\theta) = \\mathcal{L}_{\\mathrm{PPO}}(\\theta) +\\mu \\mathcal{L}_{\\mathrm{NTP}}(\\theta) $$其中 $\\mu$ 是一个超参数，用来平衡PPO loss和NTP loss。\nExperiments 模型使用Qwen-32B来进行训练, 大部分细节与VC-PPO和DAPO一致，这里不再赘述。最后与DAPO以及R1的对比结果如下：\nAblation study 针对本文使用的模块，作者进行了消融实验，结果如下：\n从实验结果可以看到：\nvalue pretraining 和 decoupled GAE 可以显著提高模型的表现 clip-higer 可以提升模型的探索能力 length-adaptive GAE 可以平衡模型在短文本和长文本上的表现 Training Dynamics 与DAPO类似，作者也分析了VAPO的训练动态，结果如下：\n从上面三张图可以看到：\nVAPO相比于DAPO来说，其训练更加稳定 从response length来看，VAPO的response length更长，说明VAPO的length scaling更强 从reward score来看，VAPO的reward score更高，说明VAPO的reward signal提供的指导信息更多 从generation entropy来看，在训练末期，VAPO的generation entropy更低，说明VAPO的生成多样性要更低一些，但这也说明了VAPO的生成更加稳定。 作者认为这个时候稳定性更重要。 Conclusion 作者在DAPO和VC-PPO的基础上，提出了VAPO，一个基于value-based 的 RL 训练方法，VAPO 通过结合 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。 后续，作者又提出了Seed-thiking-1.5的技术报告。可以说，这一系列论文的连贯性是非常高的。\nReference Arxiv VAPO: Efficient and Reliable Reinforcement Learning for Advanced Reasoning Tasks Notes onDAPO Notes on VC-PPO ","date":"2025-04-17T09:41:51+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-vapo/","title":"Notes on VAPO"},{"content":"Abstract 字节Seed团队提出了 Value-Calibrated PPO (VC-PPO), 用于解决PPO的value initialization bias 以及 reward signal decay 问题. 具体来讲：\nVC-PPO增加了 value pretraining 来解决 value initialization bias 的问题 VC-PPO分离了 actor 和 critic 的GAE的计算，避免了 reward signal decay 的问题 Introduction 已有的reasoning model的训练方法，主要包括两个个stage:\nSFT: 这个阶段主要使用了一些标注好的long CoT数据，初步激活模型的reasoning能力(参考KImi-VL) RL: 这个阶段使用收集到的数据使用RL算法进行训练，任务包括math, code, logic reasoning等 已有PPO算法在处理Long CoT任务时，存在的问题在 DAPO 中已经介绍过了，GRPO的解决方式为 使用leave-one-out estimate来替换value model. 但是GRPO相比于PPO能够提供token级别的奖励来说，只能提供response level的奖励，因此限制了模型的性能。\nPreliminary Preliminary包括MDP, RLHF, PPO三个部分，RLHF和PPO我们在 DAPO 中已经介绍过了，这里不再赘述。\nToken-level MDP 给定prompt $x$ 和 response $y$, 我们可以将 $x$ 和 $y$ 分解为 token 序列，比如 $y = y_1, y_2, \\cdots, y_n$, 其中 $y_i\\in\\mathcal{A}$, $\\mathcal{A}$ 是我们的词表。\n我们将 token-level MDP定义为：\n$$ \\mathcal{M} = \\langle \\mathcal{S}, \\mathcal{A}, P, r, d_0, \\omega \\rangle $$其中：\n$\\mathcal{S}$ 是状态空间，表示当前的token序列， $t$时刻的状态可以表示为 $s_t = (x, y_1, y_2, \\cdots, y_t)$ $\\mathcal{A}$ 是动作空间，表示下一个token， $t$时刻的动作可以表示为 $a_t = y_{t+1}\\in\\mathcal{A}$ $P$ 是状态转移概率，表示在 $t$时刻，从状态 $s_t$ 转移到状态 $s_{t+1}$ 的概率 $r$ 是奖励函数，表示在 $t$时刻，从状态 $s_t$ 采取动作 $a_t$ 转移到状态 $s_{t+1}$ 的奖励 $d_0$ 是初始状态分布，表示初始状态 $s_0$ 的概率分布 $\\omega$ 是终止状态分布，表示终止状态 $s_n$ 的概率分布，通常表示 \u0026lt;eos\u0026gt; token 方法 首先作者分析了一下为什么 PPO 在处理 long CoT 任务时效果不佳。\nPPO的传统设置为：\n将GAE的参数 $\\lambda$ 设置为 0.95 使用一个 reward model 来初始化 value model 一方面，作者认为，$\\lambda=0.95$ 是为了减少模型在 Mujoco 以及 Atari 等环境中的variance，但是对于Long CoT任务，这个设置会导致模型缺乏足够的探索能力。\n另一方面，作者认为 reward model 和 value model 虽然都是提供关于 response 的信息，但是他们之间还是存在一些差距的。作者给出了使用PPO来进行 long CoT 任务相关的实验结果\n可以看到，随着训练的进行，模型的context length以及在AIME benchmark上的表现都出现了下降。\n在本文中，作者主要是在 verifiable tasks 上进行实验，答案的正确性与 response length 长度关系不大，因此, response length 可以反应模型的训练情况。 作者绘制了训练过程中， value 和 advantage 与 token position 的关系图\n从上图可以看到，value 和 advantage 更倾向与给最初的 token 更高的 bias, 作者认为出现这个情况的原因是 value model 和 reward model 的目标并不匹配。 reward model 的目标是给 $\\omega$ 也就是 \u0026lt;eos\u0026gt; token reward, 对于一开始的 token, reward model会给出比较低的奖励；而 value model 的目标是给整个 response 一个奖励，因此，value model 会倾向于给最初的 token 更高的奖励。我们对 GAE 进行改写得到：\n$$ \\hat{A}_t = \\sum_{i=t}^{T-t-1} \\lambda^{i} \\left( r_{t+i} + V(s_{t+i+1}) - V(s_{t+i}) \\right) $$从上式可以看到，一开始 $r_{t+i}$ 的值会比较小，而 $V(s_{t+i})$ 的变化比较大，因此，一开始的 token 的 advantage 会比较大，这个 bias 会持续影响整个 trajectory。\n为了解决这个问题，作者提出了 value-pretraining, 也就是对value model 进行离线与训练，直到其收敛到一个具体的 policy 上。具体训练步骤为：\n基于一个policy， 如 $\\pi_{\\mathrm{sft}}$ 进行采样，然后更新value model ($\\lambda=1.0$) 基于收集到的数据训练 value model, 直到 value loss 或者 explain variance 收敛 接下来，在训练 value model 时，我们还需要考虑 variance reduction 的问题。作者首先改写了 GAE 的公式：\n$$ \\hat{A}_t = \\begin{cases} \\sum_{i=0}^{T-t-1} \\lambda^{i} \\left( r_{t+i} + V(s_{t+i+1}) - V(s_{t+i}) \\right)+V(s_t) \u0026 \\text{if } \\lambda \u003c 1.0 \\\\ \\sum_{i=0}^{T-t-1} r_{t+i} \u0026 \\text{if } \\lambda=1.0 \\end{cases} $$可以看到，当 $\\lambda \u003c 1.0$ 并且 $T-t-1$ 比较大时，\u0026lt;eos\u0026gt; token 的reward就非常接近于0了，作者通过实验验证了这一点。\n可以看到，当我们降低 $\\lambda$ 时，reward signal 的稀疏性会显著增加，从而提高了模型的训练难度。\n但是，我们又不能不使用 variance reduction, 因为这会导致训练的不稳定性。 作者从 TD error 的角度分析了 variance:\n$$ \\begin{aligned} \\mathrm{Var}[A_{t}^{\\lambda}] \u0026= \\mathrm{Var}\\left[\\sum_{i=0}^{T-t-1} \\lambda^{i}\\delta_{t+i}\\right] \\\\ \u0026= \\sum_{i=1}^{T-t-1} \\lambda^{2i} \\mathrm{Var}[\\delta_{t+i}] + 2\\sum_{i=1}^{T-t-1}\\sum_{j=0}^{i-1} \\lambda^{i+j} \\mathrm{Cov}[\\delta_{t+i}, \\delta_{t+j}] \\end{aligned} $$因为 $\\lambda\\in[0,1]$, 因此未来的 TD error 会衰减的更快，因此就降低了 advantage 的 variance.\n那么如何在降低variance的同时，又不至于使得 reward signal 过于稀疏呢？作者的解决方案是分离GAE的计算，也就是将 GAE 的计算分为两部分：\n$$ G_{t:t+h} = \\begin{cases} \\sum_{i=0}^{h-1} r_{t+i} + \\bar{V}(s_{t+h}) \u0026 \\text{if } t+h","date":"2025-04-14T17:36:15+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-vc-ppo/","title":"Notes on VC-PPO"},{"content":"Introduction 字节Seed团队和清华合作提出了DAPO，一种全开源的，基于PPO的强化学习方法，一用于提升LLM的reasoning能力。\nPreliminary Preliminary包括PPO，GRPO还有KL divergence\nPPO PPO的训练目标为：\n$$ \\mathcal{J}_{\\mathrm{PPO}}(\\theta) = \\mathbb{E}_{(q,a)\\sim\\mathcal{D},o_{\\leq t}\\sim \\pi_{\\theta_{old}}(\\cdot\\mid q)}\\left[ \\min\\left(r_t(\\theta)\\hat{A}_t,\\mathrm{clip}\\left(r_t(\\theta), 1-\\epsilon, 1+\\epsilon\\right)\\hat{A}_t\\right) \\right] $$其中\n$$ r_t(\\theta) = \\frac{\\pi_{\\theta}(o_t\\mid q, o_{\u003c t})}{\\pi_{\\theta_{old}}(o_t\\mid q, o_{\u003c t})} $$$(q,a)$ 是从数据集 $\\mathcal{D}$ 采样的QA pair，$\\epsilon\u003e0$ 是一个超参数，$\\hat{A}_t$ 是 $t$时刻的优势估计 (advantage estimator). 给定 value function $V$ 以及 reward function $R$, $\\hat{A}_t$ 通过计算GAE得到：\n$$ \\hat{A}_t^{\\mathrm{GAE}(\\gamma, \\lambda)}=\\sum_{k=0}^{\\infty}(\\gamma\\lambda)^k\\delta_{t+k} $$其中\n$$ \\delta_k = R_k + \\gamma V(s_{k+1})-V(s_k),\\quad 0\\leq \\gamma,\\lambda\\leq 1 $$GRPO 相比于PPO，GRPO不依赖于value function, 因此不需要使用reward model. GRPO通过一组输出来估计value $V(s)$, 然后进行更新。具体来说，给定 QA pair $(q,a)$, 我们从$\\pi_{\\theta_{old}}$中采样$G$个输出 $\\{o_i\\}_{i=1}^G$, 接下来我们基于reward $\\{R_i\\}_{i=1}^G$ 使用如下表达式来估计group-level reward:\n$$ \\hat{A}_{i,t} = \\frac{r_i - \\mathrm{mean}(\\{R_i\\}_{i=1}^G)}{\\mathrm{std}(\\{R_i\\}_{i=1}^G)} $$最后，GRPO的训练目标与PPO类似，只不过将 $\\hat{A}_t$ 替换为 $\\hat{A}_{i,t}$, 然后在分组上进行了归一化：\n$$ \\mathcal{J}_{\\mathrm{GRPO}}(\\theta) = \\mathbb{E}_{(q,a)\\sim\\mathcal{D},\\{o_i\\}_{i=1}^G\\sim \\pi_{\\theta_{old}}(\\cdot\\mid q)}\\left[ \\frac{1}{G}\\sum_{i=1}^G\\frac{1}{|o_i|}\\sum_{t=1}^{|o_i|}\\min\\left(r_{i,t}(\\theta)\\hat{A}_{i,t},\\mathrm{clip}\\left(r_{i,t}(\\theta), 1-\\epsilon, 1+\\epsilon\\right)\\hat{A}_{i,t}\\right) \\right] $$其中，\n$$ r_{i,t}(\\theta) = \\frac{\\pi_{\\theta}(o_{i,t}\\mid q, o_{i,\u003c t})}{\\pi_{\\theta_{old}}(o_{i,t}\\mid q, o_{i,\u003c t})} $$KL divergence 在传统的RLHF框架中，我们再PPO的基础上，增加了一个KL divergence正则项，用于约束新旧策略的差异。具体来说，给定旧策略 $\\pi_{\\theta_{old}}$ 和新策略 $\\pi_{\\theta}$，我们实际上优化的损失函数为\n$$ \\mathcal{J}_{\\mathrm{RLHF}}(\\theta) = \\mathcal{J}_{\\mathrm{PPO}}(\\theta) - \\beta\\mathrm{KL}\\left(\\pi_{\\theta_{old}}(\\cdot\\mid q)\\|\\pi_{\\theta}(\\cdot\\mid q)\\right) $$其中，$\\beta$ 是一个超参数，用于平衡PPO损失和KL divergence损失。上面的PPO损失函数也可以改为GRPO损失函数。\n作者在本文中认为，reasoning model和RLHF的训练目标是不一样的，RLHF加上KL divergence正则项，会使得模型在推理时过于保守，偏向于explotition而reasoning model则需要exploration。因此，作者在本文中去掉了这一项。\nRule-based reward modeling 作者基于final accuracy来作为outcome reward, 避免模型训练出现reward hacking问题。reward function 如下：\n$$ R(\\hat{y},y) = \\begin{cases} 1, \u0026 \\text{if is\\_equivalent}(\\hat{y},y) \\\\ -1, \u0026 \\text{otherwise} \\end{cases} $$DAPO DAPO基于GRPO改进，其优化的目标函数为：\n$$ \\mathcal{J}_{\\mathrm{DAPO}}(\\theta) = \\mathbb{E}_{(q,a)\\sim\\mathcal{D},\\{o_i\\}_{i=1}^G\\sim \\pi_{\\theta_{old}}(\\cdot\\mid q)}\\left[ \\frac{1}{\\sum_{t=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\min\\left(r_{i,t}(\\theta)\\hat{A}_{i,t},\\mathrm{clip}\\left(r_{i,t}(\\theta), 1-\\epsilon_{low}, 1+\\epsilon_{high}\\right)\\hat{A}_{i,t}\\right) \\right] s.t. \\quad 0\u003c \\vert \\{o_i\\mid \\text{is\\_equivalent}(o_i,a)\\} \\vert \u003c G $$其中，$\\alpha$ 是一个超参数，用于约束输出数量, $r_{i,t}(\\theta)$ 和 $\\hat{A}_{i,t}$ 的定义与GRPO相同。\n接下来就是DAPO算法的几个关键点：\nClip-Higher 作者认为，PPO和GRPO存在entropy collapse问题，也就是policy的entropy会迅速的下降，导致最终每个group的输出基本都差不多，模型很难探索到新的知识。因此，作者提出了Clip-Higher策略，也就是说，让模型能够充分探索。\n作者举了一个例子，当 $\\epsilon=0.2$ 时，如果一个group的输出为 $[0.01, 0.9]$, 那么在clip之后，最大的更新幅度为 $[0.01, 0.9]\\times 0.2=[0.002, 0.18]$. 这对于概率比较小的输出来说，很难帮助模型提升。因此，作者修改了其阈值。进一步地，作者分离了clip的阈值，分别使用 $\\epsilon_{low}$ 和 $\\epsilon_{high}$ 来约束更新幅度。作者通过实验验证了这个例子，结果如下图所示：\n最后，加入Clip-Higher策略的DAPO的训练目标为，与GRPO的训练目标相比，我们使用 $\\epsilon_{low}$ 和 $\\epsilon_{high}$ 来约束更新幅度。\n$$ \\mathcal{J}_{\\mathrm{DAPO}}(\\theta) = \\mathbb{E}_{(q,a)\\sim\\mathcal{D},\\{o_i\\}_{i=1}^G\\sim \\pi_{\\theta_{old}}(\\cdot\\mid q)}\\left[ \\frac{1}{\\sum_{t=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\min\\left(r_{i,t}(\\theta)\\hat{A}_{i,t},\\mathrm{clip}\\left(r_{i,t}(\\theta), 1-\\epsilon_{low}, 1+\\epsilon_{high}\\right)\\hat{A}_{i,t}\\right) \\right] s.t. \\quad 0\u003c \\vert \\{o_i\\mid \\text{is\\_equivalent}(o_i,a)\\} \\vert \u003c G $$在实际训练中，作者使用了一个比较大的 $\\epsilon_{high}$ 来保证概率较小的token也能有较大的概率被采样到。 Clip-Higher的实验结果如下图所示：\nDynamic Sampling 作者认为，在GRPO中，如果一个group的输出全都是对的，那么其advantage会趋近于0，结果导致 policy 不会得到更新，这样就降低了采样效率。\n为了解决这个问题，坐着提出了over-sample以及filtering策略，来过滤accracy等于1或者0的prompts,也就是DAPO中的约束项：\n$$ s.t. \\quad 0\u003c \\vert \\{o_i\\mid \\text{is\\_equivalent}(o_i,a)\\} \\vert \u003c G $$通过增加这个约束项，我们可以保证每个group的输出不会趋同，从保证模型能够得到更新，以提高采样效率。实验结果如下图所示：\nToken-level policy gradient loss GRPO采用了一个sample-level的loss计算方式，也就是GRPO首先在每个sample中计算每个token的损失并进行平均，然后在不同的sample中在此进行平均。在这种方式下，每个sample对最终的loss贡献是相同的。这样就导致两个问题：\n对于高质量的long answer，通过在token层面进行平均，因为answe比较长，因此整体loss会比较低，也就会降低这个sample的贡献，导致模型很难学习到长答案。 一些非常长的sample，其往往包含gibberish或者repetitive tokens，这些sample对模型训练是有害的 因此，为了解决这个问题，作者提出了Token-level policy gradient loss，也就是我们改变了求和方式\n$$ \\frac{1}{G}\\sum_{i=1}^G\\frac{1}{|o_i|}\\sum_{t=1}^{|o_i|}\\left(\\cdot\\right)\\to \\frac{1}{\\sum_{t=1}^G|o_i|}\\sum_{i=1}^G\\sum_{t=1}^{|o_i|}\\left(\\cdot\\right) $$通过这种方式，我们让长回答的loss贡献更大，从而提高模型学习长回答的能力。另一方面，现在每个token对整体loss的贡献是相同的，一些关键token的loss也会被放大，从而提高模型的表现。\n实验结果如下图所示：\nOverlong reward shaping 与Kimi-k1.5类似，DAPO也增加了length penalty，用于惩罚过长的回答。作者首先使用了一个overlong filtering技巧，来mask掉truncated samples的loss，作者发现通过这个技巧，可以提升模型的训练稳定性以及表现，实验结果如下图所示\n作者还提出了soft overlong punishment， 用于reshape truncated samples的reward，表达式如下：\n$$ R_{length}(y) = \\begin{cases} 0, \u0026 \\text{if } |y|\\leq L_{\\max}-L_{cache} \\\\ \\frac{(L_{\\max}-L_{cache})-|y|}{L_{cache}}, \u0026 \\text{if } L_{\\max}-L_{cache}\u003c|y|\\leq L_{\\max} \\\\ -1, \u0026 \\text{if } |y|\u003eL_{\\max} \\end{cases} $$算法 DAPO的算法流程如下：\n实验 数据集 作者构建了DAPO-Math-17K数据集用于DAPO的训练，该数据集从AoPS得到，包含了17K prompts，每个prompt都对应一个整数作为答案。\n训练细节 作者以GRPO作为baseline， $G=16$, $L_{\\max}=16384$, $L_{cache}=4096$, $\\epsilon_{low}=0.2$, $\\epsilon_{high}=0.28$\n评估时，使用AIME作为benchmark，以 avg@32作为指标，temperature设置为1.0, topp设置为0.7。\n实验结果 DAPO与GRPO的对比如下图所示： Ablation study 作者探究了每一个部分对最终表现的影响，结果如下：\ntraing dynamics 作者还探究了mean response length, reward score, generation entropy以及mean probability随训练轮数的变化，其中：\nmean response length: 在一定程度上反应了模型训练的稳定性和表现 reward score: 反应模型的表现，作者认为，给定一个reliable reward signal， LLM可以很好的fit到训练数据集上，但是作者发现最终的reward与val score相关性比较大，很可能是因为模型过拟合了 generation entropy \u0026amp; mean probability: 代表了模型的探索能力，通过实验结果可以看到，DAPO初期的探索能力比较强，但是随着训练的进行，探索能力下降，利用能力增强。 结果如下：\n结论 作者基于GRPO提出了DAPO，一种全开源的，基于PPO的强化学习方法，用于提升LLM的reasoning能力。作者首先分析了GRPO损失函数存在的不足，然后进行了针对性改进，包括Clip-Higher, Dynamic Sampling, Token-level policy gradient loss以及Overlong reward shaping。作者通过实验验证了DAPO的有效性，并探究了DAPO的训练动态。\nReference DAPO: An Open-Source LLM Reinforcement Learning System at Scale ","date":"2025-04-09T21:40:33+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-dapo/","title":"Notes on DAPO"},{"content":"介绍 2025年3月26号，Qwen团队发布了Qwen2.5 omni，Qwen2.5 omni是一个多模态模型，支持文本、音频、视频、图像等多个模态的输入，支持文本，音频的输出。 作者提出了TMRoPE来对齐不同的模态，然后还是用了Thinker-Talker的架构来同时生成文本和音频。 Thinker是一个大语言模型，Talker是一个dual-track自回归模型，Talker基于Thinker的hideen states来生成audio token，最后通过一个sliding-window DiT来解码audio token\n现有omni-model存在的问题：\n缺乏一个系统的，联合训练多个不同模态的方法 需要处理不同模态输出时互相干扰的问题 不能实时理解或者流式输出多模态信息 Qwen2.5 omni的解决方法：\n使用TMRoPE来对齐不同的模态。作者将audio以及video frame按照时间顺序组织成一个交替的架构，然后使用TMRoPE来对齐 使用Thinker-Talker的架构来同时生成文本和音频。Thinker负责文本输出，Talker负责音频的流式输出。 在multimodal encoder中使用Block-wise streaming处理技巧来实现实时理解；在输出时，实现了一个dual-track自回归架构来生成audio token，以及一个DiT模型来解码audio token Qwen2.5 omni的贡献：\n提出了Qwen2.5-omni,可以实时理解多模态信息，并流式输出文本和音频 提出了TMRoPE，一个可以对齐不同模态的RoPE算法 提出了Thinker-Talker的架构，来完成实时理解和音频输出 在多个benchmark上取得了SOTA 架构 总览 Qwen2.5-omni的架构如下图所示 其中Talker类似于人的嘴，负责基于脑中的信息来生成对话；Thinker类似于人的大脑，负责理解输入的信息，并生成对话的上下文。\n输入 text: Qwen的tokenizer audio： Qwen2-Audio的tokenizer vision： Qwen2.5-VL的vision tokenizer 视频以及TMRoPE。 对于视频来说，Qwen2.5-omni采取了和Qwen2.5-VL一样的MRoPE算法。只是这里的temporal ID对应40ms\n音频：对于音频来说，Qwen2.5-omni也是以40ms进行采样然后encoding。\n视频+音频：对于视频+音频来说，Qwen2.5-omni将视频和音频的token进行交替的排列，来保证时间上信息一致\n输出 Text：text由Thinker直接输出，这和LLM的输出方式是一样的。 Speech：Qwen2.5-omni首先构建了qwen-tts-tokenizer，一个speech codec，用于表示speech的关键信息。 然后基于这些关键信息，Talker通过自回归的方式生成audio tokens以及text tokens。 流式输出 对于流式输出来说，现在的模型存在latency，具体影响因素如下：\n处理多模态输入的延迟 TTFT （time to first token） 将第一段speech token解码为audio的时间 模型架构导致的latency，如模型参数等 为了解决这个问题，Qwen2.5-omni采取了以下措施：\nSupport prefilling. 作者修改了audio以及vision encoder来在temporal dimension上支持block-wise attention. 具体来讲，audio encoder只关注2秒内的信息，而vision encoder和Qwen2.5-VL一样，使用了一个MLP patch merger来压缩视觉token Streaming Codec generation. 为了提高流式输出的效率，作者还是用了基于Flow-Matching的DiT，输出的code收线使用Flow-Matching转化为一个mel-spectrogram. 然后再通过一个BigVGAN来重建得到对应的waveform. 作者在这里修改了DiT的attention，将其receptive field 限制为4个block，包括lookback of 2 blocks以及lookahead of 1 block. Pre Training 模型参数初始化：\nLLM： 从Qwen2.5进行初始化 vision encoder：从Qwen2.5-VL的vision encoder进行初始化 audio encoder：从Whisper-large-v3进行初始化 Pretraining和Qwen2.5-VL的训练过程类似，分成了3个stage：\n冻结LLM的参数，仅训练vision encoder和audio encoder， 该阶段使用了audio-text以及image-text pairs来进行训练。 数据：image-text, video-text, video-audio, audio-text以及text corpus， 跟Qwen2-Audio类似，作者将hierarchical tags用自然语言prompt进行替换 训练所有参数，使用更多的多模态数据进行训练。数据包括800B token的image和Video相关数据，300B token的audio数据，100B token的video-audio相关数据。 将模型的上下文扩展到32K。前两个阶段的上下文长度为8192，本阶段使用了long video data等数据来将模型的上下文扩展到32K。 Post Training Thinker 数据格式为ChatML，数据类型包括pure text-based dialogue data, visual-modality conversation data, audio-modality conversation data and mix-modality conversation data.\nTalker 包括三个阶段\nICL training: 训练Talker学习context continuation DPO training: 提升speech generation的stability multi-speaker instruction fine-tuning: 提升模型的naturalness和controllability Evaluation Qwen2.5-omni在两类benchmark上进行来的评测，分别时多模态理解（X-\u0026gt;text）以及音频生成(X-\u0026gt; Speech)\n我们这里主要关注一下speech understanding以及speech generation的benchmark.\n总结 Qwen2.5-omni相当于是结合了Qwen2.5-VL以及Mini-omni，跟mini-omni2的输入输出是类似的。不同的点在于，Qwen2.5-omni使用了Thinker-Talker的架构，然后还使用了TMRoPE来对齐不同的模态。总的来说，感觉模型还是更偏重于audio的理解与生成。\nReference Qwen2.5-omni Mini-omni Mini-omni2 ","date":"2025-04-01T10:29:00+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen2.5-omni/","title":"Notes on Qwen2.5 omni"},{"content":"Introduction A simple note to understand Sigmoid Loss in SigLip 1. Supported by DeepSeek2\nBinary cross entropy loss Suppose we want to solve the binary classification problem, with label $y\\in\\{0, 1\\}$, a common option is to use binary cross entropy loss:\n$$\\mathcal{L}(x, y) = -[y\\log (\\sigma(z)) + (1-y)\\log (1-\\sigma(z))]$$where $z=f_\\theta(x)$ is the logits predicted by our model $f_\\theta$, and $\\sigma$ is the sigmoid function:\n$$\\sigma(z) := \\frac{1}{1 + e^{-z}}$$Let $\\sigma(\\cdot)$ be the sigmoid function, then we have:\n$$ \\sigma(-z) = \\frac{1}{1 + e^{z}} = \\frac{e^{-z}}{1 + e^{-z}} = 1 - \\frac{1}{1 + e^{-z}} = 1- \\sigma(z) $$Now we substitute $\\sigma(-z)=1-\\sigma(z)$ into the loss function, we obtain:\n$$\\mathcal{L}(x, y) = -[y\\log (\\sigma(z)) + (1-y)\\log (\\sigma(-z))]$$Note that $y\\in\\{0, 1\\}$ thus for each instance, there are two cases:\nIf $y=0$, then $\\mathcal{L}(x, y) =-\\log (\\sigma(-z))$ If $y=1$, then $\\mathcal{L}(x, y) =-\\log (\\sigma(z))$ Now we want to use a unified expression to express these two cases. Note that this requires fitting a curve that passes two points $(0, -1)$ and $(1, 1)$. The simplest curve is a straight line $y=2x-1$. So, we can further simplify the loss expression into:\n$$\\mathcal{L}(x, y) = -\\log\\left[\\sigma((2y-1)z)\\right]$$Sigmoid Loss in SigLip Now we recall the sigmoid loss in SigLip:\n$$\\mathcal{L}(\\{\\bm{x}, \\bm{y}\\}_{i=1}^N)=-\\frac{1}{N}\\sum_{i=1}^N\\sum_{j=1}^N\\log \\frac{1}{1+\\exp\\left[z_{ij}(-t\\bm{x}_i\\cdot \\bm{y_j}+b)\\right]}$$where $t, b$ are learnable parameters, and $z_{ij}=1$ if $i=j$ and $z_{ij}=-1$ otherwise.\nTo understand Sigmoid loss, notice that $z_{ij}=2\\mathbb{I}_{i=j}-1$, which exactly matches the form we derived earlier.\nWhy Use Sigmoid Loss? More stable: avoids $\\log 0$. More efficient: Compute Sigmoid once. More Precise: one line of code without condition checking. References SigLip\u0026#160;\u0026#x21a9;\u0026#xfe0e;\nDeepSeek\u0026#160;\u0026#x21a9;\u0026#xfe0e;\n","date":"2025-03-28T14:55:50+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/understanding-sigmoid-loss-in-siglip/","title":"Understanding Sigmoid Loss in SigLip"},{"content":"Aya Vision是一个多模态大语言模型，包含8B, 32B两个size，支持23种语言。Aya Vision基于 Aya Expanse大语言模型。\n模型架构 Aya Vision的模型架构如下图所示\nVision Encoder: SigLip2-patch14-384 Vision-text connector: 2 layer MLP LLM: Aya Expanse 8B/ 32B 训练 训练包含两个stage：\nVision-language alignment: 仅训练vision-text connector，基于image-text pairs进行训练 SFT：训练connector和LLM，基于合成的多语种数据进行训练 多语种数据 为了提高模型的多语种能力，作者先基于English的高质量数据集合成了annotation，然后作者讲这些数据转化为22中语言对应的文本\nModel merging 最后为了提高模型在纯文本任务上的表现，作者还使用了model merging的技巧。具体做法就是merge使用的base language model和SFT之后的vision-language model\nReferences Aya Vision Blog ","date":"2025-03-17T17:58:24+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-aya-vision/","title":"Notes on Aya Vision"},{"content":"作者提出了Gemma3系列大模型，包括1B, 4B, 12B, 27B四个size。4B, 12B, 27B三个size均支持多模态，128K的上下文长度以及140种语言。\n方法 数据处理 数据格式 1 2 3 4 5 6 7 8 [BOS]\u0026lt;start_of_turn\u0026gt;user Who are you?\u0026lt;end_of_turn\u0026gt; \u0026lt;start_of_turn\u0026gt;model My name is Gemma!\u0026lt;end_of_turn\u0026gt; \u0026lt;start_of_turn\u0026gt;user What is 2+2?\u0026lt;end_of_turn\u0026gt; \u0026lt;start_of_turn\u0026gt;model 2+2=4.\u0026lt;end_of_turn\u0026gt; pretrain和SFT的区别在于,pretrain时模型输出以\u0026lt;eos\u0026gt;结束,SFT时模型输出以\u0026lt;end_of_turn\u0026gt;结束.\n图片处理 输入的图片都会被resize到896x896，如果图片精度过大或者不是正方形，则会通过Pan \u0026amp; Scan技巧裁剪为多个子图，然后每个子图分别进行resize。核心处理代码为：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 crop_size_w = int(math.ceil(width / num_crops_w)) crop_size_h = int(math.ceil(height / num_crops_h)) # Don\u0026#39;t apply PaS if crop size is too small. if min(crop_size_w, crop_size_h) \u0026lt; pan_and_scan_min_crop_size: return [] crop_positions_w = [crop_size_w * i for i in range(num_crops_w)] crop_positions_h = [crop_size_h * i for i in range(num_crops_h)] image_crops = [ image[pos_h : pos_h + crop_size_h, pos_w : pos_w + crop_size_w] for pos_h, pos_w in itertools.product(crop_positions_h, crop_positions_w) ] 模型架构 模型包括4个size，分别是1B, 4B, 12B, 27B。1B的模型是单模态的，4B, 12B, 27B的模型是多模态的。对于多模态模型来说：\nVision encoder: Siglip-400M Projection layer: linear layer LLM: Gemma 3 模型参数 Model Vision Encoder Embedding Parameters Non-embedding Parameters context length multilingual training data 1B 0 302M 698M 32K English 2T tokens 4B 417M 675M 3,209M 128K 140+ languages 4T tokens 12B 417M 1,012M 10,759M 128K 140+ languages 12T tokens 27B 417M 1,416M 25,600M 128K 140+ languages 14T tokens Attention layers 为了提高效率，作者将部分layer的self-attention替换为sliding window attention。论文里是将6 layers为一组，替换一组中最后一层为sliding window attention，其余层为self attention。判断某layer是否为sliding window attention的代码为：\n1 2 # config.sliding_window_pattern = 6 self.is_sliding = bool((layer_idx + 1) % config.sliding_window_pattern) Long context 作者将global self-attention的RoPE的base frequency从10K提升到了1M，对于sliding window attention，RoPE的base frequency保持10k不变。\n预训练 数据在模型参数里进行了汇总 tokenizer使用的是Gemini2.0的tokenizer，基于SentencePiece，vocab大小为262K 使用知识蒸馏的方法进行训练，每个token按照教师模型的概率采样256个logits，然后使用cross-entropy loss进行训练 4B, 12B, 27B的模型先在32K的context length上进行训练，然后在128K的context length上进行训练 Quantization Aware training 作者提供了quantized版本的模型，模型基于预训练好的模型使用QAT方法进行SFT 5000 steps左右。最后是各个模型的内存占用情况\nRaw (GB) Quantized (GB) Model bf16 Int4 Int4(blocks=32) SFP8 1B 2.0 0.5 0.7 1.0 +KV 2.9 1.4 1.6 1.9 4B 8.0 2.6 2.9 4.4 +KV 12.7 7.3 7.6 9.1 12B 24.0 6.6 7.1 12.4 +KV 38.9 21.5 22.0 27.3 27B 54.0 14.1 15.3 27.4 +KV 72.7 32.8 34.0 46.1 后训练 数据过滤：筛选掉包含PII，不安全的或者有毒的输出等。留下上下文依赖高，幻觉小的数据\n训练包括升级版的知识蒸馏和基于RL的finetuning，其中RL的reward来自于weight averaged reward models, code execution feedback, ground-truth rewards\n实验 表现 Gemma3的表现如下图所示\n与PaliGemma 2的表现对比\n消融实验 Local attention layers sliding_window_pattern对模型的表现影响不大 sliding window size对模型的表现也不是很大，如下图 使用sliding window attention可以降低KV cache的内存占用 Long context ablation 结论为long context会降低模型的性能\n知识蒸馏 训练token个数比较少的时候，使用小的教师模型效果更好；训练token个数比较多的时候，使用大的教师模型效果更好。\nPan \u0026amp; Scan 使用图片原始的aspect ratio训练的模型效果更好。\nmemorization memorization指模型输出的文本与训练数据中文本的重复率。结果发现Gemma3的memorization要更低一些。\n结论 sliding window attention在Qwen2.5-VL里已经验证过有效,这里在Gemma3上同样验证了有效性. 模型架构与PaliGemma系列基本一致，只是attention改变了，然后LLM从Gemma 2升级到了Gemma 3。 参考文献 Gemma3 Technical Report transformers-Gemma3 code ","date":"2025-03-15T11:15:29+08:00","image":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gemma3/cover.png","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-gemma3/","title":"Notes on Gemma3"},{"content":"Timeline 时间 模型系列 Feature 2023.08 Qwen-VL Vision-centric understanding Multi-lingual Multi-image Fine-grained visual understanding 2024.09 Qwen2-VL Image understanding Video understanding (20min) Agent capability Multilingual support 2025.02 Qwen2.5-VL Document parsing Object grounding Long video understadning and grouding (1 hour) Agent functionality Models name size Vision Adapter LLM 1.0 Qwen-VL 9.6B ViT(1.9B) Cross-attention(0.08B) Qwen-LLM(7B) 2.0 Qwen2-VL-2B 2B ViT(675M) MLP(0.0341B) Qwen2-LLM(1.5B) Qwen2-VL-7B 7B MLP(0.0446B) Qwen2-LLM(7.6B) Qwen2-VL-72B 72B MLP(0.0682B) Qwen2-LLM(72B) 2.5 Qwen2.5-VL-3B 3B ViT(675M) MLP(0.1817B) Qwen2.5-LLM(3B) Qwen2.5-VL-7B 7B MLP(0.3179B) Qwen2.5-LLM(7.6B) Qwe2.5n-VL-72B 72B MLP(0.7267B) Qwen2.5-LLM(72B) Remark\nQwen2-VL的MLP使用的是LayerNorm Qwen2.5-VL的MLP使用的是RMSNorm 出了Qwen-VL之外，所有模型均有对应的Instruct版本，这里为了方便没有列出。Qwen-VL对应的是Qwen-VL-chat。 Data \u0026amp; training Model series stage data Vision Adapter LLM Qwen-VL pretraining 1.4B ✅ ✅ Multi-task pretraining 96.8M ✅ ✅ ✅ SFT 350K ✅ ✅ Qwen2-VL pretraining 600b tokens ✅？ ✅？ Multi-task pretraining 800b tokens ✅？ ✅？ ✅？ SFT ✅？ ✅？ Qwen2.5-VL Visual Pre-Training 1.5T token ✅？ ✅？ Multimodal Pre-Training 2T token ✅ ✅ ✅ Long-Context Pre-Training 0.6T token ✅？ ✅？ ✅？ SFT 2M samples ✅ ✅ DPO ✅ ✅ Remark\nQwen2-VL按照技术报告的说法是follow了Qwen-VL的训练方式，因此这里也使用了同样的表示 带问号的地方代表不确定，为个人猜测。请谨慎参考。 难以确定的原因是Qwen-VL系列将projection layer和ViT作为一个模块，因此比较难区分是否训练了projection layer 技术报告的训练框架图\nReferences\nQwen-VL Qwen2-VL Qwen2.5-VL ","date":"2025-03-09T15:11:29+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/overview-of-qwen-vl-series/","title":"Overview of Qwen-VL series"},{"content":"Introduction 通义团队在3月6号发布了QwQ-32B，一个基于RL的，32B参数的reasoning model，QwQ-32B的表现可以与DeepSeek-R1相比\n模型架构 QwQ-32B基于Qwen2.5-32B开发。主要通过基于outcome-based rewards的RL进行训练。训练过程包括两个stage\nStage 1：本阶段在math和coding domain上进行RL的训练，作者使用了一个verifier来保证最终数学问题结果的正确性，使用了一个code execution server来保证最终生成代码的正确性。 Stage 2：本阶段在通用领域上进行RL的训练。模型基于general reward model和一些rule-based verifier进行训练。应该是类似 DeepSeek-R1 的规则。 实验结果 实验结果如下图所示\nReference blog demo ","date":"2025-03-08T09:46:16+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwq-32b/","title":"Notes on QwQ-32B"},{"content":"我们知道，基于decoder-only transformer的LLM的训练目标是最小化next-token-prediction loss，即给定sequence $x=(x_1,\\dots, x_n)\\in D$，我们的目标为求解以下优化问题\n$$ \\min_{\\theta} -\\sum_{x\\in D}\\log P_{\\theta}(x_i|x_1,\\dots,x_{i-1}) $$这里 $\\theta$ 就是我们的模型参数，$D$ 是我们的训练数据集。\n无数模型通过实际效果告诉我们，这个优化目标可以很好地训练出具有良好泛化能力的大语言模型。但是，我们的问题是，为什么这个优化目标可以训练出智能的模型？ 本文将从压缩即智能的角度来理解这个问题。\n压缩即智能 一个例子 我们首先来看一个简单的例子。给定如下三个0-1字符串：\n1 2 3 01010101010101010101 01001000100001000001 01101000101010100101 我们该如何描述这三个字符串的规律？显然，第一个字符串最简单，它是01字符串重复得到的结果；第二个字符串稍微复杂一些，它在每个1之前插入重复次数的0；第三个字符串则最复杂，它是我随手写的一个字符串，基本没有任何规律，因此，我们只能直接存储这个字符串。\n这个例子告诉我们，一个字符串的规律越简单，我们越容易描述它，因此，我们越容易压缩它。实际上，大语言模型做的也是类似的事情。它们的核心思想是，压缩即智能。\n结论 本文中，我们从压缩即智能的角度来理解大语言模型的原理。我们发现，大语言模型的next-token-prediction其实就是压缩。我们通过压缩让大语言模型学习到了语言中的规律，从而让模型具有了智能。\n参考文献 ","date":"2025-03-06T17:57:51+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/compression-is-intelligence/","title":"compression is intelligence"},{"content":"Introduction 2025年2月20号Qwen团队发布了Qwen2.5 VL技术报告，Qwen2.5 VL包括3B，7B， 72B三个size。Qwen2.5-VL主要在架构，数据上进行了改进。通过评测，Qwen2.5-VL在多个benchmark上取得了SOTA。\nQwen2.5 VL认为已有模型的缺点为：\n计算复杂度高 上下文理解能力有限 细粒度visual perception能力不足 在不同上下文长度下表现不一致 Qwen2.5 VL的贡献为：\n使用了一个从零开始训练的ViT作为vision encoder，并且在ViT中使用了window attention，来提高计算效率 使用了dynamic FPS sampling，用于处理不同采样率的视频输入 将MRoPE扩展到了temporal domain上，进一步提高了模型在与时间相关任务上的表现 使用了更高质量的数据集，其中预训练阶段使用了4.1T的token Qwen2.5 VL的主要亮点为：\n优秀的document parsing能力 精确的object grounding能力 针对长视频的理解和细粒度grounding能力 针对UI的agent functionality 模型架构 总览 Qwen2.5 VL和Qwen2 VL的架构基本一致，包括Vision Encoder，Language Encoder，以及projector layer三个部分，其架构图如下：\nLLM： 使用Qwen2.5 LLM作为LLM，并且将1D-RoPE升级为了MRoPE Vision Encoder： 使用一个从零开始训练的ViT架构，patch_size为14，position embedding为2D-RoPE， attention为window attention和self-attention的混合，其中，只有四层使用的是self-attention.对于输入的图片，ViT会将图片resize到28的整数倍。 Projector Layer： 使用的是一个两层的MLP 模型的参数配置如下图\nVision Encoder Vision encoder的主要改进点为：\n使用了window attention来提升计算效率，window attention的size为 $112\\times 112$， 对应为 $8\\times 8$ 个patch. 这样做的好处是可以不用对图片做scaling 使用了2D RoPE来捕捉空间信息，使用3D RoPE来捕捉视频输入的时间信息 与LLM的结构进行对齐，包括使用RMSNorm替换LayerNorm，使用SwiGLU替换ReLU 输入处理 对于图片输入，Qwen2.5 VL使用原始图片的空间信息来构建坐标，而不是将坐标normalize到[0, 1000]之间，这样让模型可以处理不同精度的图片输入。\n对于视频输入，因为使用了3D RoPE，因此Qwen2.5 VL可以处理不同帧率的视频输入，这样就避免了不同帧率视频对模型视频理解带来的影响。这一点和Apollo里的想法是一样的。具体来说，Qwen2.5 VL首先将连续的两帧group到了一起，然后使用了temporal ID来将position和视频所对应的时间进行对齐。这里可以看Qwen2.5 VL的代码：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 Temporal (Time): 3 patches, representing different segments of the video in time. Height: 2 patches, dividing each frame vertically. Width: 2 patches, dividing each frame horizontally. We also have some important parameters: fps (Frames Per Second): The video\u0026#39;s frame rate, set to 1. This means one frame is processed each second. tokens_per_second: This is a crucial parameter. It dictates how many \u0026#34;time-steps\u0026#34; or \u0026#34;temporal tokens\u0026#34; are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. temporal_patch_size: The number of frames that compose one temporal patch. Here, it\u0026#39;s 2 frames. interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] text temporal position_ids: [101, 102, 103, 104, 105] text height position_ids: [101, 102, 103, 104, 105] text width position_ids: [101, 102, 103, 104, 105] Here we calculate the text start position_ids as the max vision position_ids plus 1. 这里fps为1，表示每一帧对应一秒，tokens_per_second为25，表示每秒包含25个token，temporal_patch_size为2，表示每个temporal patch包含2个frame。因此一个patch里面，就包含了2frame, 对应50tokens. 然后前面提到连续两帧会被group到一起，因此每个temporal patch对应4个spatial patches. 其position_ids为：\n1 [(t, 0, 0), (t, 0, 1), (t, 1, 0), (t, 1, 1)] 训练 预训练 数据 预训练阶段使用了4.1T的token，包括image captions, interleaved image-text data, optical character recognition (OCR) data, visual knowledge (e.g., celebrity, landmark, flora, and fauna identification), multi-modal academic questions, localization data, document parsing data, video descriptions, video localization, and agent-based interaction data. 作者详细介绍了以下几种数据：\nInterleaved image-text data： 主要1. 提高模型的上下文学习能力；2. 保持模型的text-only能力；3.包含一些通用信息。数据清洗包括：1. 基于text-only quality过滤； 2. 基于image-text相关性过滤；3. 基于image-text互补程度过滤；4.基于information density balance过滤. Grounding data：作者使用了Grounding DINO, SAM等模型来生成一些grounding data.为了提升模型在open-vocabulary detection上的能力，作者将训练数据集扩展到了1万个object category上，作者还使用了一些point-based object grounding data来提升模型的能力 Document Parsing data：作者基于text,image, music sheets和chemical formulas合成了一些HTML格式的数据，然后根据tag和bounding box来提升模型document parsing的能力 OCR data：作者使用了开源，合成和in-house的数据集，数据集主要提到multi-lingual以及1M的chart-type数据 Video data：作者通过pipeline构建了long video caption来提升模型长视频的理解能力 Agent data：作者收集了一些mobile device和网页的screenshots,然后基于agent框架来合成控制轨迹 训练 如上图所示，Qwen2.5 VL的预训练阶段包括了三个阶段：\nVisual Pretraining:这一阶段使用了image-caption, knowledge, OCR数据集，旨在提升ViT提取视觉特征的能力 multimodal pretraining:这一阶段在第一阶段的基础上增加了pure-text, interleaved data, VQA, Video grounding, agent data, 旨在提升模型处理复杂视觉信息的能力 long-context pretraining: 这一阶段，在第二阶段的基础上，增加了long video, long agent, long document data，旨在提升模型处理长上下文的能力 后训练 数据 SFT阶段使用大概2M的样本进行训练，其中纯文本和多模态数据占比为1:1，语言主要是中文和英文。\n为了保证数据质量，作者提供了一个数据清洗的pipeline，包括：\ndomain-specific categorization: 作者基于Qwen2-VL-72B构建了Qwen2-VL-Instag，用于将QA pair分为8个大类别，30个小类别 Domain-tailed filtering:作者使用了rule-based和model-based方法来提升数据的质量 rule-based filtering: 重复性检测，格式检测等 model-based filtering: 使用基于Qwen2.5-VL训练的reward model来从多个维度评估QA pair的质量 为了进一步提升模型的推理能力，作者还是用了rejection sampling来refine 数据集。\n训练 post-training阶段分为SFT和DPO两个小阶段，这个阶段都会冻结VIT. SFT阶段使用大多数的训练数据，而DPO阶段专注于image-text data和pure text data，以更好地进行对齐。具体做法就是基于Grounding truth，使用checkpoint来评估数据的质量，然后只保留答案正确的数据来训练。\n评测 通用VQA 文档理解和OCR 空间理解 视频理解和Grounding Agent Reference Github Paper ","date":"2025-03-04T10:46:42+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-qwen2.5-vl/","title":"Notes on Qwen2.5 VL"},{"content":"问题描述 在使用 git clone 命令时，可能会遇到认证错误。\n1 2 3 4 \u0026gt; git clone https://github.com/user-name/some-repository.git 正克隆到 \u0026#39;some-repository\u0026#39;... remote: Invalid username or password. fatal: \u0026#39;https://github.com/user-name/some-repository.git/\u0026#39; 鉴权失败 问题原因\u0026amp;解决方案 本地git配置问题，需要重新认证本地和GitHub的连接。一般是因为本地Git无法和Github的SSH服务器连接。\n方案1：SSH 本方法将Github视为一个SSH服务器，本机通过SSH连接到Github。该方法主要需要在本地生成SSH密钥，并将其添加到Github的SSH密钥列表中。最后通过SSH的方式进行Git操作\n本地生成SSH秘钥 1 ssh-keygen 默认的秘钥地址是~/.ssh/id_rsa，对应的公钥就是~/.ssh/id_rsa.pub。如果需要指定秘钥地址，可以使用-f选项，对于指定的秘钥地址，需要通过以下命令将秘钥地址加入到搜索列表中.\n1 2 3 4 5 # 1. Start the SSH agent \u0026gt; eval \u0026#34;$(ssh-agent -s)\u0026#34; # 2. Add your custom key file \u0026gt; ssh-add path/to/id_rsa 将公钥添加到Github的SSH密钥列表中。使用 cat ~/.ssh/id_rsa.pub 复制公钥，然后打开Github-\u0026gt;设置-\u0026gt;SSH and GPG keys-\u0026gt;New SSH key，将公钥粘贴到Key中，然后点击Add SSH key。将公钥Key中并保存。\n本机通过SSH连接到Github，可以通过以下方式验证\n1 2 3 4 \u0026gt; ssh -T git@github.com # ssh -vT git@github.com # 查看ssh连接的秘钥文件搜索列表 # ssh -i ~/.ssh/id_rsa git@github.com # 自定义密钥的地址，如果使用默认的密钥地址，则不需要指定密钥地址，如果之前已经加入到搜索列表中，则也不需要指定密钥地址 Hi user-name! You\u0026#39;ve successfully authenticated, but GitHub does not provide shell access. 克隆仓库：在clone的时候选择SSH的方式（不是HTTPS的方式），即repo的地址为git@github.com:user-name/some-repository.git 设置remote repo：git remote set-url origin git@github.com:user-name/some-repository.git 通过以上步骤，即可正常使用Git进行操作。\n方案2：Personal Access Token 该方法通过生成一个Personal Access Token来认证本地和GitHub的连接。该方法需要在Github的设置中生成一个Personal Access Token，然后通过该Token来连接remote repo\n生成Personal Access Token. 打开Github-\u0026gt;设置-\u0026gt;Developer settings-\u0026gt;Personal access tokens-\u0026gt;Token(classic)-\u0026gt;Generate new token，在Note中输入本机相关信息，在Expiration中选择token失效日期，在Permissions中选择repo（必选，其他可选），然后点击Generate token。将生成的\u0026lt;token\u0026gt;复制到本地。 克隆仓库：git clone https://user-name:\u0026lt;token\u0026gt;@github.com/user-name/some-repository.git 设置remote repo：git remote set-url origin https://user-name:\u0026lt;token\u0026gt;@github.com/user-name/some-repository.git 通过以上步骤，即可正常使用Git进行操作。\n参考资料 Error: Permission denied (publickey) Managing your personal access tokens ","date":"2025-02-22T10:51:27+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/git-authentication-error/","title":"Git authentication error"},{"content":"论文提出了Kimi k1.5， 一个基于强化学习训练的多模态推理模型。作者介绍了Kimi k1.5的训练方法，以及infra上的优化。作者的主要贡献如下：\n作者发现，模型的上下文长度可以有效提升LLM的推理能力。 作者提出了基于online mirror descent的强化学习训练方法。 作者发现可以通过long context scaling和policy optimization来提升模型的推理能力。 作者提出了long2short的推理方法，可以有效提升short CoT模型的推理能力。 Architecture 论文中包括两个模型，一个是k1.5 base model， 其是一个多模态大模型。另一个是Kimi-1.5 reasoning model（简称为k1.5），k1.5是基于kimi-1.5 base model， 通过强化学习训练得到的推理模型。\nKimi-1.5 base model k1.5 base model是一个基于transformer的多模态大模型，论文没有给出详细的模型结构，只有如下的示意图\nk1.5 base model的训练包括三个阶段：\nVision-language pretraining stage: 这一阶段的目的是让模型拥有语言和视觉的表征能力。训练时，模型先在纯文本的数据上进行训练，然后提高加入图文交错的数据进行训练，直到训练数据中图文交错的数据占比达到30%。模型的更新方式为先冻结LLM更新visual tower, 然后解冻LLM更新所有参数。 Vision-language cooldown stage: 这一阶段的目的是让模型保持多模态理解能力。作者发现，使用合成数据可以提高模型在这一阶段的表现。因此，作者使用闭源大语言模型，基于math, knowledge和code domain的pretraining data来合成了一些QA pair数据对，然后讲这些QA pair数据对加入到训练数据中。 Long-context activation stage: 这一阶段的目的是让模型在长上下文的数据上进行训练，以提升模型在长上下文上的理解能力。这一阶段包含了40%的full attention data以及60%的partial attention data。其中，full attention data是基于真实数据和合成的QA以及summary数据得到的，partial attention data是在cooldown data找那个通过均匀采样得到的。这一阶段，模型的上下文长度从4096逐步增加到131072。 Kimi k1.5 Kimi k1.5是基于k1.5 base model， 通过进一步训练得到的推理模型。k1.5的训练包括四个阶段：\npretraining：这一步就是k1.5 base model的训练， 包括三个小阶段。前面已经介绍过了。 vanilla SFT：这一步的目的是使得k1.5 base model具有指令跟随能力。对于non-reasoning任务，作者先使用人类标注产生一个seed dataset，然后基于seed dataset训练一个seed model，最后根据收集的prompt来使用seed model来生成回答。对于reasoning任务，作者使用了rejection sampling来扩展SFT数据集 long-CoT SFT： 这一步的目的是让模型能够像人类一样进行推理，能够掌握最基本的推理策略，即：planning，evaluation，reflection和exploration。从而为RL训练提供一个良好的初始化。 RL：这一步就是通过RL来提高模型的推理能力。我们在下一节中进行详细的介绍。 RL Problem Definition 给定一个训练数据集 $\\mathcal{D} = \\{(x_i, y_i^\\star)\\}_ {i=1}^n$， 其中 $x_ i$ 是问题， $y_{i}^\\star$ 是ground truth。我们希望找到一个模型 $\\pi_{\\theta}$ ，来解决这个问题。通常问题比较难，因此我们使用chain of thought（CoT）的方法来解决这个问题。具体做法就是让模型输出中间步骤 $z=(z_1, z_2, ..., z_m)$， 来连接问题 $x_i$ 和答案 $y$。其中， $z_j$ 是模型在第 $j$ 步的推理结果，即 $z_t\\sim\\pi_{\\theta}(x_i, z_{","date":"2025-02-08T10:09:52+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-kimi-k1.5/","title":"Notes on Kimi k1.5"},{"content":"Introduction Screen is a terminal multiplexer, which allows you to create multiple sessions in a single terminal window.\nInstallation 1 sudo apt-get install screen Usage 1 screen [-option] [command] Options:\n-S: specify the session name -d: detach from the session -r: reattach to the session -c: execute the command in the session -L: list all sessions -wipe: wipe out all sessions -x: execute the command in the session -p: specify the port number -m: specify the mode -t: specify the title -v: specify the version -h: display help information -v: display version information -q: quit the session Reference Screen usage ","date":"2025-02-05T15:14:43+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/screen-usage/","title":"Screen usage"},{"content":"Introduction 微软在12月12号发布了Phi-4的技术报告,主要关注了数据的构建，特别是。Phi-4是一个14B的大语言模型，主要在STEM相关的QA任务以及推理相关的任务上表现比较好。Phi-4主要在三个方面进行了改进：\n在pre-training和mid-training阶段使用了合成数据进行训练 Data Pre-training Architecture Post-training DPO SFT References Blog Arxiv ","date":"2024-12-16T17:33:52+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-phi-4/","title":"Notes on Phi-4"},{"content":"Introduction A multimodal large language model (MLLM) usually consists of three parts: an encoder $E$ that ingests the information from different modality, a large language model (LLM) that is corresponds to complete various of downstream tasks given multimodal input such as image and text, and an adaption layer $C$ that aligns features of different modality to word embedding space of the LLM. Below is an example MLLM adopting aforementioned architecture: LLaVA [1]\nEfforts have been made to improve the performance of MLLMs. In this post, we aim to review the design of adaption layer and its potential effect on the downstream tasks.\nMethod Suppose the hidden size of the LLM is $d$, the feature produced by encoder $E$ is $V\\in\\mathbb{R}^{P\\times d_v}$, where $P$ is the number of features (number of visual patches if $E$ is an visual encoder) and $d_v$ is the channel dimension. The adaption layer $C$ then aligns the feature $V$ with the word embedding space with $x=C(V)\\in\\mathbb{R}^{Q\\times d}$, where $Q$ is the number of tokens. As we can see, $C$ is actually a mapping from $\\mathbb{R}^{P\\times d_v}$ to $\\mathbb{R}^{Q\\times d}$.\nBased on relationship between $d_v$ and $d$, we can divide projection layers into two types:\nFeature-preserving adaption layer, where $P=Q$ Feature-compressing adaption layer, where $P\u003eQ$. Feature-preserving adaption layer Feature-preserving adaption layer does not change the number of features extracted by $E$. It is used by LLaVA [1] and LLaVA 1.5 [2]. In LLaVA, the adaption layer is a linear layer [2], which is given by $$ x = VW^T, \\text{ where } W\\in\\mathbb{R}^{d\\times d_v}$$ the code reads as:\n1 2 # linear layer adaption_layer = nn.Linear(config.hidden_size, config.num_features) In LLaVA 1.5 , the adaption layer is a two-layer MLP, which is adopted be various of following works. It is given by $$ x = \\phi(VW_1^T)W_2^T$$ where $W_1\\in\\mathbb{R}^{d\\times d_v}$, $W_2\\in\\mathbb{R}^{d\\times d}$, $\\phi$ is a activation function, specified as nn.GELU(). The code reads as:\n1 2 3 4 5 6 # two-layer MLP adaption_layer = nn.Sequential( nn.Linear(config.num_features, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.hidden_size) ) Feature-compressing adaption layer The feature compression adaption layers can be categorized into three types:\naverage pooling attention pooling convolution mapping They usually comprise two steps:\nreduce the number of features from $P$ to $Q$ with a pooling operation: $$ f' = \\mathcal{P}(f)\\in\\mathbb{R}^{Q\\times d_v} $$ project compressed features $f'$ to word embedding space with a transformation $\\mathcal{T}$: $$ x = \\mathcal{T}(f')\\in\\mathbb{R}^{Q\\times d} $$ Average pooling This type of adaption layers use an average pooling as $\\mathcal{P}$ to reduce the number of tokens, followed by a two-layer MLP as $\\mathcal{T}$, which is the same as LLaVA 1.5: $$ f'_i = \\frac{1}{n}\\sum_{j=1}^{n}f_{(i-1)n+j}, i=1,\\dots,Q $$Perceiver Resampler This type of adaption layers use an cross-attention layer as $\\mathcal{P}$, the transformation $\\mathcal{T}$ is also the same as LLaVA 1.5. $$ K = W_kf\\in\\mathbb{R}^{d_c}, V=W_vf\\in\\mathbb{R}^{d_c}, f'=\\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d_c}}\\right)V\\in\\mathbb{R}^{Q\\times d_v} $$ where $W_k, W_v\\in\\mathbb{R}^{d_c\\times d_v}$ and $Q\\in\\mathbb{R}^{Q\\times d_c}$ is a learnable query.\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 class PerceiverResampler(nn.Module): def __init__(self, num_queries, hidden_size, num_features, num_heads): self.num_queries = num_queries self.hidden_size = hidden_size self.num_features = num_features self.query_tokens = nn.Parameter(torch.zeros(self.num_queries, self.num_features), requires_grad=True) self.query_tokens.data.normal_(mean=0.0, std=0.02) self.attention = nn.MultiheadAttention(hidden_size, num_heads) self.layer_norm_kv = nn.LayerNorm(hidden_size) self.layer_norm_q = nn.LayerNorm(hidden_size) def forward(self, x, attention_mask=None): x = self.layer_norm_kv(x) x = x.permute(1, 0, 2) N = x.shape[1] q = self.layer_norm_q(self.query_tokens) q = q.unsqueeze(1).repeat(1, N, 1) out = self.attention(q, k, v, attention_mask=attention_mask)[0] out = out.permute(1, 0, 2) adaption_layer = nn.Sequential( PerceiverResampler(num_queries, hidden_size, num_features, num_heads), MLP(hidden_size, intermediate_size, hidden_size) ) C-Abstractor This type of adaption layers use a combination of convolution layer and averaging pooling as $\\mathcal{P}$. $\\mathcal{T}$ is defined as an additional convolution layers. $$ f_i' = \\frac{1}{n}\\sum_{j=1}^n w_jf_{(i-1)n+j},\\quad x_i = \\sum_{k=-K}^Kw_k'f_{i+k}' $$ where $W=[w_1,\\dots,w_n]^T\\in\\mathbb{R}^n$ and $W'=[w_1,\\dots,w_n]^T\\in\\mathbb{R}^{2K}$ are the weights of the convolution layers.\nD-Abstractor aa\nUsages Comparisons References LLaVA LLaVA 1.5 LLaVA adaption layer code survey ","date":"2024-11-09T09:53:43+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/an-overview-of-adaption-layer-in-multimodal-large-language-models./","title":"An overview of adaption layer in multimodal large language models."},{"content":"Given a string expression representing an expression of fraction addition and subtraction, return the calculation result in string format.\nIntuition Simulate the calculation of fraction addition and subtraction process.\nApproach First, to simplify addition and subtraction, we move the - to numerator part, for example, 3/4-2/3 becomes 3/4+(-2)/3, this step makes all operations become addition.\nSecond, we use a initial value 0/1 to record result to prevent from parsing the first fraction. Thus 3/4-2/3 is actually equivalent to 0/1+3/4-2/3.\nNow, in each loop, we record the sign, the numerator, the denominator. Then we compute the result with the previous result with the formula\n$$ \\frac{a}{b} + \\frac{c}{d} = \\frac{ad-bc}{bd} $$after computation, we use gcd() function to make the resulted fraction irreducible.\nComplexity Time complexity: Iterate the string once. $$O(n)$$ Space complexity: No extra spaces needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Solution { public: string fractionAddition(string expression) { int n = expression.size(); int numerator1 = 0, dominator1 = 1; int numerator2 = 0, dominator2 = 1; int sign = 1; for (int i = 0; i \u0026lt; n;) { if (expression[i] == \u0026#39;-\u0026#39;) { sign = -1; ++i; continue; } if (expression[i] == \u0026#39;+\u0026#39;) { sign = 1; ++i; continue; } numerator2 = 0; while (i \u0026lt; n \u0026amp;\u0026amp; \u0026#39;0\u0026#39; \u0026lt;= expression[i] \u0026amp;\u0026amp; expression[i] \u0026lt;= \u0026#39;9\u0026#39;) { numerator2 = numerator2 * 10 + (expression[i] - \u0026#39;0\u0026#39;); ++i; } numerator2 *= sign; // move sign to numerator sign = 1; // reset sign ++i; // division operator dominator2 = 0; while (i \u0026lt; n \u0026amp;\u0026amp; \u0026#39;0\u0026#39; \u0026lt;= expression[i] \u0026amp;\u0026amp; expression[i] \u0026lt;= \u0026#39;9\u0026#39;) { dominator2 = dominator2 * 10 + (expression[i] - \u0026#39;0\u0026#39;); ++i; } numerator1 = (numerator1 * dominator2 + numerator2 * dominator1); dominator1 = dominator1 * dominator2; if (numerator1 \u0026amp;\u0026amp; dominator1) { int gcd_num = gcd(numerator1, dominator1); numerator1 = numerator1 / gcd_num; dominator1 = dominator1 / gcd_num; } } // corner case if (numerator1 == 0) { return \u0026#34;0/1\u0026#34;; } return to_string(numerator1) + \u0026#34;/\u0026#34; + to_string(dominator1); } }; Reference leetcode 592 ","date":"2024-08-23T20:16:54+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/592.-fraction-addition-and-subtraction/","title":"592. Fraction Addition and Subtraction"},{"content":"Given an integer, flip all bits in its binary representation.\nIntuition Use XOR operation to complete this task.\nApproach We can use the XOR operation to complete this task. However, notice that we cannot simply use num ^ INT_MAX since the leading 1s are still 1s if num is too small. Instead, we need to compute the number of digits in num, and use the corresponding masks.\nNotice that if 2^30 \u0026lt; num \u0026lt;= 2^31-1, in this case the corresponding mask will exceeds the integer ranges, thus we use INT_MAX directly in this case.\nComplexity Time complexity: Compute the digits of a number. $$O(\\log n)$$ Space complexity: No extra spaces needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 class Solution { public: int findComplement(int num) { int num_bits = 0; int temp = num; while (temp) { temp \u0026gt;\u0026gt;= 1; ++num_bits; } if (num_bits == 31) { return num ^ INT_MAX; } int mask = (1 \u0026lt;\u0026lt; num_bits) - 1; return num ^ mask; } }; Reference leetcode 476 ","date":"2024-08-22T21:36:01+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/506.-relative-ranks/","title":"506. Relative Ranks"},{"content":"Approach We can use a dp matrix, where element dp[i][j] represents the minimum turn to print substring s[i...j].\nThen, there are two cases:\ncase 1: We print s[i] separately, now we have dp[i][j] = 1 + dp[i + 1][j]. case 2: There is a char s[k] == s[i], then the string s[i...k] can be obtained by printing some new substrings on the range s[i...k], now we have dp[i][j] = dp[i][k-1]+dp[k+1][j] Combining case 1 and case 2, we can now write the update formula for dynamic programming:\n1 2 3 4 dp[i][j] = 1 + dp[i + 1][j]; for (int k = i + 1; k \u0026lt;= j; ++k) { dp[i][j] = min(dp[i][j], dp[i][k-1] + dp[k+1][j]); } the base case is when i \u0026gt; j, d[pi][j] = 0 and when i=j, dp[i][j] = 1.\nComplexity Time complexity: iterate the matrix once, each iteration takes $O(n)$ time. $$O(n^3)$$ Space complexity: use a dp matrix of size $n\\times n$ to store the result $$O(n^2)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 class Solution { public: int dfs(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; \u0026amp;dp, int start, int end, const string \u0026amp;s) { if (start \u0026gt; end) return 0; if (start == end) return 1; if (dp[start][end] != -1) return dp[start][end]; int ans = 1 + dfs(dp, start + 1, end, s); for (int k = start + 1; k \u0026lt;= end; ++k) { if (s[k] != s[start]) continue; ans = min(ans, dfs(dp, start, k - 1, s) + dfs(dp, k + 1, end, s)); } dp[start][end] = ans; return ans; } int strangePrinter(string s) { int n = s.size(); vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dp(n, vector\u0026lt;int\u0026gt;(n, -1)); return dfs(dp, 0, n - 1, s); } }; References Leetcode 664 ","date":"2024-08-21T15:16:41Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/664.-strange-printer/","title":"664. Strange Printer"},{"content":"A number is ugly if its prime factors is a subset of $\\{2, 3, 5\\}$. We are required to find the $n$-th ugly number\nIntuition Each ugly number is generated from a previous ugly number by multiplying $2$, $3$ or $5$.\nApproach We use three pointers index_2, index_3 and index_5 to record the index of previous ugly number. For example, $4$ is generated from $2$ by multiplying $2$, so the index index_2 is the index of 2.\nThen, we take the minimum of three generated numbers:\n1 dp[i] = min(dp[index_2] * 2, dp[index_3] * 3, dp[index_5] * 5) Finally, we need to update the pointers so that there is no duplication.\nComplexity Time complexity: iterate once. $$O(n)$$ Space complexity: Use a vector to store the result. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class Solution { public: int nthUglyNumber(int n) { vector\u0026lt;int\u0026gt; result{1}; int index_2 = 0, index_3 = 0, index_5 = 0; for (int i = 1; i \u0026lt; n; ++i) { int result_2 = result[index_2] * 2; int result_3 = result[index_3] * 3; int result_5 = result[index_5] * 5; int current_ugly_num = min(result_2, min(result_3, result_5)); if (current_ugly_num == result_2) ++index_2; if (current_ugly_num == result_3) ++index_3; if (current_ugly_num == result_5) ++index_5; } return result; } }; References Leetcode 264 ","date":"2024-08-18T10:37:20Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/264.-ugly-number-ii/","title":"264. Ugly Number II"},{"content":"TLDR This paper proposes a Multimodal Large Language Model VITA (Video, Image, Text, Audio). VITA supports non-awakening interaction and audio interruption for better interactive experience. VITA aims to be an open-sourced version of GPT-4o.\nIntroduction Features of GPT-4o:\na unified framework that processes text, vision, and audio signals in an end-to-end manner, the capability to enable natural multimodal human-computer interaction. Similar to Mini-GPT4, this paper tries to proposed an open-sourced version of GPT-4o.\nMethod Model The architecture of VITA is shown as follows: LLM: Mixtral $8\\times 7$ B Visual Encoder: InternViT-300M-448px Audio Encoder: Mel Filter Bank block To support audio interruption, the author uses two model at the same time, where the generation model is responsible for handling user queries and the other model monitors the environment. The other models starts to work is there is an audio interruption.\nData multimodal instruction tuning Training data of multimodal instruction tuning is given as follows:\nImprovements are made:\nThe questions are randomly (about half) replaced with their audio versions, using TTS technique such as GPT-SoVITS Different system prompts are set to avoid conflicts between different types of data To support human-AI interaction, the noisy audio data are also constructed. Noisy audio samples are generated from existed QA data. These negative sample texts aim to improve the ability of VITA to not respond to non-query-related content.\nTo distinguish three types of queries, the author uses three state tokens:\nToken \u0026lt;1\u0026gt; denotes that the question input is the query audio Token \u0026lt;2\u0026gt; denotes that the question input is the noisy audio. Token \u0026lt;3\u0026gt; signifies the question of pure text. Training pipeline Training pipeline of VITA consists of three stages:\nNon-awakening Interaction There are following requirements and solutions:\nReal-time Tracking of Environmental Sounds. This paper uses SileroVAD to complete the Voice Activity Detection (VAD) task.\nFiltering out noisy audio. This is done by making use of token \u0026lt;2\u0026gt;.\nAudio Interrupt Interaction There are following requirements and solutions:\nReal-time Tracking and Filtering of External Queries. This is done by use another VITA model as stated in Model section. Evaluation Conclusion The paper points out three limitations of VITA:\nEnhancement of Foundational Capabilities. Refinement of Noisy Audio Construction. Building end-to-end TTS in conjunction with LLM. Reference Arxiv paper Github ","date":"2024-08-13T14:36:45+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-vita/","title":"Notes on VITA"},{"content":"Given an array, find all subsets in the array\nIntuition Use DFS to iterate over all subsets and record them.\nApproach A subset can be represented as a binary number of n digits. Each digit is either 0 or 1. We can use DFS to iterate over all possible such format of binary numbers.\nComplexity Time complexity: iterate all binary numbers of length n once. $$O(2^n)$$ Space complexity: $$O(2^n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class Solution { public: void dfs(const vector\u0026lt;int\u0026gt;\u0026amp; nums, int index, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; result, vector\u0026lt;int\u0026gt;\u0026amp; s) { if (index == nums.size()) { result.push_back(s); return; } // digit is 1 s.push_back(nums[index]); dfs(nums, index + 1, result, s); s.pop_back(); // digit is 0 dfs(nums, index + 1, result, s); } vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; subsets(vector\u0026lt;int\u0026gt;\u0026amp; nums) { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; result; vector\u0026lt;int\u0026gt; s; dfs(nums, 0, result, s); return result; } }; Reference Leetcode ","date":"2024-05-21T22:16:49+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/78.-subsets/","title":"78. Subsets"},{"content":"Given a matrix, we wish to find a path that from start to end, such that the path is as far as from the dangerous position.\nIntuition We can convert the problem into finding a path with minimum weights in a graph.\nApproach We first convert the matrix into a graph, with node as position (i, j) and weight w[i,j]\\min_k(|i-thief[k][0]| + |j - thief[k][1]|), where thief[k]=(thief[k][0], thief[k][1]) is the position of the thief thief[k]. To do this, we can use DFS, starting from each thief and iterate through the grids.\nThen, we need to find a path from the start (0, 0) to target (n - 1, n - 1). This can be done via Dijkstra\u0026rsquo;s Algorithm. We use a priority queue to keep track of minimum to-be-visited nodes, this ensures that the newly added nodes are always with the maximum safe factor.\nComplexity Time complexity: iterate all elements twice, the second trial requires to maintain the priority queue. $$O(n^2\\log n)$$ Space complexity: store safe factors matrix. $$O(n^2)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 class Solution { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dirs{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; public: int maximumSafenessFactor(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid) { int n = grid.size(); if (grid[0][0] == 1 || grid[n - 1][n - 1] == 1) return 0; vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; scores(n, vector\u0026lt;int\u0026gt;(n, INT_MAX)); queue\u0026lt;pair\u0026lt;int, int\u0026gt;\u0026gt; q; for (int i = 0; i \u0026lt; n; ++i) { for (int j = 0; j \u0026lt; n; ++j) { if (grid[i][j] == 0) continue; q.push(make_pair(i, j)); scores[i][j] = 0; } } while (!q.empty()) { const auto node = q.front(); int x = node.first, y = node.second; for (const auto\u0026amp; dir : dirs) { int new_x = x + dir[0], new_y = y + dir[1]; if (new_x \u0026lt; 0 || new_x \u0026gt;= n || new_y \u0026lt; 0 || new_y \u0026gt;= n) continue; if (scores[new_x][new_y] \u0026lt;= 1 + scores[x][y]) continue; scores[new_x][new_y] = 1 + scores[x][y]; q.push(make_pair(new_x, new_y)); } q.pop(); } vector\u0026lt;vector\u0026lt;bool\u0026gt;\u0026gt; visited(n, vector\u0026lt;bool\u0026gt;(n, false)); priority_queue\u0026lt;pair\u0026lt;int, pair\u0026lt;int, int\u0026gt;\u0026gt;\u0026gt; pq; pq.push(make_pair(scores[0][0], make_pair(0, 0))); while (!pq.empty()) { auto node = pq.top(); pq.pop(); int safe_factor = node.first; auto pos = node.second; if (pos.first == n - 1 \u0026amp;\u0026amp; pos.second == n - 1) return safe_factor; visited[pos.first][pos.second] = true; for (const auto\u0026amp; dir : dirs) { int new_x = pos.first + dir[0], new_y = pos.second + dir[1]; if (new_x \u0026lt; 0 || new_x \u0026gt;= n || new_y \u0026lt; 0 || new_y \u0026gt;= n) continue; if (visited[new_x][new_y]) continue; int score = min(safe_factor, scores[new_x][new_y]); pq.push(make_pair(score, make_pair(new_x, new_y))); visited[new_x][new_y] = true; } } return -1; } }; References Leetcode ","date":"2024-05-15T20:24:56+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2812.-find-the-safest-path-in-a-grid/","title":"2812. Find the Safest Path in a Grid"},{"content":"Given a matrix, whose element representing the number of golds. Find a path such that the sum is maximized and without crossing the grids that has not gold elements.\nIntuition Use backtracking to find all possible paths, and update the results.\nApproach We use two variables current and result to store results, current stores the sum of golds from start to current position, result stores the final result.\nComplexity Time complexity: In worst case, each grid contains gold, for each position, there are $\\binom{m+n}{m}$ possible paths, so the overall complexity is $$O\\left(mn\\binom{m+n}{m}\\right)$$ Space complexity: No extra spaces needed (without considering recursive stack) $$O(1)$$Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Solution { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dirs{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; public: void backtracking(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid, int i, int j, int \u0026amp;current, int \u0026amp;result) { if (grid[i][j] \u0026lt;= 0) return; current += grid[i][j]; // update result result = max(result, current); // mark as visited int temp = grid[i][j]; grid[i][j] = -1; for (const auto \u0026amp;dir: dirs) { int row = i + dir[0], col = j + dir[1]; if (row \u0026lt; 0 || row \u0026gt;= grid.size() || col \u0026lt; 0 || col \u0026gt;= grid[0].size()) { continue; } backtracking(grid, row, col, current, result); } // retrieve the state grid[i][j] = temp; current -= grid[i][j]; } int getMaximumGold(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid) { int result = 0; for (int i = 0; i \u0026lt; grid.size(); ++i) { for (int j = 0; j \u0026lt; grid[0].size(); ++j) { if (grid[i][j] == 0) continue; int current = 0; backtracking(grid, i, j, current, result); } } return result; } }; Reference Leetcode 1219 ","date":"2024-05-14T20:39:03+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1219.-path-with-maximum-gold/","title":"1219. Path with Maximum Gold"},{"content":"Given a binary matrix, we can flip one column or one row, the goal is to flip zero or more times such that the sum of the number represented by the rows are maximized.\nIntuition The leading 1s are always important than the trailing 1s. So we make sure that 1 appears before 0s.\nApproach The challenge is to determine when to flip the rows and flip the columns. From the intuition, we know that:\na row is flipped only if its first bit is 0, after flipping, the number becomes larger and cannot be flipped again. a column is flipped only if the number of 1s are smaller than 0s. So we flip the rows first and the columns second.\nComplexity Time complexity: iterate matrix twice. $$ O(mn) $$ Space complexity: no extra space needed. $$ O(1) $$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 class Solution { public: int matrixScore(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid) { int m = grid.size(), n = grid[0].size(); for (int i = 0; i \u0026lt; m; ++i) { if (grid[i][0] == 0) { // flip row for (int j = 0; j \u0026lt; n; ++j) { grid[i][j] = 1 - grid[i][j]; } } } // check column by column for (int j = 1; j \u0026lt; n; ++j) { int count = 0; for (int i = 0; i \u0026lt; m; ++i) { count += grid[i][j]; } if (count \u0026lt; (m + 1) / 2) { // flip column for (int k = 0; k \u0026lt; m; ++k) { grid[k][j] = 1 - grid[k][j]; } } } int result = 0; for (int i = 0; i \u0026lt; m; ++i) { for (int j = 0; j \u0026lt; n; ++j) { if (grid[i][j] == 0) continue; result += (1 \u0026lt;\u0026lt; (n - j - 1)); } } return result; } }; References Leetcode ","date":"2024-05-13T20:52:02+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/861.-score-after-flipping-matrix/","title":"861. Score After Flipping Matrix"},{"content":"Given an integer array of size $n$ containing prime integers, it can form $n(n-1)/2$ fractions, we are required to find the $k$-th smallest prime fraction.\nIntuition We can use a priority queue to store prime integers, then we maintain the priority queue.\nApproach1: Brute force Complexity Time complexity: iterate the array once and maintain the priority queue. $$O(n^2\\log k)$$ Space complexity: the size of the priority queue $$O(k)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class Solution { class Compare { public: bool operator()(const vector\u0026lt;int\u0026gt;\u0026amp; a, const vector\u0026lt;int\u0026gt;\u0026amp; b){ return 1.0 * a[0] / a[1] \u0026lt; 1.0 * b[0] / b[1]; // the root is the biggest } }; public: vector\u0026lt;int\u0026gt; kthSmallestPrimeFraction(vector\u0026lt;int\u0026gt;\u0026amp; arr, int k) { priority_queue\u0026lt;vector\u0026lt;int\u0026gt;, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;, Compare\u0026gt; pq; int n = arr.size(); for (int r = 1; r \u0026lt;= k + 1; ++r) { for (int i = 0; i \u0026lt; n; ++i) { if (i + r \u0026gt;= n) break; pq.push(vector\u0026lt;int\u0026gt;{arr[i], arr[i + r]}); if (pq.size() \u0026gt; k) { pq.pop(); } } } return pq.top(); } }; Approach2: Simplification Notice that Approach1 requires iterating over all fractions, can we reduce the time complexity?\nThe solution is by considering the relative order, if we write a matrix whose element a[i][j]=nums[i]/nums[j] (i\u0026lt;j), then we know that a[i][i+1]\u0026gt;...\u0026gt;a[n-1][n] since the array nums are increasing. So, the smallest fraction are in a[1][2], ..., a[n-1][n]. If we take the smallest fraction, and add its successive elements (same column, last row), then we can find the second smallest fraction and so on. This solution requires iterating over $\\max(n, k)$ fractions and $O(n)$ spaces.\nComplexity Time complexity: iterate $\\max(n, k)$ elements and maintain the priority queue. $$O(\\max(n, k)\\log n)$$ Space complexity: the size of the priority queue $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 class Solution { public: vector\u0026lt;int\u0026gt; kthSmallestPrimeFraction(vector\u0026lt;int\u0026gt;\u0026amp; A, int K) { // (fraction, (i, j)) priority_queue\u0026lt;pair\u0026lt;double, pair\u0026lt;int, int\u0026gt;\u0026gt;\u0026gt; pq; for (int i = 0; i \u0026lt; A.size(); ++i) { pq.push({-1.0 * A[i] / A.back(), {i, A.size() - 1}}); } while (--K) { auto t = pq.top().second; q.pop(); --t.second; pq.push({-1.0 * A[t.first] / A[t.second], {t.first, t.second}}); } return {A[pq.top().second.first], A[pq.top().second.second]}; } }; Reference leetcode ","date":"2024-05-10T22:39:54+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/786.-k-th-smallest-prime-fraction/","title":"786. K-th Smallest Prime Fraction"},{"content":"Given a graph and an integer $k$, where the nodes represent values, we can choose an edge, and perform XOR operations on its nodes corresponding to $k$. The goal is the find the maximum sum of the values after 0 or more XOR operations.\nIntuition Since XOR operations satisfies the property that a XOR b XOR b = a, we can record the gain after XOR operation on an edge and finally obtain the result.\nApproach We use total_sum to record the sum of the values of the original trees.\nThen for each node, we perform the XOR operation and record the change if the change \u0026gt; 0, and this change requires one operation, which we add to count. Meanwhile, we use positive_min and negative_max to record the minimum absolute change for reverting use.\nNow after all nodes are computed, we need to compute the result.\nIf count is even, it means the operations satisfied the requirement that the nodes of an edge changes simultaneously If count is odd, then there is one invalid operation and we need to revert the operation. To make the final sum maximum, we can either subtract the positive_min or add negative_max, the result is then obtained by taking the maximum of them. Complexity Time complexity: iterate the array once. $$O(n)$$ Space complexity: No extra spaces are needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 class Solution { public: long long maximumValueSum(vector\u0026lt;int\u0026gt;\u0026amp; nums, int k, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; edges) { long long total_sum = 0; int count = 0; int positive_min = INT_MAX, negative_max = INT_MIN; for (int num : nums) { int new_num = num ^ k; total_sum += num; int change = new_num - num; if (change \u0026gt; 0) { positive_min = min(positive_min, change); total_sum += change; ++count; } else { negative_max = max(negative_max, change); } } if (count % 2 == 0) return total_sum; return max(total_sum - positive_min, total_sum + negative_max); } }; Reference Leetcode ","date":"2024-05-09T20:21:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/3068.-find-the-maximum-sum-of-node-values/","title":"3068. Find the Maximum Sum of Node Values"},{"content":"Given an integer array representing the happiness of the children, select $k$ children such that the sum of their happiness is maximized.\nIntuition Since happiness of all rest children after choosing one child will decrease by 1, their relative order will still the same. So this problem is actually requiring us to select $k$ most happy children.\nApproach Use sorting.\nComplexity Time complexity: Sort the array. $$O(n\\log n)$$ Space complexity: No extra spaces are needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 class Solution { public: long long maximumHappinessSum(vector\u0026lt;int\u0026gt;\u0026amp; happiness, int k) { sort(happiness.begin(), happiness.end()); int n = happiness.size(); long long result = 0; for (int i = 0; i \u0026lt; k; ++i) { result += max(happiness[n - 1 - i] - i, 0); } return result; } }; Reference Leetcode ","date":"2024-05-09T20:21:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/3075.-maximize-happiness-of-selected-children/","title":"3075. Maximize Happiness of Selected Children"},{"content":"ROUGE stands for Recall-Oriented Understudy for Gisting Evaluation. It includes several automatic evaluation methods that measure the similarity between summaries.\nPreliminaries ROUGE-N: N-gram co-occurrence statistics Definitions Given any string $y=y_1y_2\\cdots y_K$, where $y_i,i\\in{1,\\dots,K}$ are characters and an integer $n\\geq1$, we define the set of $n$-gram to be\n$$ G_n(y) = \\{ y_1\\cdots y_n, y_2\\cdots y_{n+1}, \\dots, y_{K-n+1}\\cdots y_K\\} $$NOTE that this is a set with unique elements, for example, $G_2(abab)=\\{ab, ba\\}$.\nGiven any two strings $s$ and $y$, we define substring count $C(s,y)$ to tbe the number of appearances of $s$ as a substring of $y$. For example, $C(ab, abcbab)=2$ since $ab$ appear in $abcbab$ twice in position $1$ and $5$.\nBase Version ROUGE-N is an n-gram recall between a candidate summary $\\hat{y}$ and a set of reference summaries $S=\\{y_1,\\dots,y_n\\}$. ROUGE-N is defined as follows:\n$$ \\text{ROUGE-N}(\\hat{y}, S) = \\frac{\\sum_{i=1}^n\\sum_{s\\in G_N(y_i)}C(s,\\hat{y})}{\\sum_{i=1}^n\\sum_{s\\in G_N(y_i)}C(s, y_i)} $$Features of ROUGE-N:\nthe denominator increases as we add more references, since there might exists multiple good summaries. A candidate summary that contains words shared by more references is favored by the ROUGE-N measure. Multiple references When there are multiple references for a candidate summary, it is suggests to use the following formula:\n$$ \\text{ROUGE-N}(\\hat{y}, S) = \\arg\\max_{i=1,\\dots,n}\\text{ROUGE-N}(\\hat{y}, \\{y_i\\}) $$the above formula is favored for the following reasons:\nThere is no single \u0026ldquo;best\u0026rdquo; reference summary, the multiple reference formula allows ROUGE-N to take into account of all the possible reference summaries and provide a more accurate measure of the quality of the generated summary. The multiple reference formula is more robust. If a reference summary contains a typo or a grammatical error, this can affect the ROUGE-N score. The multiple reference formula can provide a more comprehensive evaluation of the generated summary, since it can allow ROUGE-N to evaluate the generated summary against a wider range pf possible reference summaries. ROUGE-L A sequence $Z=[z_1,\\dots,z_m]$ is a subsequence of another sequence $X=[x_1,\\dots,x_n]$ if there exists a strict increasing sequence $[i_1,\\dots,i_k]$ of indices of $X$ such that for all $j=1,\\dots,k$, we have $x_{i_j}=z_j$.\nGiven two sequences $X$ and $Y$, the longest common subsequences (LCS) of $X$ and $Y$ is a common subsequences with maximum length.\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 def longest_common_subsequence(x: str, y: str) -\u0026gt; int: \u0026#34;\u0026#34;\u0026#34;compute the length of LCS of x and y Args: x (str): a string of length m y (str): a string of length n Return: int: the length of LCS of x and y \u0026#34;\u0026#34;\u0026#34; m, n = len(x), len(y) dp = [[0 for _ in range(m + 1)] for _ in range(n + 1)] for i in range(1, m + 1): for j in range(1, n + 1): if x[i] == y[j]: dp[i][j] = dp[i - 1][j - 1] + 1 else: dp[i][j] = max(dp[i - 1][j], dp[i][j - 1]) return dp[m][n] Sentence-level LCS The intuition of sentence-level LCS is that the longer the LCS of two summary sentences is, the more similar the two summaries are.\nGiven two summaries $X$ of length $m$ and $Y$ of length $n$, assuming $X$ is a reference summary sentence and $Y$ is a candidate summary sentence, the LCS-based recall, precision and F-measure are defined as follows:\n$$ R_{LCS} = \\frac{LCS(X, Y)}{m}, P_{LCS} = \\frac{LCS(X, Y)}{n}, R_{LCS} = \\frac{(1+\\beta^2)R_{LCS}P_{LCS}}{R_{LCS}+\\beta^2P_{LCS}} $$the above formula is called ROUGE-L. $\\beta$ is a hyperparameter.\nFeatures of ROUGE-L are listed as follows:\nIt doesn\u0026rsquo;t require consecutive matches but in-sequence matches, which is more reasonable than n-grams. It automatically includes longest in-sequence common n-grams, there fore no predefined n-gram length is necessary. It\u0026rsquo;s value is less tan or equal to the minimum of unigram F-measure of $X$ and $Y$. The disadvantage of ROUGE-L is that it only counts the main in-sequences words, therefore, other alternative LCSes and shorter sequences are not reflected in the final score. Summary-level LCS We can apply sentence-level LCS-based F-measure score to summary level. Given a reference summary of $u$ sentences $\\{r_1,\\dots,r_u\\}$ containing a total of $m$ words and a candidate summary of $v$ sentences $\\{c_1,\\dots,c_v\\}$ containing a total of $n$ words, the summary-level LCS-based recall, precision and F-measure are defined as follows:\n$$ R_{LCS} = \\frac{\\sum_{i=1}^u\\max_{j}LCS(r_i, c_j)}{m}, P_{LCS} = \\frac{\\sum_{i=1}^v \\max_{j}LCS(r_i, c_j)}{n}, R_{LCS} = \\frac{(1+\\beta^2)R_{LCS}P_{LCS}}{R_{LCS}+\\beta^2P_{LCS}} $$ROUGE-W The basic LCS has a problem that it doesn\u0026rsquo;t differentiate LCSes of different spatial relations within their embedding sequences.\nTo improve the basic LSC method, we can simply remember the length of consecutive matches encountered so fat to a regular two dimensional dynamic program table computing LCS. we call this weighted LCS (WLCS) and use $k$ to indicate the length of the current consecutive matches ending at words $x_i$ and $y_j$.\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 def weighted_longest_common_subsequence(x: str, y: str, f: Callable): \u0026#34;\u0026#34;\u0026#34;use dynamic programming to compute WLCS with a weighted function f Args: x (str): a candidate summary containing m words y (str): a reference summary containing n words f (Callable): a function satisfies f(x+y) \u0026gt; f(x) + f(y) Return: the WLCS score Reference: https://aclanthology.org/W04-1013 \u0026#34;\u0026#34;\u0026#34; x, y = x.split(), y.split() m, n = len(x), len(y) c = [[0 for _ in range(m + 1)] for _ in range(n + 1)] w = [[0 for _ in range(m + 1)] for _ in range(n + 1)] for i in range(1, m + 1): for j in range(1, n + 1): if x[i] == y[j]: # the length of consecutive matches at (i - 1, j - 1) k = w[i - 1][j - 1] c[i][j] = c[i - 1][j - 1] + f(k + 1) - f(k) # remember the length of consecutive matches at (i - 1, j - 1) w[i][j] = k + 1 else: if c[i - 1][j] \u0026gt; c[i][j - 1]: c[i][j] = c[i - 1][j] w[i][j] = 0 # no match at (i, j) elif c[i][j] = c[i][j - 1]: w[i][j] = 0 # no match at (i, j) return c[m][n] where c[i][j] stores the WLCS score ending at word x[i] of x and y[i] of y. w stores the length of consecutive matches at c[i][j]. f is a function of consecutive matches at c[i][j].\nRecall, precision, F-score based on WLCS can be computed as follows:\n$$ R_{WLCS} = f^{-1}\\left(\\frac{WLCS(X, Y)}{f(m)}\\right), P_{WLCS} = f^{-1}\\left(\\frac{WLCS(X, Y)}{f(n)}\\right), R_{LCS} = \\frac{(1+\\beta^2)R_{WLCS}P_{WLCS}}{R_{WLCS}+\\beta^2P_{WLCS}} $$where $f$ is the inverse function of $f$. We call the WLCS-based F-measure as ROUGE-W. Usually, a function $f$ that has a close form inverse is preferred.\nROUGE-S Skip-bigram is any pair of words in their sentence order, allowing for arbitrary gaps. Skip-bigram co-occurrence statistics measure the overlap of skip-bigrams between a candidate translation and a set of reference translations. A sentence with $n$ words will have $\\binom{n}{2}=n(n-1)/2$ skip-bigrams.\nRecall, precision, F-score based on skip-bigram can be computed as follows:\n$$ R_{\\mathrm{SKIP2}} = \\frac{\\mathrm{SKIP2}(X, Y)}{m}, P_{\\mathrm{SKIP2}} = \\frac{\\mathrm{SKIP2}(X, Y)}{n}, R_{\\mathrm{SKIP2}} = \\frac{(1+\\beta^2)R_{\\mathrm{SKIP2}}P_{\\mathrm{SKIP2}}}{R_{\\mathrm{SKIP2}}+\\beta^2P_{\\mathrm{SKIP2}}} $$where $\\mathrm{SKIP2}(X, Y)$ is the number of skip-bigram matches between $X$ and $Y$. The F-score is called ROUGE-S.\nROUGE-SU One problem of ROUGE-S is that it doesn\u0026rsquo;t given any credit to a candidate sentence if the sentence doesn\u0026rsquo;t have any word pair co-occurring with its references.\nTo fix this problem, we extend ROUGE-S with the addition of unigram as counting unit. The extended version is called ROUGE-SU. We can also obtain ROUGE-SU from ROUGE-S by adding a begin-of-sentence marker at the beginning of candidate and reference sentences.\nReference ROUGE: A Package for Automatic Evaluation of Summaries ROUGE Eval google-research rouge ","date":"2024-05-09T17:35:20+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/rouge-recall-oriented-understudy/","title":"ROUGE (Recall-Oriented Understudy)"},{"content":"Given an array of scores, assign different ranks based on their position in the sorted array\nIntuition Sort the array and assign based on the sorted array.\nApproach in C++, we can use the property of the container map to solve this problem, the map is constructed so that the key is the score and the value is the index of the score.\nComplexity Time complexity: construct the map $$O(n\\log n)$$ Space complexity: store the scores and indexes $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 class Solution { public: vector\u0026lt;string\u0026gt; findRelativeRanks(vector\u0026lt;int\u0026gt;\u0026amp; score) { map\u0026lt;int, int\u0026gt; m; for (int i = 0; i \u0026lt; score.size(); ++i) { m[score[i]] = i; } vector\u0026lt;string\u0026gt; result(score.size()); int index = 1; for (auto iter = m.rbegin(); iter != m.rend(); ++iter, ++index) { if (index == 1) result[iter-\u0026gt;second] = \u0026#34;Gold Medal\u0026#34;; else if (index == 2) result[iter-\u0026gt;second] = \u0026#34;Silver Medal\u0026#34;; else if (index == 3) result[iter-\u0026gt;second] = \u0026#34;Bronze Medal\u0026#34;; else result[iter-\u0026gt;second] = to_string(index); } return result; } }; Reference leetcode 506 ","date":"2024-05-08T20:17:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/506.-relative-ranks/","title":"506. Relative Ranks"},{"content":"Given a linked list, we are required to remove some nodes, such that for each node in the result linked list, the value of the node is the greatest from the node to the end of the linked list.\nIntuition We construct the result linked list from right to left, that is, the last node in the linked list is kept, then the pointer goes from right to left util there is a node with greater value. This process can be done via post traversal as tree.\nComplexity Time complexity: iterate the list once. $$O(n)$$ Space complexity: use a set to store numbers. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 /** * Definition for singly-linked list. * struct ListNode { * int val; * ListNode *next; * ListNode() : val(0), next(nullptr) {} * ListNode(int x) : val(x), next(nullptr) {} * ListNode(int x, ListNode *next) : val(x), next(next) {} * }; */ class Solution { public: void reverse_traversal(ListNode* head, stack\u0026lt;int\u0026gt; \u0026amp;s) { if (!head) return; reverse_traversal(head-\u0026gt;next, s); if (s.empty() || head-\u0026gt;val \u0026gt;= s.top()) { s.push(head-\u0026gt;val); } } ListNode* removeNodes(ListNode* head) { stack\u0026lt;int\u0026gt; s; reverse_traversal(head, s); ListNode* dummyhead = new ListNode(0, nullptr); ListNode* cur = dummyhead; while (!s.empty()) { cur-\u0026gt;next = new ListNode(s.top(), nullptr); cur = cur-\u0026gt;next; s.pop(); } return dummyhead-\u0026gt;next; } }; `` # References - [Leetcode](https://leetcode.com/problems/remove-nodes-from-linked-list/description/) ","date":"2024-05-06T19:47:28+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2487.-remove-nodes-from-linked-list/","title":"2487. Remove Nodes From Linked List"},{"content":"Given a linked list representing a non-negative integer, we are required to double this integer and convert it back to a linked list.\nIntuition We can just simulate the process, that is:\nRetrieve the integer represented by the linked list Double the integer Construct the result linked list from the result integer. Approach We can just simulate the process as above. However, if the integer is very large, retrieve the integer may cause overflow. To simplify this process, we can use a stack to store the digits of the original integer, then we double the integer by operating on the top of the stack.\nIn each step, we pop an element from the stack, doubling it and adding it with the carry digit. Then we construct the result linked list with inserting the new node in the front of the head.\nComplexity Time complexity: iterate all digits twice. $$O(n)$$ Space complexity: store the result linked list. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 /** * Definition for singly-linked list. * struct ListNode { * int val; * ListNode *next; * ListNode() : val(0), next(nullptr) {} * ListNode(int x) : val(x), next(nullptr) {} * ListNode(int x, ListNode *next) : val(x), next(next) {} * }; */ class Solution { public: ListNode* doubleIt(ListNode* head) { stack\u0026lt;int\u0026gt; s; ListNode* cur = head; while (cur) { s.push(cur-\u0026gt;val); cur = cur-\u0026gt;next; } int carry = 0; while (!s.empty()) { int num = s.top() * 2 + carry; ListNode* pre = new ListNode(num % 10, cur); carry = num / 10; cur = pre; s.pop(); } if (carry) { ListNode* pre = new ListNode(carry, cur); cur = pre; } return cur; } }; References Leetcode ","date":"2024-05-06T19:47:28+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2816.-double-a-number-represented-as-a-linked-list/","title":"2816. Double a Number Represented as a Linked List"},{"content":"Given a linked list and the node to be deleted, delete the node without accessing the head of the linked list\nIntuition This deletion is same as deleting a node from an array.\nApproach We use two pointers, pre and node to represent the previous and current node of the linked list, and we update the value of pre and node at each iteration. Finally, we delete the last node in the linked list.\nComplexity Time complexity: iterate the linked list once. $$O(n)$$ Space complexity: no extra space needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 /** * Definition for singly-linked list. * struct ListNode { * int val; * ListNode *next; * ListNode(int x) : val(x), next(NULL) {} * }; */ class Solution { public: void deleteNode(ListNode* node) { ListNode* pre = nullptr; while (node-\u0026gt;next) { node-\u0026gt;val = node-\u0026gt;next-\u0026gt;val; pre = node; node = node-\u0026gt;next; } pre-\u0026gt;next = nullptr; } }; References Leetcode ","date":"2024-05-05T20:13:45+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/237.-delete-node-in-a-linked-list/","title":"237. Delete Node in a Linked List"},{"content":"Given an array weight where weight[i] representing the weight of person i, now given the capacity of the boat and the constraint that each boat can carry at most two people, find the minimum number of boats to carry all people.\nIntuition Always pair the lightest person abd the heaviest person to a boat.\nApproach We first sort the array weight in ascending order. Then we use two pointers left=0 and right=n-1 to iterate through the array. In each step, there are two cases:\nIf weight[left]+weight[right] \u0026gt; limit, then we cannot find a peer who can take one boat with right, in this case, right occupies a single boat alone. If weight[left]+weight[right] \u0026lt;= limit, then these two people can take one boat. Complexity Time complexity: sort the array and iterate the array once. $$ O(n\\log n) $$ Space complexity: no extra space needed. $$ O(1) $$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 class Solution { public: int numRescueBoats(vector\u0026lt;int\u0026gt;\u0026amp; people, int limit) { sort(people.begin(), people.end()); int left = 0, right = people.size() - 1; int num = 0; while (left \u0026lt;= right){ if (people[right] + people[left] \u0026lt;= limit) ++left; --right; ++num; } return num; } }; References Leetcode ","date":"2024-05-04T09:06:48+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/881.-boats-to-save-people/","title":"881. Boats to Save People"},{"content":"Given two strings containing digits and dot, compare them in the form of a version number.\nIntuition Since the leading zeros are required to be ignored, we use dot character as separator and compute the integer of each part and compare them individually.\nApproach We use two index i and j to iterate over s1 and s2, we move the index i to the next dot character and compute the integer num1 we have found. Same operations for j to obtain the integer num2. Then num1 and num2 are compared.\nComplexity Time complexity: iterate each string once. $$O(n)$$ Space complexity: no extra space needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 class Solution { public: int compareVersion(string version1, string version2) { int i = 0, j = 0; int m = version1.size(), n = version2.size(); while (i \u0026lt; m || j \u0026lt; n) { int num1 = 0, num2 = 0; while (i \u0026lt; m \u0026amp;\u0026amp; version1[i] != \u0026#39;.\u0026#39;) { num1 = 10 * num1 + (version1[i++] - \u0026#39;0\u0026#39;); } while (j \u0026lt; n \u0026amp;\u0026amp; version2[j] != \u0026#39;.\u0026#39;) { num2 = 10 * num2 + (version2[j++] - \u0026#39;0\u0026#39;); } if (num1 \u0026lt; num2) return -1; if (num1 \u0026gt; num2) return 1; ++i, ++j; } return 0; } }; References Leetcode ","date":"2024-05-03T15:13:23+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/165.-compare-version-numbers/","title":"165. Compare Version Numbers"},{"content":"This post is a notes on understanding how transformer works in an algorithm perspective.\nIntroduction Transformer is a neural architecture that is used for neural language processing. Transformer receives an embedding matrix, which represents a sentence as input, and outputs a matrix of the same size as the embedding matrix, then the output can be used for downstream tasks.\nNotation We denote $V=[N_V]:=\\{1,\\dots,N_V\\}$ as the vocabulary of tokens or words or characters. We denote $\\bm{x}=x[1...n]:=x[1]...x[n]\\in V^n$ be a sequence of tokens, for example, a sentence or a paragraph. Given a matrix $M\\in\\mathbb{R}^{m\\times n}$, $M[i,:]\\in\\mathbb{R}^n$ is the $i$-th row of $M$, $M[:, j]\\in\\mathbb{R}^m$ is the $j$-th column of $M$. Tokenization Tokenization determines how the text are represented. Given a piece of text, for example, \u0026quot;I have an apple\u0026quot;, we seek to find a proper way to represent this sentence.\nCharacter level tokenization. In this setting, $V$ is the English alphabet plus punctuation. This tends to yield very long sequences (depends on the character contained in the raw text). Word level tokenization. In this setting, $V$ is the set of all English words plus punctuation. Word level tokenization is straightforward, but it tends to required a very large vocabulary and cannot handle new vocabulary at test time. Subword tokenization. This is the most common way used in nowadays, $V$ is the set containing the commonly used word segments like \u0026ldquo;ing\u0026rdquo;, \u0026ldquo;est\u0026rdquo;. This can be computed via Byte-Pair Encoding (BPE) algorithm. We suppose the length of the input text is $L$, if the length input text exceeds $L$, we chunk it. After tokenization, each element in the vocabulary is assigned to a unique index $i\\in\\{1,\\dots,N_V-3\\}$, and a number of special tokens are added to the vocabulary. For example:\nmask_token$=N_V-2$, used in masked language modeling bos_token$=N_V-1$ and eos_token$=N_V$, these two tokens are used to represent the beginning and the end of the sequence. Finally, a piece of raw text is represented as a sequence of indices, often called token IDs corresponding to its subwords, preceded by bos_token and followed by eos_token.\nEmbedding The embedding layer is used to represent each token as a vector that contains richer semantic information. The embedding contains two parts:\ntoken embedding, where each token is embedded into a vector space positional embedding, where embeds the position information of the tokens. Token embedding Given a sequence of token ID, we now need to represent each token as a vector in $\\mathbb{R}^d$.\nThe simplest way is to use one-hot embedding, where each token $i$ is represented a vector $[0,\\dots,1,\\dots,0]\\in\\mathbb{R}^{N_V}$ whose elements are all $0$ excepts that $i$-th position is equal to $1$. However, the problem is that the vocabulary size $N_V$ is two large.\nTo solve this problem, we can train a learnable embedding model, of which parameter is a matrix $W_e\\in\\mathbb{R}^{d\\times N_V}$, its $i$-th row corresponds to vector representation of the token $i$:\n$$ \\bm{e} = W_{e}[:, i]\\in\\mathbb{R}^d $$Position embedding There is a problem in token embedding, that is, it doesn\u0026rsquo;t contain consider the order of tokens. In latter, we show that the self-attention mechanism is equivariant to a permutation matrix $X\\Pi$ of data $X$, where $\\Pi$ is a permutation matrix, that is,\n$$ \\mathrm{Sa}(X\\Pi) = \\mathrm{Sa}(X)\\Pi $$the above equation indicates that the self-attention layer learn no position information at all!\nTo solve this problem, we add a positional embedding to token embedding. There are two kinds of embeddings:\nAbsolute positional embeddings. In this setting, a matrix $W_P\\in\\mathbb{R}^{d\\times N}$ is learned or design to indicate the position of tokens. Mathematically, we have $$ \\bm{e}_p = W_p[:, \\mathrm{index}(i)]\\in\\mathbb{R}^{d} $$where $\\mathrm{index}(i)$ is the index of token $i$ in the input sequence. 2. Relative positional embeddings. We leave this in latter notes. Compared to absolute positional embeddings, relative positional embeddings uses offset information, which performs well when the input sequence is too long.\nThe final embedding of a token $i$ is given by\n$$ \\bm{e} = W_e[:, i] + W_p[:, \\mathrm{index}(i)]\\in\\mathbb{R}^{d} $$Attention The idea of attention mechanism is: Given a sequence of token, to predict the current token, which token should I pay attention to? For example, I opened the door with my ___, we may answer key, password or fingerprint etc. This is because we notice that we opened the door, so to predict the next token, we should make use of the information. What attention mechanism does is to quantify this process and make them parallel and learnable.\nSingle query attention We first consider a simple example. Given the embedding of the current token $\\bm{e}\\in\\mathbb{R}^d$ and the list of context tokens $[\\bm{e}_1,\\dots,\\bm{e}_N]\\in\\mathbb{R}^{d\\times N}$, the attention is given as follows:\ncompute query vector: $\\bm{q}=W_q\\bm{e}+b_q\\in\\mathbb{R}^{d}$ compute key vectors: for $i=1,\\dots,L$, $\\bm{k}_i=W_k\\bm{e}_i+b_k\\in\\mathbb{R}^{d}$ compute value vectors: for $i=1,\\dots,L$, $\\bm{v}_i=W_v\\bm{e}_i+b_v\\in\\mathbb{R}^{d}$ compute attention weights: let $\\bm{s}=[\\bm{q}^T\\bm{k}_1,\\dots,\\bm{q}^T\\bm{k}_L]\\in\\mathbb{R}^{N}$, then: $$ \\bm\\alpha = \\mathrm{softmax}\\left(\\frac{\\bm{s}}{\\sqrt{d}}\\right)\\in\\mathbb{R}^{N}$$ 5. compute vector representation of the token and context combined:\n$$ \\bm{v}'= \\sum_{i=1}^N\\alpha_i\\bm{v}_i\\in\\mathbb{R}^{d} $$where $W_q,W_k,W_v\\in\\mathbb{R}^{d\\times d}$, $b_q,b_k,b_v\\in\\mathbb{R}$.\nGeneral attention To extend the single query attention to general form, we consider the embedding matrix $X\\in\\mathbb{R}^{D\\times N}$, the context matrix $Z\\in\\mathbb{R}^{d\\times C}$ and a mask matrix $M\\in\\mathbb{R}^{D\\times D}$, then the attention is computed as follows:\ncompute query matrix: $$ Q=W_qX+\\bm{b}_q\\in\\mathbb{R}^{D\\times N}$$ 2. compute key matrix:\n$$ K=W_kZ+\\bm{b}_k\\in\\mathbb{R}^{D\\times C}$$ 3. compute value vectors:\n$$ V=W_vZ+\\bm{b}_v\\in\\mathbb{R}^{D\\times C}$$ 4. compute attention weights:\n$$\\mathrm{Sa}(X) = \\mathrm{softmax}\\left(M\\odot \\frac{K^TQ}{\\sqrt{D}}\\right) \\in\\mathbb{R}^{C\\times N} $$ where $\\odot$ is the element-wise product. 5. output the updated representations of tokens in $X$, folding the information from tokens in $Z$\n$$ \\tilde{V} = V\\odot \\mathrm{Sa}(X)\\in\\mathbb{R}^{D\\times N} $$There are two kinds of mask matrices depending on which attention we are using:\nBidirectional attention, in this case, $M=\\bm{1}\\bm{1}^T\\in\\mathbb{R}^{C\\times N}$. Undirectional attention, in this case, $M[i,j] = \\bm{1}_{i\\leq j}$, where $\\bm{1}_{i\\leq j}$ is the indicator function. Multi-head Attention The previous describes the operation of a single head. In practice, transformers run multiple attention heads in parallel and combine their outputs, this is called multi-head attention. The idea behind multi-head attention can be summarized as follows:\nIn high dimensional space, two vectors are usually far from each other, with multiple attention heads, we can reduce the dimension of the representation. With multiple attention heads, each heads may focus on specific semantics of the representation. For example, one head focuses on positiveness, and one another head focuses on noun/verb semantics. For simplicity, we denote the single-head attention as $\\mathrm{attention}(X, Z\\mid M)$, suppose we have $H$ heads, then we compute $\\tilde{V}_i$ for each heads:\n$$ \\tilde{V}_i = \\mathrm{attention}(X, Z\\mid M)\\in\\mathbb{R}^{D\\times N},i=1,\\dots,H $$ then attention representation are concatenated together:\n$$ V = [\\tilde{V}_1^T, \\dots, \\tilde{V}_H^T]^T\\in\\mathbb{R}^{HD\\times N} $$combined via an output matrix $W_o\\in\\mathbb{R}^{D\\times HD}$:\n$$ \\tilde{V} = W_oV+\\bm{b}_o\\in\\mathbb{R}^{D\\times N} $$We denote the multi head attention as $\\mathrm{MhSa}(X, Z\\mid M)$.\nTransformer layer After computing the multi head attention, we can now construct a transformer layer, which can also be stacked as convolution neural networks. A transformer layer can be constructed by the following operations:\nMulti head attention (residual), $X\\gets X + \\mathrm{MhSa}(X, Z\\mid M)$. Layer norm, $X \\gets \\mathrm{LayerNorm}(X)$ Multi layer perception $\\bm{x}_i\\gets \\bm{x}_i + \\mathrm{mlp}(\\bm{x}_i), i=1,\\dots,N$ Layer norm, $X \\gets \\mathrm{LayerNorm}(X)$ where $\\mathrm{LayerNorm}$ is the layer norm operation. $\\mathrm{mlp}$ is a multi layer perception, usually it consists of one hidden layer of size $4D$, that is, then umber of neurons in three layers are $D, 4D, D$.\nUsually, a large language model consists of multiple transformer layers.\nUnembedding The unembedding learns to convert a vector representation of a token and its context $\\bm{e}$ into a distribution over the vocabulary elements.\n$$ \\bm{p} = \\mathrm{softmax}(W_u\\bm{e})\\in \\Delta(V)\\subseteq \\mathbb{R}^d $$ where $\\Delta(V)$ is a simplex over the set $V$.\nReference Understanding Deep Learning Chapter 12 Formal Algorithms for Transformers ","date":"2024-05-02T13:13:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/formal-algorithms-for-transformer/","title":"Formal Algorithms for Transformer"},{"content":"TLDR This paper proposed an open-sourced, GPT-4 liked multimodal language model (MLLM), mini-GPT4, aiming to provide inspiration of how to build an MLLM.\nIntroduction GPT-4 has exhibited remarkable vision language abilities, however, due to it is not open sourced, it is hard to explore the construction and how the GPT-4 is trained.\nTo solve this problem, this paper builds mini-GPT4, which utilizes an advanced LLM, Vicuna and BLIP-2, to build an open sourced MLLM. The result show that mini-GPT4 possesses numerous capabilities similar to those demostrated by GPT-4.\nMethod The architecture of mini-GPT4 is shown as follows:\nmini-GPT4 consists of three parts:\nVision Encoder same as used in BLIP-2, which is built upon a ViT backbone and a Q-former network. A single projection layer, which aligns the encoded visual features with the Vicuna language model. Language decoder: Vicuna. Training mini-GPT4 only trains the linear projection layer, this includes two stages:\nFirst pre-training stage: in this stage, the language decoder and vision encoder are remain frozen, only the linear projection layer is trained. The training dataset contains Conceptual Caption, SBU and LAION. The problem is that the output may be incoherent like repetitive words or sentences. Fine-tune dataset generation: the pre-trained mini-GPT4 is used to generate image-text pair as fine-tune data, to improve the text quality, the generated text is polished with the help of ChatGPT. Finally, the dataset is manually refined. Second-stage fine-tuning: Fine-tune the mini-GPT-4 with the high quality fine-tune data. Reference Github Paper Homepage ","date":"2024-05-02T13:13:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/minigpt-4-enhancing-vision-language-understanding-with-advanced-large-language-models/","title":"MiniGPT-4-Enhancing Vision-Language Understanding with Advanced Large Language Models"},{"content":"This post introduces how to understand t-SNE.\nIntroduction t-SNE an dimension reduction algorithm, which projects high-dimensional data into low-dimensional space. Thus the algorithm can be used to visualize the data distribution.\nTo understand how t-SNE works, we first review the SNE algorithms, then we introduce the t-SNE algorithm.\nSNE Method Stochastic Neighbor Embedding, or SNE, is the previous version of t-SNE.\nThe basic idea behind SNE is that: data points that are close in high-dimensional space should be close in lower-dimensional space too.\nFormally speaking, given a data set $X\\in\\mathbb{R}^{D\\times N}$ consisting of $N$ data points, with each data point lies in $D$ dimensional space. Our goal is to reduce the data points into $d\u003c\u003c D$ dimensional space $Y\\in\\mathbb{R}^{d\\times N}$, that is, we seek to find a map $f:\\mathbb{R}^{D\\times N}\\to \\mathbb{R}^{d\\times N}$ such that $f(X)=Y$. Usually, $d=2$ or $d=3$ for visualization use.\nSNE measures \u0026ldquo;close\u0026rdquo; in a probabilistic way. The similarity is represented by converting Euclidean distance between data points to condition probabilities:\n$$ p_{j\\mid i} = \\frac{\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_j\\Vert^2/(2\\sigma_i^2)\\right)}{\\sum_{k\\neq i}\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_k\\Vert^2/(2\\sigma_i^2)\\right)} $$the above equation can be interpreted as the probability of point $\\bm{x}_j$ being a neighbor of point $\\bm{i}$ is proportional to the distance between them. $\\sigma_i$ is the variance of the Gaussian distribution that is centered on data point $\\bm{x}_i$. We introduce the method for determining $\\sigma_i$ later.\nSimilarly, we can construct a probability distribution $q$ based on $Y$.\n$$ q_{j\\mid i} = \\frac{\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_j\\Vert^2\\right)}{\\sum_{k\\neq i}\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_k\\Vert^2\\right)} $$where we set the variance as $1/\\sqrt{2}$ following the original paper.\n$p_{i\\mid i}$ and $q_{i\\mid i}$ are set $0$ since we are only interested in modeling pairwise similarities.\nNow we want $q_{j\\mid i}$ are as close as $p_{j\\mid i}$, that is, we want two distributions are as close as to each other. This can be measured by Kullback- Leibler divergence, which is written as:\n$$ C(P, Q) = \\sum_{i=1}^N\\mathrm{KL}(P_i\\Vert Q_i)=\\sum_{i=1}^N\\sum_{j=1}^N p_{j\\mid i}\\log \\frac{p_{i\\mid j}}{q_{i\\mid j}} $$where $P_i=[p_{1\\mid i},\\dots,p_{N\\mid i}]\\in\\mathbb{R}^N$ and $Q_i=[q_{1\\mid i},\\dots,q_{N\\mid i}]\\in\\mathbb{R}^N$.\nChoosing $\\sigma$ Now we introduce how to choose $\\sigma$. Note that $\\sigma$ determines the distribution of data points, larger $\\sigma$ indicates sparser distribution of data points. The original paper uses perplexity to measure such sparsity. It is defined as\n$$ \\mathrm{Perp}(P_i) = 2^{H(P_i)} $$where $H(P_i)$ is the Shannon entropy of $P_i$ measured in bits:\n$$ H(P_i) = -\\sum_{i=1}^N p_{j\\mid i}\\log p_{j\\mid i} $$The perplexity can be interpreted as a smooth measure of the effective number of neighbors. The performance of SNE is fairly robust to changes in the perplexity, and typical values are between 5 and 50.\nNotice that $p_{j\\mid i}$, by setting different value on $\\mathrm{Perp}(P_i)$, we can obtain different $\\sigma_i$ via binary search.\nOptimization Our goal now becomes minimizing $C(p, q)$ over variables $\\bm{y}_1,\\dots,\\bm{y}_N\\in\\mathbb{R}^d$, given $p$ and hyperparameter $\\sigma_i, i=1,\\dots,N$. This can be done via gradient descent methods. The gradient is now given by\n$$ \\frac{d C}{d\\bm{y}_i}=2\\sum_{j=1}^N\\left(p_{j\\mid i} - q_{j\\mid i} + p_{i\\mid j}- q_{i\\mid j} \\right)(\\bm{y}_i-\\bm{y}_j)\\in\\mathbb{R}^d, \\ i=1,\\dots,N $$We can write it in matrix form and add a momentum term:\n$$ Y^{t+1} = Y^t + \\beta\\frac{dC}{dY} + \\alpha_t\\left(Y^{t-1}-Y^{t-2}\\right) $$where $\\beta$ is the step size and $\\alpha_t$ is momentum parameter,\n$$ Y^t = [\\bm{y}_1^t,\\dots, \\bm{y}_N^t]\\in\\mathbb{R}^{d\\times N} ,\\ \\frac{dC}{dY} = \\left[\\frac{d C}{d\\bm{y}_1},\\dots,\\frac{d C}{d\\bm{y}_N}\\right]\\in\\mathbb{R}^{d\\times N} $$t-SNE Symmetric SNE The first difference between t-SNE and SNE is the probability, t-SNE uses symmetric version of SNE to simplify computations.\nDifferent from SNE, symmetric SNE uses a joint probability instead of a condition probability:\n$$ C(P, Q) = \\sum_{i=1}^N\\mathrm{KL}(P\\Vert Q)=\\sum_{i=1}^N\\sum_{j=1}^N p_{ij}\\log \\frac{p_{ij}}{q_{ij}} $$where $p_{ij}$ and $q_{ij}$ are defined as\n$$ p_{ij} = \\frac{\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_j\\Vert^2/(2\\sigma^2)\\right)}{\\sum_{k\\neq r}\\exp\\left(-\\Vert\\bm{x}_r-\\bm{x}_k\\Vert^2/(2\\sigma^2)\\right)},q_{ij} = \\frac{\\exp\\left(-\\Vert\\bm{x}_i-\\bm{x}_j\\Vert^2\\right)}{\\sum_{k\\neq r}\\exp\\left(-\\Vert\\bm{x}_r-\\bm{x}_k\\Vert^2\\right)} $$the problem if joint probability $p_{ij}$ is that if there is an outlier $\\bm{x}_i$, then $p_{ij}$ will be extremely small for all $j$. This problem can be solved by defining $p_{ij}$ from the conditional probability $p_{i\\mid j}$ and $p_{j\\mid i}$\n$$ p_{ij} = \\frac{p_{j\\mid i}+p_{i\\mid j}}{2N} $$This ensures that\n$$ \\sum_{j=1}^N p_{ij} \u003e \\frac{1}{2N} $$for all $\\bm{x}_i$, in result, each data point makes a significant contribution to the cost function.\nIn this case, the gradient of the cost function is now given by\n$$ \\frac{d C}{d\\bm{y}_i}=4\\sum_{j=1}^N\\left(p_{ij} - q_{ij}\\right)(\\bm{y}_i-\\bm{y}_j)\\in\\mathbb{R}^d, \\ i=1,\\dots,N $$t-SNE Experiments show that symmetric SNE seems to produce maps that are just as good as asymmetric SNE, and sometimes even a little better.\nHowever, there is a problem with SNE, that is, the crowding problem, which manifests as a tendency for points in the low-dimensional space to be clustered too closely together, particularly in high-density regions of the data.\nThe causes of the crowding problems are:\nData points in high dimensional space tend to far from each other, which makes the distance information less useful. SNE aims to preserve the local structure of the data points, but it can struggle with non-linear relationships. The projected data points will be closed to each other due to this reason. The optimization algorithm used by SNE can get stuck in local minimum. To alleviate the crowding problem, t-SNE is introduced in the following way: In the high-dimensional space, we convert distances into probabilities using a Gaussian distribution. In the low-dimensional map, we can use a probability distribution that has much heavier tails than a Gaussian to convert distances into probabilities. This allows a moderate distance in the high-dimensional space to be faithfully modeled by a much larger distance in the map and, as a result, it eliminates the unwanted attractive forces between map points that represent moderately dissimilar data points.\nt-SNE uses student t-distribution in low-dimensional map:\n$$ q_{ij} = \\frac{\\left(1+\\Vert\\bm{y}_i-\\bm{y}_j\\Vert^2\\right)^{-1}}{\\sum_{k\\neq r}\\left(1+\\Vert\\bm{y}_i-\\bm{y}_j\\Vert^2\\right)^{-1}} $$A Student t-distribution with a single degree of freedom is used, because it has the particularly nice property that $\\left(1+\\Vert\\bm{y}_i-\\bm{y}_j\\Vert^2\\right)^{-1}$ approaches an inverse square law for large pairwise distances $\\Vert\\bm{y}_i-\\bm{y}_j\\Vert$ in the low-dimensional map.\nCompared to Gaussian distribution, t-distribution is heavily tailed。\nA computationally convenient property of t-SNE is that it is much faster to evaluate the density of a point under a Student t-distribution than under a Gaussian because it does not involve an exponential, even though the Student t-distribution is equivalent to an infinite mixture of Gaussians with different variances.\nThe gradient of t-SNE is given by\n$$ \\frac{d C}{d\\bm{y}_i}=4\\sum_{j=1}^N\\left(p_{ij} - q_{ij}\\right)(\\bm{y}_i-\\bm{y}_j)\\left(1+\\Vert\\bm{y}_i-\\bm{y}_j\\Vert^2\\right)^{-1}\\in\\mathbb{R}^d, \\ i=1,\\dots,N $$The advantages of t-SNE gradients over SNE are given by:\nThe t-SNE gradient strongly repels dissimilar data points that are modeled by a small pairwise distance in the low-dimensional representation. Second, although t-SNE introduces strong repulsions between dissimilar data points that are modeled by small pairwise distances, these repulsions do not go to infinity. The algorithm is given as follows:\nOptimization There are some optimizations that can be used to improve performance of t-SNE:\nEarly compression, which is used to force the map points to stay close together at the start of the optimization Early exaggeration, which is used to multiply all of the pi j’s by, for example, 4, in the initial stages of the optimization Implementation Reference Visualizing Data using t-SNE ","date":"2024-05-02T13:13:12+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-t-sne/","title":"Notes on t-SNE"},{"content":"Given a string ,count the number of wonderful substrings, where a wonderful string is a string that at most one character appears in an odd number of times.\nIntuition We use the prefix sum and bitwise representation to solve the problem. The idea is that every string can be represented by a state, and we can count the number of wonderful substrings by performing operations among states\nApproach According to the hint and the problem description, we have:\nthere are $10$ possible variables, which represent a to j we only care about if the variable appears in odd number of times or in even number of times, which means that each variable has $2$ possible states. so the number of overall possible states is $2^10=1024$. We define the prefix substring as prefix[i] = s[1...i].\nFor each substring, we can represent the state of the prefix as an index of the array, that is, state(prefix[i]) is a number between 0 and 1023, with each byte representing if the corresponding character appears in odd number of times (1) or in even number of times (0). A wonderful string is then defined as a string whose state representation is 0 or a power of 2.\nnow the substring s[i...j] is defined as prefix[j]-prefix[i], in state representations, there are three cases:\nif one character appears in odd number of times in both prefix substring, then in the result substring, the character appears in even number of times, which is equivalent to (1, 1) -\u0026gt; 0 if one character appears in even number of times in both prefix substring, then in the result substring, the character appears in even number of times, which is equivalent to (0, 0) -\u0026gt; 0 If one character appears in even number of times in one prefix substring, and appears in odd number of times in the other prefix substring, then in the result substring, the character appears in even number of times, which is equivalent to (1, 0) -\u0026gt; 1 so, according to analysis, the minus operation in string corresponding to the XOR operation in states. That is, the state of s[i...j] is given by state(s[i...j])=state(prefix[j]) XOR state(prefix[i]).\nNote that there is also a simplification since if a XOR b = c, then a XOR c = b. Instead of using a for loop to compute all substrings that ended with character j:\n1 2 3 4 5 for (int i = 0; i \u0026lt;j; ++i) { // substring s[i...j] int substr_state = state[j] ^ state[i]; // check if substr is a wonderful substring } with the property of XOR, wo directly check if state(prefix[j]) XOR 2^k exists, that is:\n1 2 3 4 for (int i = 0; i \u0026lt; 10; ++i) { int state = state[j] ^ (1 \u0026lt;\u0026lt; i); // check if state exists and add them. } Complexity Time complexity: iterate string once $$O(n)$$ Space complexity: stores the state representation. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 class Solution { public: long long wonderfulSubstrings(string word) { vector\u0026lt;int\u0026gt; bits(1024); bits[0] = 1; long long result = 0; int prefix = 0; for (char c: word) { prefix ^= 1 \u0026lt;\u0026lt; (c - \u0026#39;a\u0026#39;); result += bits[prefix]; for (int i = 0; i \u0026lt; 10; ++i) { result += bits[prefix ^ (1 \u0026lt;\u0026lt; i)]; } ++bits[prefix]; } return result; } }; References:\nLeetcode ","date":"2024-05-02T10:51:29+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1915.-number-of-wonderful-substrings/","title":"1915. Number of Wonderful Substrings"},{"content":"Given a list containing integers, find the largest positive integer that its negative also exists in the list.\nIntuition There are two ways to solve the problem.\nOne way is to use two sum, we find the all pairs of integers such that their sum is zero and the one is the negative of the other. The second way is to use two pointer, we move left and right pointer utils their absolute values are equal. Approach 1: two sum Similar to Two Sum, for each num in nums, we store its negative -num in the hash table, however, notice that the added term can be determined by num, we can use a set instead.\nComplexity Time complexity: iterate the list once. $$O(n)$$ Space complexity: use a set to store numbers. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 class Solution { public: int findMaxK(vector\u0026lt;int\u0026gt;\u0026amp; nums) { set\u0026lt;int\u0026gt; s; int result = -1; for (int num: nums) { if (s.find(-num) != s.end()) { result = max(result, abs(num)); } else { s.insert(num); } } return result; } }; Approach 2: two pointer We first sort the lists, then we use left=0 and right=length(nums) pointer to iterate the list, there are three cases:\nif nums[left] == -nums[right], then we directly return nums[right] since this is the largest one (note that the list is sorted) if nums[left] \u0026lt; -nums[right], then we update left to left + 1 since (nums[left], -nums[left]) cannot be found in the list. if nums[left] \u0026gt; -nums[right], then we update right to right - 1 since (-nums[right], nums[right]) cannot be found in the list. Complexity Time complexity: sort the array and iterate the list once. $$O(n\\log n)$$ Space complexity: no extra spaces are needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 class Solution { public: int findMaxK(vector\u0026lt;int\u0026gt;\u0026amp; nums) { sort(nums.begin(), nums.end()); int left = 0, right = nums.size() - 1; while (left \u0026lt;= right) { if (nums[left] == -nums[right]) return nums[right]; else if (nums[left] \u0026lt; -nums[right]) ++left; else --right; } return -1; } }; `` # References - [Leetcode](https://leetcode.com/problems/largest-positive-integer-that-exists-with-its-negative/description/) - [Leetcode Two Sum](https://leetcode.com/problems/two-sum/) ","date":"2024-05-02T10:28:05+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2441.-largest-positive-integer-that-exists-with-its-negative/","title":"2441. Largest Positive Integer That Exists With Its Negative"},{"content":"Given a string ans a specified character, reverse the prefix that is ended with the specified character.\nIntuition Simulate the process.\nApproach find the first occurrence of the specified character, then reverse the prefix.\nComplexity Time complexity: find the first occurrence of the specified character. $$O(n)$$ Space complexity: no extra spaces are needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 class Solution { public: string reversePrefix(string word, char ch) { auto pos = word.find(ch); if (pos == word.npos) return word; reverse(word.begin(), word.begin() + pos + 1); return word; } }; ","date":"2024-05-01T09:01:22+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2000.-reverse-prefix-of-word/","title":"2000. Reverse Prefix of Word"},{"content":"Given a tree, return a vector, with each elements of the vector is the sum of the distances between the corresponding node and all other nodes.\nIntuition Intuitively, I want to use DFS + memorization to solve the problem, however, such method exceeds te time limit.\nThen, I refer to some solutions, and solve the problem.\nApproach DP -\u0026gt; TLE In the first, I am thinking that we can use dp[i][j] to represent the distance between the node i and the node j. Initially, dp are initialized as follows:\nIf there is an edge between i and j, then dp[i][j]=dp[j][i]=1. If there is an edge between i and j, then dp[i][j]=dp[j][i]=INFINITY. we traversal from node 0 to node n-1, for each node i, we compute the distance between i and all other nodes j. There are two cases:\ndp[i][j] != INFINITY, then we directly returns dp[i][j] dp[i][j] == INFINITY, then it means the distance between node i and node j hasn\u0026rsquo;t been computed, we then update it as follows: $$ dp[i][j] = \\min_{k\\in N(i)} (1 + dp[k][j]) $$where $N(i)$ is the nodes that adjacent to node i. The min operation is used here since the node k and the node j may not connected (without passing node i).\nThe code is given as follows. However, the time complexity is $O(n^2)$, which exceeds the time limit.\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 class Solution { public: int dfs(const vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; adjs, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; \u0026amp;dp, vector\u0026lt;bool\u0026gt; \u0026amp;visited, int start, int end) { if (dp[start][end] != INT_MAX) return dp[start][end]; if (visited[start]) return INT_MAX; visited[start] = true; int distance = INT_MAX; for (int next_node: adjs[start]) { distance = min(distance, dfs(adjs, dp, visited, next_node, end)); } visited[start] = false; if (distance != INT_MAX) dp[start][end] = distance + 1; return dp[start][end]; } vector\u0026lt;int\u0026gt; sumOfDistancesInTree(int n, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; edges) { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dp(n, vector\u0026lt;int\u0026gt;(n, INT_MAX)); vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; adjs(n); // initialize the distance matrix for (const auto\u0026amp; edge: edges) { dp[edge[0]][edge[1]] = 1; dp[edge[1]][edge[0]] = 1; adjs[edge[0]].push_back(edge[1]); adjs[edge[1]].push_back(edge[0]); } vector\u0026lt;bool\u0026gt; visited(n, false); for (int i = 0; i \u0026lt; n; ++i) { for (int j = i + 1; j \u0026lt; n; ++j) { dfs(adjs, dp, visited, i, j); } // the graph is undirected for (int j = 0; j \u0026lt; i; ++j) { dp[i][j] = dp[j][i]; } } vector\u0026lt;int\u0026gt; result(n); for (int i = 0; i \u0026lt; n; ++i) { for (int j = 0; j \u0026lt; n; ++j) { if (j == i) continue; result[i] += dp[i][j]; } } return result; } }; Tree + Traversal The second way is not easy to figure out. We decompose it into two parts:\ncompute the distance between the root node and all other nodes. Convert the root node from one to another and update the result. Sum of distance from root node to all other nodes. Consider one example tree with root set as 0:\n1 2 3 4 5 0 / \\ 1 2 / | \\ 3 4 5 We first define dist[i] as the distance of all child nodes of node i to node i. Then we have:\n$dist[i] = 0$ if i is a leaf node. $dist[i] = \\sum_{j\\in C(i)}dist[j] + |C(i)|$ if i is not a leaf node, where $C(i)$ is the offspring of node i We can now compute the sum of distances between root 0 to all other nodes, which is result[0].\nNow we need to compute all other results. Repeating the above process is too time consuming, we need to reduce the time complexity. We are seeking a way to compute result[i] from result[j], where $j\\in C(i)$.\nNote that for $k\\in C(j)$, when we compute result[i], we are computing distance from $k$ to $i$, so we need to reduce by 1 since their height decreases (the root changes from i to j). On the other hand, all other nodes, which are not offspring of node j, is added by 1 since the height of them increases. Thus, the transformation reads:\n$$ result[j] = result[i] - |C(j)| + n - |C(j)| $$Complexity Time complexity: visited all nodes twice. $$ O(n) $$ Space complexity: use two vectors of size $n$ to store the result. $$ O(n) $$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class Solution { public: void post_order_traversal(const vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; adjs, vector\u0026lt;int\u0026gt;\u0026amp; count, vector\u0026lt;int\u0026gt;\u0026amp; result, int node, int parent) { for (int next_node : adjs[node]) { if (next_node == parent) continue; post_order_traversal(adjs, count, result, next_node, node); count[node] += count[next_node]; result[node] += count[next_node] + result[next_node]; } ++count[node]; } void pre_order_traversal(const vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; adjs, vector\u0026lt;int\u0026gt;\u0026amp; count, vector\u0026lt;int\u0026gt;\u0026amp; result, int node, int parent) { for (int next_node : adjs[node]) { if (next_node == parent) continue; result[next_node] = result[node] - count[next_node] + count.size() - count[next_node]; pre_order_traversal(adjs, count, result, next_node, node); } } vector\u0026lt;int\u0026gt; sumOfDistancesInTree(int n, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; edges) { vector\u0026lt;int\u0026gt; count(n), result(n); vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; adjs(n); for (const auto \u0026amp;edge: edges) { adjs[edge[0]].push_back(edge[1]); adjs[edge[1]].push_back(edge[0]); } // compute result[0] post_order_traversal(adjs, count, result, 0, -1); // compute other results from result[0] pre_order_traversal(adjs, count, result, 0, -1); return result; } }; References Leetcode Solution ","date":"2024-04-29T21:24:06+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/834.-sum-of-distances-in-tree/","title":"834. Sum of Distances in Tree"},{"content":"Given an integer array, we can flip one bit of one element of the array at every step, we asked to compute the minimum flips, such that the XOR of all elements of the array is equal to the given integer $k$.\nIntuition XOR of multiple bits is always equal to either 1 or 0, and changes one of the input bits will cause the result change to the other one.\nApproach We first compute the XOR of all elements of the array, then we compare bit by bit with the XOR result and the given k, utils all bits becomes the same.\nComplexity Time complexity: compute the XOR result $$O(n)$$ Space complexity: No extra spaces are needed. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 class Solution { public: int minOperations(vector\u0026lt;int\u0026gt;\u0026amp; nums, int k) { int result = nums[0]; for (int i = 1; i \u0026lt; nums.size(); ++i) { result ^= nums[i]; } vector\u0026lt;int\u0026gt; bits_result, bits_k; while (k) { bits_k.push_back(k % 2); k /= 2; } while (result) { bits_result.push_back(result % 2); result /= 2; } int count = 0; int i = 0; for (; i \u0026lt; bits_result.size() \u0026amp;\u0026amp; i \u0026lt; bits_k.size(); ++i) { if (bits_result[i] != bits_k[i]) ++count; } for (; i \u0026lt; bits_result.size(); ++i) { if (bits_result[i] == 1) ++count; } for (; i \u0026lt; bits_k.size(); ++i) { if (bits_k[i] == 1) ++count; } return count; } }; ","date":"2024-04-29T21:16:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2997.-minimum-number-of-operations-to-make-array-xor-equal-to-k/","title":"2997. Minimum Number of Operations to Make Array XOR Equal to K"},{"content":"To reduce the gap of the performance of the model on the training dataset and the test dataset, we need use regularization methods.\nIntroduction Possible reasons for the discrepancy that the model performs worse on the test dataset than on the training dataset are:\nthe model describes statistical peculiarities of the training dataset that are not representative of the true mapping from the input to the output, that is, overfitting. the model is unconstrained in areas with no training areas, leading to suboptimal predictions Explicit Regularization Methods Consider fitting a model $f(\\bm{x}; \\phi)$ with parameter $\\phi$ using a training dataset $\\{\\bm{x}_i,y_i\\}$, we seek to minimize the loss function $L(\\phi)$:\n$$ \\hat{\\phi} = \\arg\\min_{\\phi}L(\\phi;\\{\\bm{x}_i,y_i\\}) $$Now, to bias the minimization towards certain solutions, we add an additional term:\n$$ \\hat{\\phi} = \\arg\\min_{\\phi}[L(\\phi;\\{\\bm{x}_i,y_i\\}) + \\lambda g(\\phi)] $$where $g(\\phi)$ is called the regularization term and $\\lambda\u003e0$ is a hyperparameter.\nIn the probabilistic perspective, we can construct the loss function from maximum likelihood estimation, or MLE, that is\n$$ \\hat{\\phi} = \\arg\\max_{\\phi}\\left[ \\prod_{i}\\mathrm{Pr}(y_i\\mid \\bm{x}_i,\\phi) \\right] $$The regularization term can be considered as a prior $\\mathrm{Pr}(\\phi)$， in this way, we are now using maximum a posteriori criterion:\n$$ \\hat{\\phi} = \\arg\\max_{\\phi}\\left[ \\prod_{i}\\mathrm{Pr}(y_i\\mid \\bm{x}_i,\\phi)\\mathrm{Pr}(\\phi) \\right] $$Implicit Regularization Methods Gradient descent and stochastic gradient descent are commonly used to minimize the loss functions, however, neither of them moves neutrally to the minimum of the loss function, thus the implicit regularization method is proposed to solve the problem.\nImplicit Regularization in gradient descent The change of the parameters $\\phi$ is defined by the differential equation:\n$$ \\frac{d{\\phi}}{d t} = -\\frac{dL}{d\\phi} $$gradient descent uses difference quotient with increment (or learning rate) $\\alpha$ to approximate the change of $\\phi$:\n$$ \\frac{\\phi_{t+1}-\\phi_{t}}{\\alpha}=-\\frac{dL}{d\\phi} \\Rightarrow \\phi_{t+1} = \\phi_{t} - \\alpha\\frac{dL}{d\\phi} $$However, this discretization causes deviation from the continuous path.\nTo fix the problem, and extra item is added to the loss to avoid the deviation caused by discretization:\n$$ \\tilde{L}(\\phi) = L(\\phi) + \\frac{\\alpha}{4}\\left\\Vert \\frac{dL}{d\\phi}\\right\\Vert^2 $$Implicit Regularization in stochastic gradient descent A similar approach can be applied to stochastic gradient descent, which reads as\n$$ \\tilde{L}(\\phi) = L(\\phi) + \\frac{\\alpha}{4|B|}\\sum_{i\\in B}\\left\\Vert \\frac{dL_{i}}{d\\phi}-\\frac{dL}{d\\phi}\\right\\Vert^2 $$ where $L_{B}$ is the loss on the batch $B$.\nHeuristic Methods Early stopping Early stopping means that we stop the training procedure before the model becomes overfitting. By stopping early, we prevent the model captures the corner features of the training dataset.\nEarly stopping has a single hyperparameter, the number of steps after which the training is stopped, this is chosen usually with the help of the validation dataset.\nEnsembling Ensembling means we train multiple models on the training dataset, and during the inference time, we take the average inference result of each model. The technique improves the test performance with the sacrifices of training and storing multiple models.\nThere are some ensembling methods:\nUse different random initializations. This leads the model reaches different local minimum and may help reduce the overfitting. Generate several different datasets by re-sampling the training dataset and train model on each of them.This is also known as bootstrap aggregating or bagging, this division can smooth out the data, since each model tries to predict the distribution of data that is not included in its training dataset. Dropout Drop out randomly clamps a subset of hidden units of the layer at each iteration of SGD. This makes the model depends on general feature instead of some specific feature, since the specific feature may be masked.\nAt test time, we can run the network as usual with all the hidden units active; however, the network now has more hidden units than it was trained with at any given iteration, so we multiply the weights by one minus the dropout probability to compensate. This is known as the weight scaling inference rule.\nApplying noise Dropout can be interpreted as applying multiplicative Bernoulli noise to the network activations. We can apply noise to other parts of the model during training.\nWe can add noise to the input data, this smooth out the learned function. We can also add noise to model parameters, this encourages the model to be robust to small perturbations of the weights. We can also perturb the labels. We can change the label of a portion of the training dataset, this can prevent the model from being overconfident. Bayesian inference The MLE approach tries to find a function $f(\\bm{x};\\phi)$ that fit the dataset $\\{\\bm{x}_i,y_i\\}$, this approach may be overconfident about the task since the bias of the training data construction.\nTo overcome such bias, we treats the parameters $\\phi$ as unknown variables instead of scalars. Then we find a distribution over the parameters $\\phi$ conditioned on the training data $\\{\\bm{x}_i,y_i\\}$, using Bayesian theorem:\n$$ \\mathrm{Pr}(\\phi\\mid \\{\\bm{x}_i,y_i\\}) = \\frac{\\prod_{i}\\mathrm{Pr}(y_i\\mid \\bm{x}_i,\\phi)\\mathrm{Pr}(\\phi)}{\\int\\prod_{i}\\mathrm{Pr}(y_i\\mid \\bm{x}_i,\\phi)\\mathrm{Pr}(\\phi)d\\phi} $$where $\\mathrm{Pr}(\\phi)$ us the prior probability of the parameters, and the denominator is a normalizing term.\nThen the prediction for unseen item $\\{\\bm{x},y\\}$ is given by the following infinite weighted sum:\n$$ \\mathrm{Pr}(y\\mid x, \\{\\bm{x}_i,y_i\\}) = \\int \\mathrm{Pr}(y\\mid \\bm{x},\\phi)\\mathrm{Pr}(\\phi\\mid \\{\\bm{x}_i,y_i\\})d\\phi $$This is an infinite weighted ensemble, where the weights depend on\nthe prior probability of the parameters the agreement with the data Though Bayesian approach is capable of representing the data more robust, it is hard to implement since there is no way to represent an distribution. Current implementation simplifies the distribution as Gaussian distribution and each parameter is replaced with the mean $\\mu$ and standard deviation $\\sigma$ of the Gaussian distribution.\nTransfer learning and multi-task learning In transfer learning, the model is first pre-trained before training or fine-tuning on the task we are interested in. The idea is that the model may learn some good representation of the data from the main task. Alternatively, we can think transfer learning as initializing the model parameters in a reasonable area such that the minimum is better compared to the random initialization.\nMulti-task learning is a related technique that the model is trained on multiple related tasks concurrently. In this way, the model can learn from multiple datasets and multiple objectives, this encourages the model to learn the essential part of the tasks.\nSelf-supervised learning In some cases, we do not have multiple datasets for pre-training or for multi-tasks. To solve this problem, we can use self-supervised learning to generate large amounts of label-free data. There are two families of self-supervised learning: generative and contrastive.\nIn generative self-supervised learning, part of each data example is masked, and the task is to predict the masked part. For example, given a sentence, we can mask the verb and ask for the model to predict the correct verb, the helps the model to learning semantic meaning of a sentence.\nIn contrastive self-supervised learning, we try to group related data and separated unrelated data. For example, a cat is more similar to another cat compared with a dog. In this way, the model can learn more robust representations and can be adapted to new tasks easily.\nAugmentation Augmentation aims to expand the training dataset, we can perform transformation to each training data without changing the labels, for example we can rotate, flip a image of cat. The augmentation is to teach the model to be invariant to these irrelevant data transformations.\nSummary To summarize the regularization methods, we use the following picture to depict the mechanisms. Reference Understanding Deep Learning Chapter 9 ","date":"2024-04-27T18:02:02+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/regularization-methods-in-deep-learning/","title":"Regularization methods in deep learning"},{"content":"Given a string displaying in the ring format, we can move one character at each step either in clockwise or anticlockwise (or hold still), now we need to retrieve a given string with minimum steps.\nIntuition We can construct a graph from the string and use DFS to search all possible paths, and find the path with minimum steps.\nApproach DFS -\u0026gt; TLE In the beginning, I am going to use DFS to find all possible paths and find the path with minimum steps.\nFirst, we need construct the graph, each character is adajcent to all other characters in ring, so there are n=length(ring) nodes and (n-1)^n edges. each edge has a weight representing the distance between two characters: weight(i,j)=min(abs(i-j), n - abs(i-j)) (here we use index of the character to represent the node).\nThen, we can use DFS to search all possible paths, in each point, we have a state (ring_index, key_index), representing the current index on ring and key, there are two cases:\nring[ring_index] == key[key_index], in this case, we change key_index to key_index+1 and keeps ring_index unchanged (hold still). ring[ring_index] != key[key_index], in this case, we need to rotate the string ring to make ring[ring_index] == key[key_index], this takes step and notice that there may multiple choices, so we need to go over all of them. Once key_index == len(key), we have found a path and we can now update the result.\nThe code is given as follows:\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 class Solution { public: void dfs(unordered_map\u0026lt;int, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026gt; \u0026amp;adjs, const string \u0026amp;ring, const string \u0026amp;key, int ring_index, int key_index, int \u0026amp;min_rotates, int \u0026amp;current_rotates) { if (key_index == key.size()) { min_rotates = min(min_rotates, current_rotates); return; } // character matches if (ring[ring_index] == key[key_index]) { dfs(adjs, ring, key, ring_index, key_index + 1, min_rotates, current_rotates); return; } // character doesn\u0026#39;t match for (const auto \u0026amp;next_node: adjs[ring_index]) { int next_ring_index = next_node[0], rotates = next_node[1]; if (ring[next_ring_index] != key[key_index]) continue; current_rotates += rotates; dfs(adjs, ring, key, next_ring_index, key_index, min_rotates, current_rotates); current_rotates -= rotates; } } int findRotateSteps(string ring, string key) { unordered_map\u0026lt;int, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026gt; adjs; int n = ring.size(); for (int i = 0; i \u0026lt; ring.size(); ++i) { for (int j = 0; j \u0026lt; ring.size(); ++j) { if (j == i) continue; int diff = abs(i - j); int rotates = min(diff, n - diff); adjs[i].push_back(vector\u0026lt;int\u0026gt;{j, rotates}); } } int ring_index = 0, key_index = 0; int min_rotates = INT_MAX, current_rotates = 0; dfs(adjs, ring, key, ring_index, key_index, min_rotates, current_rotates); int spell = key.size(); return min_rotates + spell; } }; Dynamic programming The problem of DFS is that, its time complexity grows exponentially if there have multiple repeat characters, which causes TLE (time limit exceeded) error.\nSo, to reduce the complexity, we can construct the solution from bottom to up. That is, we remember the path to go from the current state, and now we now go one step back, util we back to the original state.\nWe use dp[i][j] to represent start from ring[j], the minimum rotate steps we need to recover the string key[i...m], where m=length(key), the target then becomes finding out dp[0][0].\nNote that we can easily compute dp[m-1][j], j=1,\\dots,n, since there is only one character we need to recover, so we start from ring[j], and rotate util we find a character ring[k] such that ring[k]==key[m-1], the minimum steps is then updated. The update formula is then given by\n$$ dp[i][j] = \\min_{k=1,\\dots,n,\\ ring[k]=key[i]}(dp[i][j],\\ dp[i + 1][k] + step(j, k)) $$where $step(j,k)=\\min(|j-k|,\\ n-|j-k|)$.\nComplexity Time complexity: The dp matrix is of size $m\\times n$, and each time, we need to iterate over the ring once. $$O(mn^2)$$ Space complexity:The dp matrix is of size $m\\times n$ $$O(mn)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class Solution { public: int findRotateSteps(string ring, string key) { int m = key.size(), n = ring.size(); vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dp(m + 1, vector\u0026lt;int\u0026gt;(n)); for (int i = m - 1; i \u0026gt;= 0; --i) { // start from ring[j] for (int j = 0; j \u0026lt; n; ++j) { dp[i][j] = INT_MAX; // Find the feasible target ring[k] == key[i] for (int k = 0; k \u0026lt; n; ++k) { if (ring[k] != key[i]) continue; int rotates = min(abs(k - j), n - abs(k - j)); dp[i][j] = min(dp[i][j], dp[i + 1][k] + rotates); } } } return dp[0][0] + m; } }; Reference leetcode 514 Solution ","date":"2024-04-27T11:55:28+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/514.-freedom-trail/","title":"514. Freedom Trail"},{"content":"Given a matrix, find the minimum falling path sum from top to bottom, with no two adjacent rows sharing the same column.\nIntuition Same as the previous version, we only change the update formula.\nApproach We use dp[i][j] to represent the minimum falling path sum that ends with grid[i][j]. The update rules is then given by\n$$ dp[i][j] = \\min_{k=1,\\dots,n,k\\neq j}(grid[i][j] + dp[i - 1][k]) $$Complexity Time complexity: We need to iterate all elements of the matrix once, and each iterate requires to iterate its last row once, which is $O(n)$. This can be reduced to $O(n^2\\log n)$ by compute the smallest, second smallest elements of the last row. $$O(n^3)$$ Space complexity: the dp matrix is of size $n\\times n$. This can be reduced to $O(n)$ by use a $n\\times 2$ matrix since each row only relates to its last row. $$O(n^2)$$Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 class Solution { public: int minFallingPathSum(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid) { int n = grid.size(); vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dp(n, vector\u0026lt;int\u0026gt;(n, INT_MAX)); dp[0] = grid[0]; for (int row = 1; row \u0026lt; n; ++row) { for (int col = 0; col \u0026lt; n; ++col) { for (int last_col = 0; last_col \u0026lt; n; ++last_col) { if (last_col == col) continue; dp[row][col] = min(dp[row][col], grid[row][col] + dp[row - 1][last_col]); } } } int result = INT_MAX; for (int i = 0; i \u0026lt; n; ++i) { result = min(result, dp[n - 1][i]); } return result; } }; Reference Leetcode 1289 ","date":"2024-04-26T20:44:17+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1289.-minimum-falling-path-sum-ii/","title":"1289. Minimum Falling Path Sum II"},{"content":"BLEU (Bilingual Evaluation Understudy) is a widely used metric that evaluates the quality of the translated text with respect to the reference translations.\nIntroduction The formula of BLEU is defined as follows:\n$$ \\mathrm{BLEU}_ {w_n}(\\hat{S}, S) = \\mathrm{BP} \\cdot \\exp\\left(\\sum_{n=1}^Nw_n\\log p_n(\\hat{S}, S) \\right) $$where\n$\\mathrm{BP}$ represents the Brevity Penalty to penalize the translations that are shorter than the reference translations. $p_n$ represents the modified $n$-gram precision score $w_n$ represents the weight for $p_n$, it satisfies $0\\leq w_n\\leq1$ and $\\sum_{n=1}^Nw_n=1$. Now we interprets three parts in detail.\nInterpretation Definitions Given any string $y=y_1y_2\\cdots y_K$, where $y_i,i\\in{1,\\dots,K}$ are characters and an integer $n\\geq1$, we define the set of $n$-gram to be\n$$ G_n(y) = \\{ y_1\\cdots y_n, y_2\\cdots y_{n+1}, \\dots, y_{K-n+1}\\cdots y_K\\} $$NOTE that this is a set with unique elements, for example, $G_2(abab)=\\{ab, ba\\}$.\nGiven any two strings $s$ and $y$, we define substring count $C(s,y)$ to tbe the number of appearances of $s$ as a substring of $y$. For example, $C(ab, abcbab)=2$ since $ab$ appear in $abcbab$ twice in position $1$ and $5$.\nModified precision score We start from the one candidate translation $\\hat{y}$ and one reference translation $y$. The modified $n$-gram is defined as\n$$ p_n(\\hat{y}, y) = \\frac{\\sum_{s\\in G_n(\\hat{y})}\\min (C(s,\\hat{y}), C(s,y))}{\\sum_{s\\in G_n(\\hat{y})}C(s,\\hat{y})} $$The quantity measures how many $n$-grams of the reference translation $y$ appears in the candidate translation $\\hat{y}$. In case that $\\hat{y}$ is too short, we take a minimum between $C(s,\\hat{y})$ and $ C(s,y)$. Then we normalize to make $p_n(\\hat{y}, y)$ comparable among multiple translation pairs.\nNow, suppose we have a candidate translation corpus, $\\hat{S}=\\{\\hat{y}^1,\\dots,\\hat{y}^M\\}$, and for each candidate translation $\\hat{y}^i$, we have a reference translation corpus (there are multiple translations can represent the same meaning) $S_i=\\{y^{i,1},\\dots,y^{i,N_i}\\}$. We define $S=\\{S_1,\\dots,S_M\\}$, then our modified $n$-gram precision is defined as\n$$ p_n(\\hat{S}, S) = \\frac{\\sum_{i=1}^M\\sum_{s\\in G_n(\\hat{y})}\\min (C(s,\\hat{y}), \\max_{y\\in S_i}C(s,y))}{\\sum_{i=1}^M\\sum_{s\\in G_n(\\hat{y})}C(s,\\hat{y})} $$note that we have replaced $\\min (C(s,\\hat{y}), C(s,y))$ with $\\min (C(s,\\hat{y}), \\max_{y\\in S_i}C(s,y))$ since there are multiple reference translation, we use the most similar one. So this is to say: \u0026ldquo;There are multiple answer, go to choose the best one and compute the score.\u0026rdquo;\nBP (Brevity Penalty) Candidate translations longer than their references are already penalized by the modified $n$-gram precision measure. Now, to penalize those translations shorter than the reference translations, we need add an penalty term. This is when brevity penalty comes out:\n$$ \\mathrm{BP}=\\begin{cases}1\u0026\\text{ if }c \u003e r\\\\ \\exp(1-\\frac{r}{c})\u0026\\text{ if }c \\leq r \\end{cases} $$where\n$c$ is the number of words or tokens of the candidate corpus. That is, $$ c = \\sum_{i=1}^M\\mathrm{length}(\\hat{y}^i) $$ $r$ is the number of words or tokens of the effective reference corpus length, where the effective reference is defined as the reference translation whose length is as close to the corresponding candidate translation as possible. That is $$ r = \\sum_{i=1}^M\\mathrm{length}(y^{i,j}),\\text{ where } y^{i,j}=\\arg\\min_{y\\in S_i}|\\mathrm{length}(\\hat{y}^i)-\\mathrm{length}(y)| $$with this penalty term, we wish the model to output the translations with the same length as the reference translations.\nWeight The weight measures the importance of different precision score, in the original paper, the uniform weights are adopted, that is\n$$ w_i = \\frac{1}{N}, \\ \\text{ for } i=1,\\dots,N $$Final definition The final definition of the BLEU is given by\n$$ \\mathrm{BLEU}_ {w}(\\hat{S}, S) = \\mathrm{BP} \\cdot \\exp\\left(\\sum_{n=1}^\\infty w_n\\log p_n(\\hat{S}, S) \\right) $$usually, the upper-bound of the above summation can be reduced to $\\max_{i=1,\\dots,M}\\mathrm{length}(\\hat{y}^i)$.\nAnalysis Disadvantages of BLEU:\nBLEU compares overlap in tokens from the predictions and references, instead of comparing meaning. This can lead to discrepancies between BLEU scores and human ratings. BLEU scores are not comparable across different datasets, nor are they comparable across different languages. BLEU scores can vary greatly depending on which parameters are used to generate the scores, especially when different tokenization and normalization techniques are used. BLEU ignores synonym or similar expression, which causes refuses of reasonable translation. BLEU is affected by common words. Python Implementation 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 import math from typing import Set, List def compute_n_gram_set(s: str, n: int) -\u0026gt; Set[str]: return set(s[i : i + n] for i in range(len(s) - n)) def modified_n_gram_precision(S_hat: List[str], S: List[List[str]], n: int) -\u0026gt; float: if len(S_hat) \u0026lt; n: return 0.0 numerator: int = 0 denominator: int = 0 for index, y_hat in enumerate(S_hat): n_gram_set = compute_n_gram_set(y_hat, n) if not n_gram_set: continue # print(n_gram_set) for n_gram in n_gram_set: candidate_substr_count = y_hat.count(n_gram) best_ref_substr_count, best_ref_len = max( [(y.count(n_gram), len(y)) for y in S[index]], key=lambda x: x[0] ) numerator += min(candidate_substr_count, best_ref_substr_count) denominator += candidate_substr_count return numerator / denominator def brevity_penalty(S_hat: List[str], S: List[List[str]]) -\u0026gt; float: r: int = 0 c: int = 0 for index, y_hat in enumerate(S_hat): c += len(y_hat) best_match_ref = min(S[index], key=lambda x: abs(len(x) - len(y_hat))) r += len(best_match_ref) return 1.0 if c \u0026gt; r else math.exp(1 - r / c) def compute_bleu_score(S_hat: List[str], S: List[List[str]]) -\u0026gt; float: assert len(S_hat) == len(S) # take N as a sufficiently large integer # N = max(len(y_hat) for y_hat in S_hat) N = sum(len(y_hat) for y_hat in S_hat) bp = brevity_penalty(S_hat, S) precisions = [modified_n_gram_precision(S_hat, S, n) for n in range(1, N + 1)] # bleu_score = bp * exp(p_n) return bp * math.exp( sum(math.log(precision) for precision in precisions if precision != 0) ) Reference Hugging Face space Original Paper Wikipedia Documentation, Recommended ","date":"2024-04-25T22:46:53+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/bleu-bilingual-evaluation-understudy/","title":"BLEU (Bilingual Evaluation Understudy)"},{"content":"Given a string consisting of lower-case characters, find the longest subsequence such that the distance between adjacent characters in the subsequence are less than a given threshold.\nIntuition Same as the longest increasing subsequence, we can use dynamic programming to solve this problem.\nApproach We use dp[i] to represent the longest ideal subsequence that ended with s[i - 1]. We then update the dp[i] with the property of ideal subsequence: $$ dp[i] = \\max_{j=1,\\dots,i-1, \\mathrm{abs}(s[i - 1]-s[j - 1])\\leq k}(dp[i],\\ dp[j] + 1) ,\\ i=1,\\dots,n$$However, it turns out that the above solution is of complexity $O(n^2)$, which leads to Time Exceed Limit, so we need to optimize it.\nNow, consider the property of ideal sequence, we only care about those characters that is within the range (s[i - 1] - k, s[i - 1] + k). So, we can use a map record, whose key is all lowercase characters, to remember the result that is used to update dp[i], that is, for given index i:\n$$ record[l] = \\max_{j=1,\\dots,i - 1, s[j] - 'a' = l}dp[j] $$ in this way, we have $$ dp[i] = \\max_{\\mathrm{abs}((s[i]-'a') - l)\\leq k}(dp[i],\\ record[l] + 1),\\ i=1,\\dots,n $$ notice that len(record)=26, so the complexity now reduces to $O(n)$.\nComplexity Time complexity: iterate all characters once, each iteration queries record at most $26$ times. $$O(n)$$ Space complexity: stores the dp array, can be optimized to $O(1)$ by update result at each iteration. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class Solution { public: int longestIdealString(string s, int k) { int n = s.size(); vector\u0026lt;int\u0026gt; dp(n + 1); vector\u0026lt;int\u0026gt; record(26); for (int i = 0; i \u0026lt; n; ++i) { dp[i + 1] = 1; int index = s[i] - \u0026#39;a\u0026#39;; for (int j = 0; j \u0026lt; 26; ++j) { if (abs(index - j) \u0026lt;= k) { dp[i + 1] = max(dp[i + 1], record[j] + 1); } } record[index] = max(record[index], dp[i + 1]); } int result = 0; for (int i = 0; i \u0026lt; n + 1; ++i) { result = max(result, dp[i]); } return result; } }; ","date":"2024-04-25T20:34:06+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/2370.-longest-ideal-subsequence/","title":"2370. Longest Ideal Subsequence"},{"content":"Compute the $n$-th tribonacci number.\nIntuition Same as compute the $n$-th fibonacci number, we use three numbers to remember the state.\nApproach We use three numbers to represent $n-2$, $n-1$ and $n$-th tribonacci number respectively\nComplexity Time complexity: $$O(n)$$ Space complexity: $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 class Solution { public: int tribonacci(int n) { vector\u0026lt;int\u0026gt; nums{0, 1, 1}; if (n \u0026lt; 3) return nums[n]; for (int i = 2; i \u0026lt; n; ++i) { int temp = nums[2]; nums[2] += nums[0] + nums[1]; nums[0] = nums[1]; nums[1] = temp; } return nums[2]; } }; ","date":"2024-04-24T18:53:26+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1137.-n-th-tribonacci-number/","title":"1137. N-th Tribonacci Number"},{"content":"Given a tree, reorganize the tree such that the height of the tree is minimized.\nIntuition We construct the tree from bottom to top, util we find the root of the tree\nApproach We use topological sort to order the nodes of the tree, then we iteratively construct the tree from bottom to top with BFS.\nWe use a different stop criteria to avoid missing possible solutions.\nThe result contains at most two possible roots, since if there are three, then the degree of one node must be lower than the other two nodes.\nComplexity Time complexity: iterate all nodes once. $$O(n)$$ Space complexity: store the adjacent lists. $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 class Solution { public: vector\u0026lt;int\u0026gt; findMinHeightTrees(int n, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; edges) { // boundary check if (n == 1) return vector\u0026lt;int\u0026gt;{0}; vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; adjs(n); vector\u0026lt;int\u0026gt; in_degrees(n); for (const auto \u0026amp;edge: edges) { adjs[edge[0]].push_back(edge[1]); adjs[edge[1]].push_back(edge[0]); ++in_degrees[edge[0]]; ++in_degrees[edge[1]]; } // leaf nodes queue\u0026lt;int\u0026gt; q; for (int i = 0; i \u0026lt; n; ++i) { if (in_degrees[i] == 1) q.push(i); } while (n \u0026gt; 2) { // construct next layer int size = q.size(); n -= size; for (int i = 0; i \u0026lt; size; ++i) { int node = q.front(); for (int next_node: adjs[node]) { if (--in_degrees[next_node] == 1) { q.push(next_node); } } q.pop(); } } vector\u0026lt;int\u0026gt; result; while (!q.empty()) { result.push_back(q.front()); q.pop(); } return result; } }; ","date":"2024-04-23T21:02:07+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/310.-minimum-height-trees/","title":"310. Minimum Height Trees"},{"content":"Given a four-digit string, change one digit (plus or minus 1) at a time, find the minimum number of steps to go from the source to target without passing through invalid states.\nIntuition We can image this as a graph path finding problem, where we need to find a path from the source to the target with the minimum number of steps.\nApproach We use BFS to solve this problem. The graph is constructed as follows: each possible state is a node of the graph, such as \u0026quot;1234\u0026quot;, \u0026quot;4567\u0026quot;, each operation defines an edge between two nodes, for example, we can rotate third digit of \u0026quot;1234\u0026quot; to obtain \u0026quot;1244\u0026quot;, since there are two possible directions and four digits, each node has $2^4=16$ adjacent nodes. We keep track of visited nodes and add them to deadends since there are no difference between them.\nComplexity Time complexity: there are at most $10^4=10000$ nodes and we visited all nodes at most once. $$O(1)$$ Space complexity: there are at most $10^4=10000$ nodes. $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class Solution { public: vector\u0026lt;string\u0026gt; get_ajacent_nodes(const string \u0026amp;s) { vector\u0026lt;string\u0026gt; result; for (int i = 0; i \u0026lt; s.size(); ++i) { string s1 = s; string s2 = s; if (\u0026#39;1\u0026#39; \u0026lt;= s[i] \u0026amp;\u0026amp; s[i] \u0026lt;= \u0026#39;9\u0026#39;) { s1.replace(i, 1, string(1, s[i] + 1)); s2.replace(i, 1, string(1, s[i] - 1)); } else if (s[i] == \u0026#39;9\u0026#39;) { s1.replace(i, 1, \u0026#34;0\u0026#34;); s2.replace(i, 1, \u0026#34;8\u0026#34;); } else { s1.replace(i, 1, \u0026#34;1\u0026#34;); s2.replace(i, 1, \u0026#34;9\u0026#34;); } result.push_back(s1); result.push_back(s2); } return result; } int openLock(vector\u0026lt;string\u0026gt;\u0026amp; deadends, string target) { unordered_set\u0026lt;string\u0026gt; set_deadends(deadends.begin(), deadends.end()); queue\u0026lt;string\u0026gt; q; q.push(\u0026#34;0000\u0026#34;); int result = 0; while (!q.empty()) { int size = q.size(); for (int i = 0; i \u0026lt; size; ++i) { string node = q.front(); q.pop(); // dead end check if (set_deadends.find(node) != set_deadends.end()) continue; // target check if (node == target) return result; // mark as visited set_deadends.insert(node); const vector\u0026lt;string\u0026gt; \u0026amp;adjacent_nodes = get_ajacent_nodes(node); for (const string \u0026amp;adjacent_node: adjacent_nodes) { q.push(adjacent_node); } } ++result; } return -1; } }; ","date":"2024-04-22T19:01:23+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/752.-open-the-lock/","title":"752. Open the Lock"},{"content":"Meta released Llama3 at April 18, which is evaluated on several benchmarks and achieves the SOTA on open-sourced LLMs.\nIntroduction Instruct model performance The performance of Llama3 8B compared with Gemma and Mistral:\nModel Llama3 8B Gemma 7B - It Mistral \u0026amp;B Instruct MMLU (5 shot) 68.4 53.3 58.4 GPQA (0 shot) 34.2 21.4 26.3 HumanEval (0 shot) 62.2 30.5 36.6 GSM-8K (8 shot, CoT) 79.6 30.6 39.9 MATH (4 shot, CoT) 30.0 12.2 11.0 performance of Llama3 70B compared with Gemini Pro 1.5 and Claude Sonnet:\nModel Llama3 70B Gemini Pro 1.5 (Published) Claude 3 Sonnet (Published) MMLU (5 shot) 82.0 81.9 79.0 GPQA (0 shot) 39.5 41.5 (CoT) 38.5 (CoT) HumanEval (0 shot) 81.7 71.9 73.0 GSM-8K (8 shot, CoT) 93.0 91.7 (11 shot) 92.3 (0 shot) MATH (4 shot, CoT) 50.4 58.5 (Minerva prompt) 40.5 Pre-trained model performance The performance of Llama3 8B compared with Gemma and Mistral:\nModel Llama3 8B Gemma 7B (Published) Gemma 7B (Measured) Mistral 7B (Published) Mistral 7B (Measured) MMLU (5 shot) 66.6 64.3 64.4 62.5 63.9 AGIEval English (3-5 shot) 45.9 41.7 44.9 - 44.0 BIG-Bench Hard (3 shot, CoT) 61.1 55.1 59.0 - 56.0 ARC-Challenge (25 shot) 78.6 53.2(0 shot) 79.1 78.1 78.7 DROP (3 shot, F1) 58.4 - 56.3 - 54.4 performance of Llama3 70B compared with Gemini Pro 1.5 and Claude Sonnet:\nModel Llama3 70B Gemini Pro 1.0 (Published) Mixtral 8 $\\times$ 22B (Measured) MMLU (5-shot) 79.5 71.8 77.7 AGIEval English (3-5 shot) 63.0 - 61.2 BIG-Bench Hard (3 shot, CoT) 81.3 75.0 79.2 ARC-Challenge (25 shot) 93.0 - 90.7 DROP (3 shot, F1) 79.7 74.1 (variable shot) 77.6 Model Architecture Several improvements are made on Llama3 compared to llama2:\nLlama3 uses a tokenizer with a vocabulary of 128K tokens. Llama3 adopts grouped query attention (GQA) across both the 8B and 70B sizes. Llama3 uses to context window of size 8192 tokens Traning Llama3 uses 15T tokens for pre-training. Compares to Llama2, it is seven times larger and includes four times more code.\n5% data of the training dataset are non-English to support multi-lingual use cases.\nData processing includes:\nHeuristic filters NSFW filters Semantic deduplication approaches Text classifiers to predict data quality. Llama2 is used to generate training data for the text classifiers. Data mixing strategy is explored to improve the performance of Llama3.\nScaling up pretraining Llama3 developed a series of scaling laws for downstream benchmark evaluations.\nScaling laws help:\nSelect an optimal data mix and to make informed decisions on how to best use training compute. Scaling laws allow Llama3 to predict the performance of the largest models on key tasks without training the models. The authors finds our that the performance of the model continues to improve log-linearly as the training tokens increase. It is seen that Larger models can match the performance of these smaller models with less training compute, but smaller models are generally preferred because they are much more efficient during inference.\nThe authors combine three types of parallelization:\nData parallelization Model parallelization Pipeline parallelization Instruction fine-tuning The fine-tuning of Llama3 contains:\nSupervised fine-tuning Rejection sampling Proximal Policy Optimization Direct Preference Optimization Learning from perference rankings via PPO and DPO also greatly improved the performance of LLma3 on reasoning and coding tasks. Since perference ranking helps the model to select answer when it is in a dilemma.\nReference Llam3 blog Evaluation details Model Card ","date":"2024-04-22T16:22:19+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-llama3/","title":"Notes on Llama3"},{"content":"Find a available path from a given source to a given destination in a given graph.\nIntuition Use DFS to find all reachable nodes and check if the destination lie within those nodes.\nApproach We first transform the adjacency matrix to adjacency list to make BFS easier, then we use a queue to maintain the reachable nodes, to prevent from cycling, we also use a set to keep track of visited nodes. If at any point, we reach the destination, we return directly.\nComplexity Time complexity: iterate all nodes at most once $$O(n)$$ Space complexity: stores nodes of the same distance in the graph to the source $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 class Solution { public: bool validPath(int n, vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; edges, int source, int destination) { unordered_map\u0026lt;int, vector\u0026lt;int\u0026gt;\u0026gt; adjs; for (const auto \u0026amp;edge: edges) { adjs[edge[0]].push_back(edge[1]); adjs[edge[1]].push_back(edge[0]); } unordered_set\u0026lt;int\u0026gt; visited; queue\u0026lt;int\u0026gt; q; q.push(source); while (!q.empty()) { int node = q.front(); q.pop(); if (node == destination) return true; for (int next_node: adjs[node]) { if (visited.find(next_node) != visited.end()) continue; q.push(next_node); visited.insert(next_node); } } return false; } }; ","date":"2024-04-21T10:51:46+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1971.-find-if-path-exists-in-graph/","title":"1971. Find if Path Exists in Graph"},{"content":"Given a matrix where its grid component representing islands and forests, count the number of farmlands.\nIntuition Start from the top left coordinate of the farmland, Use DFS to find the bottom right coordinate of the farmland.\nApproach We define the possible directions as go_right and go_down respectively, we iterate all grids, if it is an grid of the farmland, then we use DFS to find the bottom right coordinate of the current farmland, and then we mark the farmland as visited and store the coordinates.\nComplexity Time complexity: iterate all grids once $$O(mn)$$ Space complexity: number of farmlands $$O(mn)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 class Solution { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dirs{{0, 1}, {1, 0}}; public: void dfs(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; land, int i, int j, vector\u0026lt;int\u0026gt; \u0026amp;bottom_right) { land[i][j] = 0; bool reach_end = true; for (const auto \u0026amp;dir: dirs) { int row = i + dir[0], col = j + dir[1]; if (0 \u0026lt;= row \u0026amp;\u0026amp; row \u0026lt; land.size() \u0026amp;\u0026amp; 0 \u0026lt;= col \u0026amp;\u0026amp; col \u0026lt; land[0].size() \u0026amp;\u0026amp; land[row][col] == 1) { reach_end = false; dfs(land, row, col, bottom_right); } } if (reach_end \u0026amp;\u0026amp; bottom_right.empty()) { bottom_right = vector\u0026lt;int\u0026gt;{i, j}; } } vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; findFarmland(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; land) { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; result; for (int i = 0; i \u0026lt; land.size(); ++i) { for (int j = 0; j \u0026lt; land[0].size(); ++j) { if (land[i][j] == 0) continue; vector\u0026lt;int\u0026gt; bottom_right; dfs(land, i, j, bottom_right); result.push_back(vector\u0026lt;int\u0026gt;{i, j, bottom_right[0], bottom_right[1]}); } } return result; } }; ","date":"2024-04-20T15:31:49+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/1992.-find-all-groups-of-farmland/","title":"1992. Find All Groups of Farmland"},{"content":"Given a matrix where its grid component representing islands and waters, count the number of Islands.\nIntuition Use DFS to find all connected components of the island, then count the number of islands.\nApproach We iterate all grid component, when we meet the land, we use DFS to find all connected components of the island, and mark those connected components as visited.\nComplexity Time complexity: visit all positions once. $$O(mn)$$ Space complexity: $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 class Solution { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dirs{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; public: void dfs(vector\u0026lt;vector\u0026lt;char\u0026gt;\u0026gt;\u0026amp; grid, int i, int j) { if (grid[i][j] == \u0026#39;0\u0026#39;) return; grid[i][j] = \u0026#39;0\u0026#39;; for (const auto \u0026amp;dir: dirs) { if (0 \u0026lt;= i + dir[0] \u0026amp;\u0026amp; i + dir[0] \u0026lt; grid.size() \u0026amp;\u0026amp; 0 \u0026lt;= j + dir[1] \u0026amp;\u0026amp; j + dir[1] \u0026lt; grid[0].size()) { dfs(grid, i + dir[0], j + dir[1]); } } } int numIslands(vector\u0026lt;vector\u0026lt;char\u0026gt;\u0026gt;\u0026amp; grid) { int count = 0; for (int i = 0; i \u0026lt; grid.size(); ++i) { for (int j = 0; j \u0026lt; grid[0].size(); ++j) { if (grid[i][j] == \u0026#39;0\u0026#39;) continue; dfs(grid, i, j); ++count; } } return count; } }; ","date":"2024-04-19T20:25:03+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/200.-number-of-islands/","title":"200. Number of Islands"},{"content":"Given a matrix, find the perimeter of the connected grid land\nIntuition Use DFS to find the island, then compute the perimeter.\nApproach We use result to store the result and find all connected components of the island with DFS, to compute the perimeter, to update result, we need to compute how many components that are connected the current component.\nNow, note that to prevent infinite recursion, we set grid[i][j] = -1 to mark it visited, then for each component, it may be in three states:\ngrid[i][j] = 0, it is water grid[i][j] = 1, it is a component of the island and being unvisited grid[i][j] = -1, it is a component of the island and has been visited. Complexity Time complexity: iterate all grids once. $$O(mn)$$ Space complexity: no extra space needed $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 class Solution { vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt; dirs{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; public: void dfs(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid, int i, int j, int \u0026amp;result) { if (grid[i][j] != 1) return; int num_edges = 4; grid[i][j] = -1; for (const auto\u0026amp;dir: dirs) { if (0 \u0026lt;= i + dir[0] \u0026amp;\u0026amp; i + dir[0] \u0026lt; grid.size() \u0026amp;\u0026amp; 0 \u0026lt;= j + dir[1] \u0026amp;\u0026amp; j + dir[1] \u0026lt; grid[0].size()) { if (grid[i + dir[0]][j + dir[1]] == 0) continue; --num_edges; dfs(grid, i + dir[0], j + dir[1], result); } } result += num_edges; } int islandPerimeter(vector\u0026lt;vector\u0026lt;int\u0026gt;\u0026gt;\u0026amp; grid) { int result = 0; for (int i = 0; i \u0026lt; grid.size(); ++i) { for (int j = 0; j \u0026lt; grid[0].size(); ++j) { if (grid[i][j] == 0) continue; dfs(grid, i, j, result); } } return result; } }; ","date":"2024-04-18T20:49:51+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/463.-island-perimeter/","title":"463. Island Perimeter"},{"content":"Introduction These advises are given by Patrick Riley in 2016, though it has been years util now, I think some of them are still useful.\nThe advice is organized into three general areas:\nTechnical: Ideas and techniques for how to manipulate and examine your data. Process: Recommendation on how you approach your data, what questions to ask, and what things to check. Social: How to work with others and communicate about your data and insights. Technical Look at your distribution Besides the typically used summary metrics, we should looking at a much richer representation of the distribution, such as histograms, CDFs, Q-Q plots, etc. This allows us to see some interesting features.\nConsider the outliers We should look at the outliers in our data. It\u0026rsquo;s fine to exclude them from our data or to lump them together into an unusual category, but we should make sure we know why.\nReport noise/ confidence Every estimator that you produce should have a notion of your confidence in this estimate attached to it.\nLook at examples Anytime you are producing new analysis code, you need to look at examples of the underlying data and how your code is interpreting those examples.\nSlice your data slicing help us obtain underlying features of the data subgroups easier. However, when we use slicing, we need to care about the mix shift.\nConsider practical significance Don\u0026rsquo;t be blind by statistics, watch out those may have an impact on deploying or ethical problems.\nCheck for consistency over time One particular slicing you should almost always employ is to slice by units of time. This is because many disturbances to underlying data happen as our systems evolve over time.\nProcess Separate Validation, description, and evaluation Description should be things that everyone can agree on from the data. Evaluation is likely to have much more debate because you imbuing meaning and value to the data. Confirm expt/data collection setup Before looking at any data, make sure you understand the experiment and data collection setup\nCheck vital signs Before actually answering the question you are interested in you need to check for a lot of other things that may not be related to what you are interested in but may be useful in later analysis or indicate problems in the data\nStandard first, custom second When we use metric, we should always look at standard metrics first, even if we expect them to change.\nMeasure twice, or more If you are trying to capture a new phenomenon, try to measure the same underlying thing in multiple ways. Then, check to see if these multiple measurements are consistent\nCheck for reproducibility Both slicing and consistency over time are particular examples of checking for reproducibility. If a phenomenon is important and meaningful, you should see it across different user populations and time.\nCheck for consistency with past measurements You should compare your metrics to metrics reported in the past, even if these measurements are on different user populations.\nNew metrics should be applied to old data/features first\nMake hypothesis and look for evidence Typically, exploratory data analysis for a complex problem is iterative. You will discover anomalies, trends, or other features of the data. Naturally, you will make hypotheses to explain this data. It’s essential that you don’t just make a hypothesis and proclaim it to be true. Look for evidence (inside or outside the data) to confirm/deny this theory.\nExploratory analysis benefits from end to end iteration When doing exploratory analysis, you should strive to get as many iterations of the whole analysis as possible.\nSocial Data analysis starts with questions, not data or a technique Ask question first and use tools to answer the questions.\nAcknowledge and count your filtering Acknowledge and clearly specify what filtering you are doing Count how much is being filtered at each of your steps the best way to do the latter is to actually compute all your metrics even for the population you are excluding Ratios should have clear numerator and denominators When you communicate results containing ratios, you must be clear about the numerator and denominator.\nEducate your consumers You are responsible for providing the context and a full picture of the data and not just the number a consumer asked for.\nBe both skeptic and champion As you work with data, you must be both the champion of the insights you are gaining as well as a skeptic.\nShare with peers first, external consumers second A skilled peer reviewer can provide qualitatively different feedback and sanity-checking than the consumers of your data can, especially since consumers generally have an outcome they want to get\nReference Practical advice for analysis of large, complex data sets ","date":"2024-04-17T22:40:11+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/practical-advice-for-analysis-of-large-complex-data-sets/","title":"Practical advice for analysis of large, complex data sets"},{"content":"Given a binary tree with value on each node representing a lowercase letter, find the lexicographically smallest string starting from the leaf to root node\nIntuition We use DFS to find all strings from the root node to the leaf nodes, then reverse the string and compare it withe the largest string.\nApproach We use current to represent the string starting from the root node to the current node and we use result to store the currently best result. When we reach the leaf node, we compare the current with result and update result.\nComplexity Time complexity: iterate all nodes once. $$O(n)$$ Space complexity: the string is related to the height of the tree. $$O(\\log n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 /** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: void dfs(TreeNode* root, string \u0026amp;current, string \u0026amp;result) { current.push_back(\u0026#39;a\u0026#39; + root-\u0026gt;val); // leaf node if (!root-\u0026gt;left \u0026amp;\u0026amp; !root-\u0026gt;right) { // update result reverse(current.begin(), current.end()); if (current \u0026lt; result) result = current; reverse(current.begin(), current.end()); } else { if (root-\u0026gt;left) dfs(root-\u0026gt;left, current, result); if (root-\u0026gt;right) dfs(root-\u0026gt;right, current, result); } current.pop_back(); } string smallestFromLeaf(TreeNode* root) { string result(8501, \u0026#39;z\u0026#39;); string current; dfs(root, current, result); return result; } }; ","date":"2024-04-17T21:15:03+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/988.-smallest-string-starting-from-leaf/","title":"988. Smallest String Starting From Leaf"},{"content":"concatenate the digit of path from the root to the leaf, and sum over the concatenated numbers.\nIntuition Use DFS to find all the paths, use a num variable to store the number concatenated, then sum over num.\nApproach We use num to represent the number from the root to the current node, if the current node is a leaf node, we add num to sum. Finally, we return sum as the result.\nComplexity Time complexity: iterate all nodes once. $$O(n)$$ Space complexity: $$O(1)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 /** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: void dfs(TreeNode *root, int \u0026amp;sum, int \u0026amp;num){ if (!root) { return; } num = num * 10 + root-\u0026gt;val; if (!root-\u0026gt;left \u0026amp;\u0026amp; !root-\u0026gt;right){ sum += num; }else{ dfs(root-\u0026gt;left, sum, num); dfs(root-\u0026gt;right, sum, num); } num /= 10; } int sumNumbers(TreeNode* root) { int sum = 0; int num = 0; dfs(root, sum, num); return sum; } }; ","date":"2024-04-14T21:06:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/129.-sum-root-to-leaf-numbers/","title":"129. Sum Root to Leaf Numbers"},{"content":"Given the root of a binary tree, return the sum of all left leaves.\nIntuition Use DFS to iterate over all nodes, if it is a left leaf, sum it to the result.\nApproach For every node, we care about one thing: whether its left child is a leaf node or not. If it is, then we add it.\nComplexity Time complexity: iterate all nodes once. $$ O(n) $$ Space complexity: Store the sum. $$ O(1) $$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 /** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: bool is_leaf_node(TreeNode* node) { if (!node) return false; return !(node-\u0026gt;left || node-\u0026gt;right); } int sumOfLeftLeaves(TreeNode* root) { if (!root) return 0; int result = 0; if (is_leaf_node(root-\u0026gt;left)) { result += root-\u0026gt;left-\u0026gt;val; } result += sumOfLeftLeaves(root-\u0026gt;right); result += sumOfLeftLeaves(root-\u0026gt;left); return result; } }; ","date":"2024-04-14T21:06:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/404.-sum-of-left-leaves/","title":"404. Sum of Left Leaves"},{"content":"Intuition Just use the transformation algorithm from decimal to hexadecimal\nApproach Simulation by doing the following:\ncompute remain = num % 16, add remain to the result (push_back) update num = (num - remainder) / 16 repeat step 1 and step 2 until num is 0 Notice that when num \u0026lt; 0, we need use its complement.\nComplexity Time complexity: the loop depends on the number of digits of num. $$O(\\log n)$$ Space complexity: $$ O(1) $$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 class Solution { public: string toHex(long num) { if (num == 0) return \u0026#34;0\u0026#34;; if (num \u0026lt; 0) num = INT_MAX + num + 2 + INT_MAX; string result; while (num) { int divide = num % 16; num = (num - divide) / 16; if (divide \u0026lt; 10) result += (\u0026#39;0\u0026#39; + divide); else result += (\u0026#39;a\u0026#39; + (divide - 10)); } reverse(result.begin(), result.end()); return result; } }; ","date":"2024-04-14T21:06:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/405.-convert-a-number-to-hexadecimal/","title":"405. Convert a Number to Hexadecimal"},{"content":"Intuition Just list all nodes with height height - 1 and insert a new layer with the given rules.\nApproach Use a queue to store all nodes with the same height, and use BFS to update the nodes, once we reach the height height - 1, we add a new layer with the given val for each node:\nCreate a new left child with the given (val, node-\u0026gt;left, nullptr) Create a new right child with the given (val, nullptr, node-\u0026gt;right) Complexity Time complexity: iterate all nodes of height height - 1 $$O(n)$$ Space complexity: storing all nodes of height height - 1 $$O(n)$$ Code 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 /** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: TreeNode* addOneRow(TreeNode* root, int val, int depth) { if (depth == 1) { TreeNode* new_root = new TreeNode(val, root, nullptr); return new_root; } queue\u0026lt;TreeNode*\u0026gt; q; q.push(root); --depth; while (--depth) { int size = q.size(); for (int i = 0; i \u0026lt; size; ++i) { auto node = q.front(); if (node-\u0026gt;left) q.push(node-\u0026gt;left); if (node-\u0026gt;right) q.push(node-\u0026gt;right); q.pop(); } } while (!q.empty()) { auto node = q.front(); node-\u0026gt;left = new TreeNode(val, node-\u0026gt;left, nullptr); node-\u0026gt;right = new TreeNode(val, nullptr, node-\u0026gt;right); q.pop(); } return root; } }; ","date":"2024-04-14T21:06:34+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/623.-add-one-row-to-tree/","title":"623. Add One Row to Tree"},{"content":"Problems of LLM Out of date knowledge: the model cannot gain knowledge after training Humiliation: the model may generate nonsense output Specific domain: the generalized model is difficult to adapt to specific domain Enthetic problems: the model may encounter Fine-tuning Fine-tuning is used to improve performance of foundation model on specific tasks with the help with some supervised data\nFine-tuning methods can be classified into:\nBased on range of updated parameters: Full Model fine-tuning: update the parameters of the whole model Partial fine-tuning: freeze the top layer; freeze the bottom layer Based on special technology: Adapter tuning LoRA Continual Learning fine-tuning Based on input: Instruction tuning Based on objective Multi-task fine-tuning Problems of fine-tuning:\nRequires task-specific labeled data, may cause overfitting and catastrophic forgetting。 The generalization ability is limited, and fine-tuning are required when adapting to new tasks The performance may be destroyed after fine-tuning, for example, safety. RAG RAG consists of three major processes of retrieval, augmentation, and generation. The framework of RAG in LLM can be described as follows:\nRetrieval Retriever type Retrieval methods can be generally categorized into two types: sparse and dense, based on the information encoding methods.\nsparse retrieval usually relies on inverted index matching along with the raw data input, for example TF-IDF and BM25. The limitation of sparse retrieval in RAG is its no-training nature, which makes the retrieval performance heavily rely on the quality of database construction and query generation. Moreover, such fixed term-based methods only support similarity retrieval, while cannot be adapted for other retrieval considerations demanding in LLM applications, such as the diversity dense retrieval, on the contrary, embeds the query and documents into continuous vector space with certain criteria, for example, semantic similarity. Examples include BERT, Dense Passage Retriever (DPR), etc. Retrieval Granularity Retrieval granularity denotes the retrieval unit in which the corpus is indexed, e.g., document, passage, token, or other levels like entity.\nChunk retrieval. Token retrieval. Entity retrieval. Pre-retrieval and Post-retrieval Enhancement Pre-retrieval and post retrieval strategies can be added to improve the quality of the retriever.\nPre-retrieval methods include:\nQuery rewrite. This method aims to close the gaps between the input text and the needed knowledge in retrieval, to reformulate the original question into a more conducive version to retrieve. Query augmentation. This method aims to combine the original query and the preliminary generated outputs as a new query, which is further used to retrieve relevant information from the external database Post-retrieval enhancement denotes the procedure to process the extracted top-k documents from the retriever before feeding them to the generator for the sake of better alignment between the retrieval and generation stages.\nDatabase Wikipedia Domain specific database search engine Generation Parameter-Accessible Generators (White-box). Allow parameter optimization. Parameter-Inaccessible Generators (Black-box). Focus more on retrieval and augmentation processes, trying to enhance the generator by augmenting the input with better knowledge, guidances or examples for the generation. Augmentation Input layer integration Output layer integration Intermediate layer integration Retrieval Frequency If it is necessary to retrieve? Self-RAG\nretrieval frequency:\nOne-time. Every-n-token Every token RAG training Training Free Independent training Sequential training Joint training Advance RAG Module RAG Applications NLP applications QA systems: REALM Chatbot: Fact Verification: self-RAG Downstream tasks: Recommendations Software engineering Domain-specific Applications AI for science Finance: ChatDOC Limitations of RAG Long Context Window Advantages of Long Context Window:\nImprove the understanding and relativity: long context window allows model to refer to more context information when generating answers. Handling complex tasks: long context window makes handling complex tasks such as writing a long article, coding Improve users\u0026rsquo; experience: the user expects the model remember the chat history and use them to interact with the user. Disadvantages of long context window:\nOnly uses context once, Requires refeeding the data to use long context window. Cost expensive due to input price. Time expensive due to limit of tokens per second. Needle in HayStack experiment show that there are problems with long context window. Advantages of RAG:\nPrivacy projection. Allow chunking the texts and retrieve the related information more accurately Adaptive to the size of data. Accepts multiple type of data source (multimodality). Only uses a small part of the total data, which is cheaper compared with long context window . problems of RAG\nThe quality of retrieval The retrieved text cannot be aligned with the queried text. The queried text are not retrieved all. Redundancy or out-dated data may cause inaccuracy. the quality of response generation Model Humiliation Irrelevance Organize the output to make it reasonable Depends on the external information Futures:\nTrustworthy RA-LLMs Multi-lingual RA-LLMs Multi-modal RA-LLMs Quality of External Knowledge Other technologies Query transformations Sentence window retrieval Fusion retrieval/ hybrid search multi-document agents Reference A Survey on RAG Meets LLMs: Towards Retrieval-Augmented Large Language Models ","date":"2024-04-14T12:38:04+08:00","image":"https://jiangyigithub.github.io/ai.github.io/agent_performance.png","permalink":"https://jiangyigithub.github.io/ai.github.io/p/notes-on-rag/","title":"Notes on RAG"},{"content":"In AI Ascent, Andrew Ng shared his opinion about the agentic workflow, specifically, his talk focused on designs and future works on agentic workflows.\nWhat is Agent LLM based agents can be classified into two categories: non-agentic and agentic.\nNon-agentic workflow requires LLM do something in a sequence of steps, usually this requires the \u0026ldquo;zero-shot\u0026rdquo; learning ability of LLM. Agentic workflow requires LLM do something step by step, and revise its action by thinking or reflection. Andrew then posts an image shows the agentic workflow can improve the performance of LLM on HumanEval benchmark.\nAgentic Design Reasoning Design Pattern Agentic design pattern can be divided into following kinds. Where the first two design patterns are robust and the latter two are emerging.\nReflection Self-Refine: Iterative Refinement with Self-Feedback Reflexion: Language Agents with Verbal Reinforcement Learning Tool Use Gorilla: Large Language Model Connected with Massive APIs MM-REACT: Prompting ChatGPT for Multimodal Reasoning and Action Planning Chain-of-Thought Prompting Elicits Reasoning in Large Language Models HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in Hugging Face Multi-agent collaboration Communicative Agents for Software Development AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation Future Perspective The set of tasks that AI can do expand dramatically because of agentic workflows. We have to get used to delegating tasks to AI agent and wait patiently for response. Fast token generation is important. Generating more tokens even from a a lower quality LLM can achieve a good result. Reference What\u0026rsquo;s next for AI agentic workflows ft. Andrew Ng of AI Fund ","date":"2024-04-14T12:38:04+08:00","image":"https://jiangyigithub.github.io/ai.github.io/p/whats-next-for-ai-agentic-workflows/agent_performance.png","permalink":"https://jiangyigithub.github.io/ai.github.io/p/whats-next-for-ai-agentic-workflows/","title":"What's next for AI agentic workflows"},{"content":"Best practice for machine learning.\nIntroduction Google posts a guide on how to uses machine learning in practice. It represents a style for machine learning, similar to Google C++ Style Guide.\nIn overview, to make great products:\nDo machine learning like the great engineer your are, not like the great machine learning expert you are.\nMost algorithms we are facing are engineering problems instead of machine learning algorithms. A basic approach Google recommends is:\nMake sure your pipeline is solid end to end Start with a reasonable objective Add a common-sense features in a simple way Make sure that your pipeline stays solid. Google separates rules with respect to different stages.\nBefore machine learning These rules help us understand whether the time is right for building a machine learning system.\nRule #1: Do not be afraid of lunch a product without machine learning.\nThis rule tells us that is not absolutely necessary, if rule-based methods work well, there is no need to develop a machine learning algorithm.\nRule #2: First, design and implement metrics.\nThis rule tells us that tracking as much as possible before we formalize what our machine learning system will do. This step helps us construct the goal of our system, that is, a metric.\nRule #3: Choose machine learning over a complex heuristic.\nConsidering using machine learning algorithms only if the heuristic algorithm doesn\u0026rsquo;t work well, since a complex heuristic is not maintainable. Meanwhile, machine-learned models are easier to update and maintain.\nML phase 1: Your First Pipeline When creating our first pipeline, we should focus on our system infrastructure\nRule #4: Keep the first model simple and get the infrastructure right.\nRemember infrastructure issues are many more than model problems when we creating the first machine learning model. We have to determine:\nHow to obtain data for our model How to evaluate the performance of our model How to integrate our model into our application Moreover, we should choose simple features to ensure that:\nThe features reach our learning algorithm correctly The model learns reasonably weights The features reach our model in the sever correctly Once we have a system that does above things reliably, we have done most of the work.\nRule #5: Test the infrastructure independently from the machine learning.\nThis rule tells us split the infrastructure and the machine learning model and test them separately to avoid dependency issue occurs. In this way, our model can be developed without worrying about environment. Specifically:\nTest getting data into algorithm. Check the data that are feed into the models and do some statistics before using the data Testing getting models out of the training algorithm. Make sure our algorithms work in the same way in our serving environment as the training environment. Rule #6: Be careful about dropped data.\nDo not drop data without thoroughly test, this may data loss.\nRule #7: Turn heuristics into features, or handle them externally.\nBefore using machine learning models, if we have tried some heuristic algorithms, then these algorithms may help us to improve the overall performance. Some ways we can use an existing heuristic algorithm:\nPreprocessing using the heuristic. If the feature is incredibly awesome, then do not try to relearn the feature. Just use the heuristic way to pre-process the data. Create a feature. We can use the heuristic way to create a new feature to help improve the machine learning performance. Mine the raw inputs of the heuristic. We can use the inputs of heuristic as features to learn the heuristic implicitly. Modify the label. Monitoring In general, such as making alerts and having a dashboard page.\nRule #8: Know the freshness requirements of our system.\nIt is important for us to know the freshness of our model, for example, how much does performance degrade if we have a model that is a day old. The freshness helps us monitor and improve the performance.\nRule #9: Detect problems before exporting models.\nThis rule tells us that evaluating the performance of the model before serving. The evaluation includes the testing on hold-out data, check AUC metric.\nRule #10: Watch for silent failures.\nSince the continuos change of data, silent failures may occur, so keep tracking statistics of the data as well as manually inspect the data on occasion help us reduce these kind of issues.\nRule #11: Give feature owners and documentation.\nKnowing who created the feature helps us gain information about data. A detailed documentation helps user understand how it works.\nYour first objective Rule #12: Don\u0026rsquo;t overthink which objective you choose to optimize.\nThere are many metrics to optimize according to Rule #2. However, it turns out in early stage, some metrics are optimized even though we not directly optimizing them.\nRule #13: Choose simple, observable and attributable metric for your first objective.\nThis rule tells us the strategy of choosing metric in the beginning. In principal, The ML objective should be something that is easy to measure and is a proxy for the \u0026ldquo;true\u0026rdquo; objective. In fact however, there is no such \u0026ldquo;true\u0026rdquo; objective, so we should keep the objective as simple as possible, it\u0026rsquo;s better if the objective is observable. Then, we can modify the objective based on the performance of the model.\nRule #14: Starts with an interpretable model makes debugging easier.\nRule #15: Separate spam filtering and quality ranking in a policy layer.\nSometimes, spam filtering confuses quality ranking, when we do quality ranking, we should clean the data.\nML phase 2: Feature engineering After we have a working end to end system with unit and system tests instrumented, Phase II begins. In this phase, we should make use of features.\nRule #16: Plan to lunch and iterate.\nThere are three reasons to lunch new models:\nYou are coming up with new features You are tuning regularization and combining old features in new ways You are tuning the objectives When lunch a new model, we should think about:\nHow easy is it to add or remove or recombine features How easy is it to create a fresh copy of the pipeline and verify its correctness. Is it possible to have two or three copies running in parallel. Rule #17: Start with directly observed and reported features as opposed to learned features.\na learned feature is a feature generated by an external system or by the learner itself. If the learned feature comes from an external system, then bias or being out-of-date may affect the model. If the learned feature comes from the learner itself, then it is hard to tell the impact of the feature.\nRule #18: Explore with features of content that generalize across contexts.\nRule #19: Use very specific features when you can.\nIt is simpler to learn millions of simple features than a few complex features.\nRule #20: Combine and modify existing features to create new features in human-understanding ways.\nCombine features may causing overfitting problems\nRule #21: The number of feature weights we can learn in a linear model is roughly proportional to the number of data you have.\nRule #22: Clean up features you are no longer using.\nIf you find that you are not using a feature, and that combining it with other features is not working, then drop it out of your infrastructure.\nHuman analysis of the system This subsection teaches us how to look at an existing model and improve it.\nRule #23: You are not a typical end user.\nCheck carefully before we deploying the model.\nRule #24: Measure the delta between models\nMake sure the system is stable when making small changes. Make sure that a model when compared with itself has a low (ideally zero) symmetric difference.\nRule #25: When choosing models, utilitarian performance trumps predictive power.\nRule #26: Look for patterns in the measured errors, and create new features.\nOnce you have examples that the model got wrong, look for trends that are outside your current feature set.\nRule #27: Try to quantify observed undesired behavior\nIf your issues are measurable, then you can start using them as features, objectives, or metrics. The general rule is \u0026ldquo;measure first, optimize second\u0026rdquo;.\nRule #28: Be aware that identical short-term behavior does not imply identical long-term behavior.\nTraining-Serving skew Training-serving skew is a difference between performance during training and performance during serving. This skew can be caused by:\nA discrepancy between how you handle data in the training and serving pipelines A change in the data between when you train and when you serve A feedback loop between your model and your algorithm. The best solution is to explicitly monitor it so that system and data changes don\u0026rsquo;t introduce skew unnoticed.\nRule #29: The best way to make sure that you train like you serve is to save the set of features used at the serving time, and then pipe those features to a log to use them at training time.\nThis can help verify the consistency between the training and serving.\nRule #30: Importance-weight sampled data, don\u0026rsquo;t arbitrarily drop it.\nImportance weighting means that if you decide that you are going to sample example X with a 30% probability, then give it a weight of 10/3. With importance weighting, all of the calibration properties discussed in Rule #14 still hold.\nRule #31: Beware that if your join data from a table at training and serving time, the data in the table may change.\nRule #32: Re-use code between your training pipeline and your serving pipeline whenever possible.\nRule #33: If you produce a model based on the data until January 5th, test the model on the data from January 6th and after.\nIn general, measure performance of a model on the data gathered after the data you trained the model on, as this better reflects what your system will do in production\nRule #37: Measure Training-Serving Skew.\nWe can divide causes of Training-Serving Skew into several parts:\nThe difference between the performance on the training data and the holdout data. In general, this will always exist, and it is not always bad. The difference between the performance on the holdout data and the \u0026ldquo;next­day\u0026rdquo; data. Again, this will always exist The difference between the performance on the \u0026ldquo;next-day\u0026rdquo; data and the live data. ML phase 3: Slowed growth, Optimization refinement, and complex models Rule #38: Don’t waste time on new features if unaligned objectives have become the issue.\nRule #39: Launch decisions are a proxy for long-term product goals.\nThe only easy launch decisions are when all metrics get better (or at least do not get worse).\nIndividuals, on the other hand, tend to favor one objective that they can directly optimize.\nRule #40: Keep ensembles simple.\nTo keep things simple, each model should either be an ensemble only taking the input of other models, or a base model taking many features, but not both.\nRule #41: When performance plateaus, look for qualitatively new sources of information to add rather than refining existing signals.\nAs in any engineering project, you have to weigh the benefit of adding new features against the cost of increased complexity.\nRule #42: Don’t expect diversity, personalization, or relevance to be as correlated with popularity as you think they are.\nRule #43: Your friends tend to be the same across different products. Your interests tend not to be.\nConclusion In early stage, make sure that the infrastructure is well constructed, the used model can be simple.\nIn main stage, focusing on the utilitarian performance and the gap between training data and test data.\nWhen utilizing features, use simple, observable features.\nWhen deploying models, watch out training-serving skew.\nReference Rules of Machine Learning ","date":"2024-04-13T19:57:47+08:00","permalink":"https://jiangyigithub.github.io/ai.github.io/p/rules-of-machine-learning/","title":"Rules of Machine Learning"},{"content":"Problem description When we use conda create --env myenv to create a new environment, it may throw an exception\n1 CondaHTTPError: HTTP 000 CONNECTION FAILED for url \u0026lt;https://repo.anaconda.com/pkgs/main/osx-64/repodata.json\u0026gt; usually, this occurs when the source channels are not domestic. In this case, we need the change the configuration.\nSolution use conda info to view the configuration file: 1 conda info focus on the user config file and open it. 2. use the following configuration to overwrite the file:\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ssl_verify: false channels: - defaults show_channel_urls: true default_channels: - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2 custom_channels: conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud pytorch-lts: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud deepmodeling: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/ ","date":"2023-08-25T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/how-to-fix-http-error-when-creating-a-new-environment./","title":"How to fix http error when creating a new environment."},{"content":"(4) # 数字 (4,) # 元组\nlogits 是模型对每个类别的“打分”\nsoftmax 才把打分变成概率\nPyTorch 要用元组 .item() 把只有一个元素的 tensor 变成普通 Python 数字\ntorch.randn(4,) # 一维 torch.randn(4, 10) # 二维 torch.randn(2, 3, 5) # 三维\nCrossEntropyLoss(logits, target) = - log_softmax(logits)[target] LogSoftmax = logits − LogSumExp LogSoftmax(z)=z−LogSumExp(z)\ntorch.randint(low, high, size)\n1 2 3 4 5 6 7 8 9 import torch def cross_entropy_loss(logits, targets): log_probs = logits - torch.logsumexp(logits, dim=-1, keepdim=True) return -log_probs[torch.arange(targets.shape[0]), targets].mean() logits = torch.randn(4, 10) targets = torch.randint(0, 10, (4,)) print(\u0026#39;Loss:\u0026#39;, cross_entropy_loss(logits, targets).item()) print(\u0026#39;Ref: \u0026#39;, torch.nn.functional.cross_entropy(logits, targets).item()) 好，我们不用公式，直接用一个具体数字例子手算一遍，你马上就懂了。\n🎯 假设有一个 3 分类问题 某一个 batch 里有 2 个样本。\n模型输出（logits）是：\n1 2 3 logits = [[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]] 真实标签：\n1 targets = [0, 1] 意思是：\n第 1 个样本 → 正确类别是 0 第 2 个样本 → 正确类别是 1 第一步：把 logits 变成概率（softmax） 第一个样本 1 [2.0, 1.0, 0.1] softmax 后大概是：\n1 [0.659, 0.242, 0.099] 意思是：\n类别 0 → 65.9% 类别 1 → 24.2% 类别 2 → 9.9% 因为真实标签是 0 所以我们只关心 0.659\n第二个样本 1 [0.5, 2.5, 0.3] softmax 后大概是：\n1 [0.106, 0.787, 0.107] 真实标签是 1 所以我们只关心 0.787\n第二步：取 log 交叉熵就是：\n真实类别概率越大 → loss 越小 真实类别概率越小 → loss 越大\n公式其实就是：\n1 loss = - log(真实类别的概率) 算一下：\n1 2 样本1: -log(0.659) ≈ 0.417 样本2: -log(0.787) ≈ 0.239 第三步：求平均 1 (0.417 + 0.239) / 2 = 0.328 这就是最终 loss。\n🔥 那你的代码在干嘛？ 它干了这 4 件事：\nlogits → 变成 log_softmax 只拿真实类别那一列 取负号 求平均 🧠 直觉理解 交叉熵本质就是一句话：\n“模型对正确答案的自信程度有多高？”\n如果正确类别概率 = 0.99 → loss 很小 如果正确类别概率 = 0.01 → loss 巨大 🎯 再用一句人话解释代码 1 return -log_probs[torch.arange(targets.shape[0]), targets].mean() 等价于\n1 2 for i in range(batch_size): result[i] = log_probs[i][targets[i]] 意思是：\n把每个样本正确类别的 log 概率拿出来， 变成负数，然后算平均。\n好问题 👍 “向量化索引”一开始确实容易晕，但其实有一条非常简单的记忆规则。\n我给你一个 一条核心口诀 + 三个固定套路，以后基本不会错。\n🌟 一条核心口诀 两个索引张量形状相同 → 就是“成对取坐标”\n1 A[rows, cols] 等价于：\n1 2 for i in range(len(rows)): A[rows[i], cols[i]] 一、交叉熵那句怎么记？ 1 log_probs[torch.arange(N), targets] 记住：\n1 2 arange(N) = 行号 targets = 列号 意思是：\n1 第 i 行，取 targets[i] 列 二、最常见的 3 种向量化规则 ✅ 规则 1：单维索引（选行） 1 A[rows] 等价于：\n1 A[rows, :] 例子：\n1 A[[0,2]] 选第 0 行和第 2 行。\n✅ 规则 2：两个索引 → 成对取元素 1 A[rows, cols] 如果：\n1 2 rows = [0,1,2] cols = [2,0,1] 那取：\n1 2 3 A[0,2] A[1,0] A[2,1] 记忆法：\n“拉链式匹配”\n✅ 规则 3：广播规则 如果一个是标量或长度1：\n1 A[rows, 1] 等价于：\n1 2 for i in rows: A[i, 1] 三、一个万能记忆图 想象：\n1 2 行索引 = 竖着的指针 列索引 = 横着的指针 当你写：\n1 A[rows, cols] 就是：\n1 rows[i] 和 cols[i] 组成一个坐标 四、交叉熵为什么必须这样？ 因为我们要做：\n[ \\log p_{y_i} ]\n也就是：\n1 2 第 i 个样本 取 第 y_i 类 向量化后变成：\n1 log_probs[arange(N), targets] 五、终极记忆公式 把这句背下来：\nA[x, y] → 取 (x[i], y[i]) 这些坐标点\n只要 x 和 y 是同形状向量，就是“配对取值”。\n六、如果你想更牢固记住 每次看到：\n1 A[rows, cols] 就在脑子里自动翻译成：\n1 2 for i in range(len(rows)): A[rows[i], cols[i]] 这样你永远不会错。\n如果你愿意，我可以再给你出 3 个小练习（非常短），你做完就彻底掌握向量化索引。\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"单层decoder layer 这就是一整层 Transformer decoder layer 的完整实现\n而且是“多模态 mixture 版本 + adaptive scaling”。\n1 \u0026#34;\u0026#34;\u0026#34;the usual norm + attn + res + norm + mlp + res\u0026#34;\u0026#34;\u0026#34; 我们逐块拆开讲清楚。\n🧠 一、整体结构（先给你大图） 标准 Transformer Decoder Layer：\n1 2 3 4 5 6 7 8 9 10 11 12 13 x ↓ LayerNorm ↓ Self-Attention ↓ Residual Add ↓ LayerNorm ↓ MLP ↓ Residual Add 你这段代码完全对应这个结构。\n🔍 二、逐段对应解释 🔹 Step 1：Pre-LN 1 2 3 4 5 6 7 hidden_states_input_norm[name] = mixtures[name].layer_func( \u0026#34;forward_norm\u0026#34;, layer_idx, \u0026#34;input_layernorm\u0026#34;, embeds_all[name], time_cond, ) 这就是：\n[ x_1 = LN(x) ]\n说明它是：\nPre-LN Transformer（现代主流结构）\n🔹 Step 2：Self-Attention 1 hidden_states_post_attn = forward_mixture_attn(...) 这就是我们刚分析过的：\n多模态 self-attention\n得到：\n[ Attn(LN(x)) ]\n🔹 Step 3：第一个 Residual 1 2 3 hidden_states_post_res[name] = ( residuals_pre_attn[name] + hidden_states_pre_res[name] ) 也就是：\n[ x_2 = x + Attn(LN(x)) ]\n这一步是：\nAttention 残差连接\n🔹 Step 4：第二个 LayerNorm 1 forward_norm(... \u0026#34;post_attention_layernorm\u0026#34;) 对应：\n[ x_3 = LN(x_2) ]\n🔹 Step 5：MLP 1 2 3 4 5 hidden_states_pos_mlp[name] = mixtures[name].layer_func( \u0026#34;mlp\u0026#34;, layer_idx, hidden_states_pre_mlp[name], ) 就是：\n[ MLP(LN(x_2)) ]\n🔹 Step 6：第二个 Residual 1 2 3 hidden_states_final[name] = ( residuals_pre_post_attn[name] + hidden_states_pre_final_res[name] ) 也就是：\n[ x_4 = x_2 + MLP(LN(x_2)) ]\n🎯 三、你这层完整公式 合起来：\n[ x_1 = LN(x) ]\n[ x_2 = x + Attn(x_1) ]\n[ x_3 = LN(x_2) ]\n[ x_4 = x_2 + MLP(x_3) ]\n完全标准 Transformer。\n🧩 四、特殊点：adaptive_scale 你这里多了：\n1 forward_adaptive_scale(...) 这个通常用于：\n条件调制（conditioning）\n结合参数：\n1 time_cond 说明这层支持：\ndiffusion-like time conditioning 或控制任务的时间步调制 或 FiLM-style scaling 通常实现形式类似：\n[ y = \\alpha(t) \\cdot x + \\beta(t) ]\n这不是标准 Transformer，而是增强版。\n🧩 五、Mixture 的意义 这里每个 name 代表一个模态：\n1 2 3 vlm proprio action 它对每个模态：\n单独 LN 单独 MLP 单独 residual 但：\nAttention 是拼接后统一算的\n这就是“共享注意力 + 独立 FFN”的结构。\n18层 decoder 现在我们把你前面所有代码串起来，完整走一遍 18 层的计算流程。\n你这段代码是整个 Transformer backbone 的主循环：\n1 2 for layer_idx in range(self.num_hidden_layers): # 18 layers embeds_all = forward_mixture_layers(...) 我们一步一步拆开。\n🧠 一、整体大结构（系统视角） 你的模型整体是：\n1 2 3 4 5 6 7 8 9 image / text / proprio / action ↓ embedding + position id ↓ 18 层 mixture transformer ↓ action decoder (linear) ↓ trajectory 它是：\n🔥 多模态 Decoder-only Transformer 🔥 18 层堆叠 🔥 支持 KV cache 🔥 支持 time conditioning 🔥 支持 final layer 特殊跳过逻辑\n🏗 二、进入第 0 层之前发生了什么？ 在进 18 层前，你已经有：\n1 2 3 4 5 embeds_all = { \u0026#34;vlm\u0026#34;: image_text_embeds, \u0026#34;proprio\u0026#34;: proprio_embeds, \u0026#34;action\u0026#34;: action_embeds, } 形状：\n1 [Batch, SeqLen, Hidden] 同时你有：\nattention_mask（causal mask） position_ids_all（每个模态自己的 position id） kv_caches（推理时缓存） 🔁 三、18 层循环发生什么？ 循环：\n1 for layer_idx in range(18): 每一层都调用：\n1 forward_mixture_layers(...) 而我们前面已经拆解过它。\n🔬 四、单层内部的完整计算流程 每一层都做：\n🔹 1️⃣ 每个模态做 Input LayerNorm 对：\n1 2 3 vlm proprio action 分别：\n[ x_i^{norm} = LN(x_i) ]\n🔹 2️⃣ 所有模态拼接做 Attention 内部发生：\n1 2 3 4 5 concat(vlm, proprio, action) ↓ multi-head self-attention ↓ 再按模态切回 得到：\n[ Attn(LN(x)) ]\n注意：\n是 causal attention 支持 kv cache 支持不同模态位置编码 🔹 3️⃣ Attention 残差 [ x = x + Attn(x) ]\n🔹 4️⃣ Post-Attn LayerNorm [ x = LN(x) ]\n🔹 5️⃣ 每个模态自己的 MLP 每个 mixture 有自己的 FFN：\n[ MLP_{vlm} ] [ MLP_{proprio} ] [ MLP_{action} ]\n🔹 6️⃣ MLP 残差 [ x = x + MLP(x) ]\n🔹 7️⃣ adaptive_scale（可选） 如果有：\n1 time_cond 那每层都会：\n[ x = \\alpha(t)x + \\beta(t) ]\n说明：\n这是条件 transformer（可能是 diffusion policy 或时间调制控制模型）\n🧩 五、Final Layer 特殊逻辑 你这里有：\n1 post_attn_skip_names = final_layer_post_attn_skip_names if is_final_layer else [] 意思是：\n最后一层可能跳过某些模态的 FFN 或 residual\n例如：\n1 final_layer_post_attn_skip_names = [\u0026#34;vlm\u0026#34;] 那：\n最后一层 vlm 不做 MLP/residual 只保留 action 分支 这通常用于：\n🔥 只输出 action 🔥 不再更新视觉或语言 token\n是一种节省计算 + 控制输出结构的设计。\n🧠 六、18 层堆叠后的状态 经过 18 层后：\n1 2 3 4 5 embeds_all = { \u0026#34;vlm\u0026#34;: ..., \u0026#34;proprio\u0026#34;: ..., \u0026#34;action\u0026#34;: ... } 这时：\naction embedding 已经融合了\nimage 信息 text 信息 state 信息 历史 action 信息 🎯 七、最后 Action Decoder 你之前给过：\n1 self.action_decoder = nn.Linear(hidden_size, action_dim) 于是：\n[ a = W h + b ]\n把：\n1 [Batch, SeqLen, Hidden] 映射成：\n1 [Batch, SeqLen, action_dim] 得到 trajectory。\n🔥 八、整个计算流程总结（从输入到输出） 完整过程：\n🔹 Step 0：Embedding 1 2 3 4 image → vision encoder → embed text → LLM embed proprio → linear action → embed 🔹 Step 1：位置编码 每个模态加：\n1 position_ids_all 🔹 Step 2：18 层 Transformer 每层：\n1 2 3 4 5 6 7 8 9 10 11 12 13 LN ↓ Mixture Attention (全模态交互) ↓ Residual ↓ LN ↓ Mixture MLP (各模态独立) ↓ Residual ↓ Adaptive Scale 堆叠 18 次。\n🔹 Step 3：只取 action 分支 最终：\n1 embeds_all[\u0026#34;action\u0026#34;] 🔹 Step 4：Linear 输出 1 action_decoder 输出：\n1 trajectory transformer 函数库 很好 👍 这段代码其实就是一个 支持 AdaLN-Zero + LoRA + GQA 的 Transformer Decoder Block 实现，而且是为 multi-mixture（多模态拼接）设计的。\n我给你分 4 层讲清楚：\n整体架构在干嘛 Mixture 是什么 MixtureDecoderLayer 是什么 MixtureAttention 里面每个函数在干嘛（重点） 一、整体结构在干嘛？ 结构是：\n1 2 3 4 5 6 Mixture ├── MixtureDecoderLayer × N (比如 18 层) ├── Self Attention ├── MLP ├── Norm / AdaLN ├── Adaptive LayerScale 本质就是一个：\n带条件调制（cond）的 Transformer Decoder\n用于：\nVLM 阶段 Flow Matching 阶段 diffusion-like 训练 二、Mixture 类是干嘛的？ 1 2 3 self.layers = nn.ModuleList( [MixtureDecoderLayer(config) for _ in range(config.num_hidden_layers)] ) 就是堆叠 N 层 decoder。\n🔹 layer_func 1 2 def layer_func(self, method_name, layer_idx, *args): return getattr(self.layers[layer_idx], method_name)(*args) 意思是：\n调用某一层的某个函数\n例如：\n1 mixture.layer_func(\u0026#34;forward_norm\u0026#34;, 3, x, cond) 就等价于：\n1 self.layers[3].forward_norm(x, cond) 🔹 attn_func 1 return getattr(self.layers[layer_idx].self_attn, method_name)(*args) 专门调用某一层 attention 的函数。\n比如：\n1 mixture.attn_func(\u0026#34;forward_q_proj\u0026#34;, 5, x) 就是：\n1 self.layers[5].self_attn.forward_q_proj(x) 这个设计是为了：\n方便 multi-mixture 控制不同 block\n三、MixtureDecoderLayer 是什么？ 就是一个标准 transformer block：\n1 2 x → LN → Attention → residual → LN → MLP → residual 但它支持：\n普通 RMSNorm 或 AdaptiveRMSNorm 或 adaLN-Zero 🔥 Adaptive 模式 1 self.adaptive_mode = config.get(\u0026#34;adaptive_mode\u0026#34;, None) 如果开了 adaptive_mode：\nLayerNorm 会变成：\n1 AdaptiveRMSNorm(hidden_size, time_hidden_size) 意味着：\nNorm 受时间 embedding 或条件 embedding 调制\n🔥 adaLN-Zero 的额外设计 1 2 self.post_adaptive_scale self.final_adaptive_scale 这是 diffusion 里常见的：\n在 residual 之前加一个 learnable scale\n数学上是：\n[ x = x + scale(cond) * F(x) ]\n优点：\n初始 scale 接近 0 训练更稳定 防止 early exploding 四、MixtureAttention（重点） 这就是一个支持：\nGQA（Grouped Query Attention） RoPE LoRA QLoRA 的 attention。\n1️⃣ 初始化 1 2 3 self.num_heads self.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads 说明它支持：\nGQA（Grouped Query Attention） 比如：\n1 2 num_heads = 16 num_kv_heads = 4 那么：\n1 2 16 query heads 4 key/value heads 每个 KV head 被 4 个 Q head 共享。\n2️⃣ Q/K/V 投影 forward_q_proj 1 query_states = self.q_proj(x) shape 变化：\n1 2 3 4 5 [B, L, hidden] → [B, L, num_heads * head_dim] → [B, num_heads, L, head_dim] 通过：\n1 .view(...).transpose(1,2) forward_k_proj / forward_v_proj 类似，但：\n1 num_key_value_heads 而不是 num_heads。\n3️⃣ Rotary Embedding 1 self.rotary_emb = GemmaRotaryEmbedding(...) forward_rotary_emb 返回：\n1 cos, sin 然后：\n1 apply_rotary_pos_emb 应用到：\n1 Q 和 K 4️⃣ repeat_kv（关键） 因为是 GQA：\nQ 头数 \u0026gt; KV 头数\n所以必须：\n1 key_states = repeat_kv(key_states, self.num_key_value_groups) 举例：\n1 2 3 num_heads = 16 num_kv_heads = 4 groups = 4 repeat 后：\n1 4 → 16 让 KV 头数和 Q 头数一致。\n5️⃣ forward_o_proj 把：\n1 [B, L, num_heads * head_dim] 映射回：\n1 [B, L, hidden_size] 五、整体 Attention 流程（完整图） 假设输入：\n1 x: [B, L, hidden] 流程：\n1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 Q = q_proj(x) → reshape → [B, Hq, L, D] K = k_proj(x) → reshape → [B, Hkv, L, D] V = v_proj(x) → reshape → [B, Hkv, L, D] apply RoPE(Q, K) repeat_kv(K,V) → [B, Hq, L, D] Attention: softmax(QK^T / sqrt(d)) V → [B, Hq, L, D] reshape → [B, L, hidden] o_proj → [B, L, hidden] 六、这个类在 multi-mixture 里的意义 关键点：\n它本身并不知道“vlm/action”。\n它只负责：\n对输入的那一段 token 做 attention。\nmulti-mixture 框架会：\n拼接不同模态 统一算 attention 再 split 回去 七、和你前面问的问题连接起来 你之前问：\naction token 为什么最后有全局语义？\n答案就在这里：\n因为：\n1 2 3 Q_action attend to K/V_vlm + K/V_prop + K/V_action 在多层堆叠后：\naction token 在每一层都融合了所有模态信息。\n八、总结一句话 这段代码本质是：\n一个支持 GQA + RoPE + AdaLN-Zero + LoRA 的 Transformer Decoder Block，用于 multi-modal 拼接场景。\n如果你愿意，我可以下一步画一个：\nFlow Matching 阶段 Q/K/V 拼接的真实张量流图\n那个会把你所有关于 18 层信息流的问题彻底讲透。\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"简单来说：没错，你的直觉非常准。\n在数学底层，神经网络的 Embedding 层确实可以被看作是一种特殊的矩阵乘法。但在工程实现上，为了效率，它被优化成了**查表（Look-up Table）**操作。\n1. 数学本质：One-hot 向量 × 权重矩阵 假设我们要把索引 $i$ 转换成一个 $d$ 维的向量。在数学上，这等同于：\n将索引 $i$ 转换成一个独热编码向量（One-hot vector） $\\mathbf{x}$。 让 $\\mathbf{x}$ 乘以 Embedding 权重矩阵 $\\mathbf{W}$。 $$\\mathbf{y} = \\mathbf{x} \\mathbf{W}$$由于 $\\mathbf{x}$ 中只有第 $i$ 个位置是 $1$，其余全是 $0$，矩阵乘法的结果会精准地取出矩阵 $\\mathbf{W}$ 的第 $i$ 行。\n2. 为什么不直接写成矩阵乘法？ 虽然数学上等价，但在代码里（比如 PyTorch 的 nn.Embedding）直接用矩阵乘法会造成极大的资源浪费：\n内存浪费：One-hot 矩阵极其稀疏（全是 0）。如果词汇表有 10 万个词，为了表示一个词，你需要一个长度为 10 万的向量，这太占空间了。 计算浪费：矩阵乘法涉及大量的 $0 \\times W_{ij}$。既然我们知道结果就是矩阵的某一行，为什么还要去算那些没意义的加法和乘法呢？ 结论： nn.Embedding 在底层直接通过索引去内存里“抓取”对应的行（即 Indexing），这比矩阵乘法快得多。\n3. 对比总结 特性 矩阵乘法 ($xW$) nn.Embedding (查表) 数学意义 完全相同 完全相同 输入要求 需要 One-hot 向量 需要 整数索引 (LongTensor) 运算效率 低 (计算大量冗余的 0) 极高 (直接内存偏移取值) 可学习性 权重随梯度更新 权重随梯度更新 💡 进阶：什么时候必须用矩阵乘法？ 虽然 nn.Embedding 很快，但在某些特定架构中（比如 Weight Tying 技巧），我们会让模型的最后一层全连接层的权重和 Embedding 层共享。\n在这种情况下，最后输出预测词概率时，我们必须用矩阵乘法：将 Hidden State 乘以 Embedding 矩阵的转置，从而计算当前向量与词汇表中每个词的相似度。\n你目前是在尝试实现某种特定的模型架构（比如 Transformer），还是只是在研究 Embedding 的底层原理？\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"Q1. Flow matching为什么比较难获得logp 这是一个非常直观的对比。简单来说，LLM（大语言模型）之所以能轻松获得对数似然，是因为它在数学本质上是一个分类器（Classifier），而基于 Flow Matching 的 VLA 在数学本质上是一个回归器（Regressor）/ 动力学系统。\n以下是深层原因的拆解：\n1. 离散空间 vs. 连续空间 LLM（离散空间）： LLM 的输出层是一个巨大的 Softmax 层。其输出是一个概率分布向量（维度通常为 32k 到 128k）。 当你输入一个序列，模型会直接给出下一个 token 出现的概率 $P(x)$。 对数似然 $\\log P(x)$ 只需要对 Softmax 的输出取个 $\\log$，这在计算上几乎是零成本的。 VLA (Flow Matching - 连续空间)： VLA 处理的是连续的动作坐标（如机械臂的 7 维关节角）。在连续空间中，概率变成了概率密度函数 (PDF)。 Flow Matching 并不直接输出 PDF 的值，它输出的是向量场 $v_t$（即指导噪声如何变成动作的“速度”）。 要把“速度”转换回“密度”，必须沿着轨迹积分，正如前文所说，这涉及极其复杂的雅可比矩阵计算。 2. 预测目标的不同 LLM 和 Flow-based VLA 在训练时的损失函数（Loss Function）就决定了它们的这种差异：\n模型类型 预测目标 损失函数 (Loss) 获取 Log-likelihood 的难度 LLM 概率 (Probability) 交叉熵 (Cross Entropy) 极易：直接读取输出层结果 VLA (FM) 速度 (Velocity/Vector Field) 向量场匹配 (L2 Regression) 极难：需解 ODE 和迹运算 3. 为什么不把 VLA 做成类似 LLM 的模式？ 既然 LLM 计算似然这么方便，为什么 VLA 不直接学 LLM 呢？其实早期的 VLA（如 RT-2）确实是这么做的，但存在以下几个严重的瓶颈：\n维度爆炸 (Curse of Dimensionality)： 如果把 7 维连续动作空间离散化（假设每个轴分 256 份），那么输出空间将达到 $256^7$，Softmax 根本算不动。 精度折损： 离散化会导致动作不平滑。如果你要把动作分得很细，Token 数量会激增，导致模型预测变得异常缓慢。 多峰分布问题： 在复杂的任务中，同一个状态可能有多种解（比如左绕或右绕障碍物）。Flow Matching 能够很好地建模这种复杂的连续多峰分布，而简单的离散化或单高斯分布很难兼顾精度和多样性。 4. 形象的类比 LLM 就像是在做“多选题”：卷子上写好了所有选项，模型只需要给每个选项打分。你问它某个答案的概率是多少，它看一眼自己的打分表就能告诉你。 Flow-Matching VLA 就像是在做“雕刻”：它拿到一块方正的大理石（高斯噪声），然后一刀一刀（向量场迭代）把它削成大卫像（最终动作）。你问它“大卫像出现的概率是多少”，它没法直接回答，它得回顾每一刀是怎么切的，并计算每一步把空间压缩了多少。 总结 LLM 的对数似然是直接输出的，因为它的动作空间（词表）是静态且有限的；而 Flow Matching 的对数似然是隐含在变换过程中的，因为它的动作空间是无限且连续的。\n这也正是为什么目前 VLA 结合强化学习是一个难题：我们既想要 Flow Matching 带来的高精度动作，又想要 LLM 那样方便计算的对数似然来跑 PPO。\n你目前研究的场景中，对动作的精度要求（如毫米级操作）是否很高？这决定了是否值得为了 PPO 而牺牲生成性能去尝试离散化方案。\nflow matching也是transformer 结构，为啥不用softmax 最后的输出头 这是一个非常关键的问题。简单直接地回答：Head 和 Loss 是“手”和“大脑指令”的关系。\n在深度学习中，Head（输出头） 决定了模型最终“长什么样”，而 Loss（损失函数） 决定了模型如何为了这个目标去“进化”。\n1. 它们之间是“一一对应”的 你不能随便配对 Head 和 Loss。Head 的输出形式必须满足 Loss 的数学要求：\n任务类型 Head 的样子 (最后输出层) 对应的 Loss (损失函数) 逻辑关系 分类 (如 VLM 吐字) Linear + Softmax CrossEntropy (交叉熵) Softmax 把输出变成概率，CrossEntropy 衡量这个概率离正确答案差多远。 回归 (如 FM 预测速度) Linear (无激活函数) MSE (均方误差) 模型直接输出一个数值，MSE 计算这个数和真实值的距离平方。 对比学习 (如 CLIP) 线性投影到固定维度 InfoNCE / Contrastive Loss Head 把图片和文字拉到同一个空间，Loss 让匹配的靠近，不匹配的远离。 2. Head 的结构直接影响 Loss 的计算 以 VLM (视觉语言模型) 为例： VLM 的 Head 是一个极其巨大的全连接层（词表大小，比如 100k 维）。\nHead 做的事： 把模型内部的隐藏特征投影到这 100k 个词上。 Loss 做的事： 它看了一眼 Head 输出的概率分布，发现模型给“猫”打了 0.1 分，给“狗”打了 0.8 分（但正确答案是猫），于是 Loss 就会产生一个很大的值，逼迫 Head 修改它的参数。 3. 回到你之前的疑问：Flow Matching (FM) vs VLM 这正是为什么 FM 算 $\\log p$ 难，而 VLM 容易的原因：\nVLM 的 Head 后面接了 Softmax：这使得它的输出天然就是一个概率分布。Loss 函数（CrossEntropy）直接取那个概率的负对数 $-\\log(P)$，就得到了 $\\log p$。计算量几乎为零。 FM 的 Head 输出的是速度向量 $v$：它没有 Softmax。它的 Loss 通常是 $||v_{\\theta} - v_{true}||^2$（一种回归 Loss）。这个 Loss 只管“像不像”，不管“概率是多少”。如果你非要算概率 $\\log p$，就得像我之前说的那样去做复杂的路径积分，这跟它的 Head 输出形式是不匹配的。 总结 Head 是为了适配 Loss 而存在的。\n如果你想用 分类 Loss，你的 Head 必须输出离散的得分（加 Softmax）。 如果你想用 生成/回归 Loss，你的 Head 必须输出连续的数值。 所以，当你问“Flow Matching 为什么不能直接加 Softmax”时，本质上是你试图用分类任务的 Head 去解回归任务的 Loss，这会导致模型生成的不是“流的方向”，而是一堆无意义的概率，从而无法训练。\n你现在是在设计一个新的模型架构，在纠结该选什么样的 Head 吗？\npi0.6* 中的强化学习 RECEP 这张截图描述的是 RECAP（一种针对 VLA 改进的强化学习方案）的核心逻辑。它非常巧妙地通过引入一个指示变量 (Indicator Variable) $I$，把复杂的强化学习问题强行“转译”成了 VLA 最擅长的生成问题。\n为了让你透彻理解，我们把这段话拆解为三个层次：\n1. 核心操作：把“优势”变成一个标签 (Label) 传统的 PPO 是根据优势函数 $A$ 的数值大小来调整概率梯度。而 RECAP 剑走偏锋：\n它定义了一个二元变量 $I$。 规则：如果某个动作的优势 $A$ 超过了一个阈值 $\\epsilon_l$，那么 $I=1$（代表这是个“好球”）；否则 $I=0$。 数学表达：这就是图中的 Delta 分布 $\\delta(\\cdot)$，意思是非常确定，要么是 0，要么是 1。 直观理解：它给每一条训练数据打了个标签：“优等生”或“普通生”。\n2. 损失函数 (Loss Function) 的玄机 图中的损失函数由两部分组成： $$\\min_\\theta \\mathbb{E} [ \\underbrace{-\\log \\pi_\\theta(a|o, \\ell)}_{\\text{第一项}} - \\alpha \\underbrace{\\log \\pi_\\theta(a|I, o, \\ell)}_{\\text{第二项}} ]$$ 第一项（基础模仿）：不看标签，让模型学习所有的行为。这保证了模型具备基本的控制能力和对环境的理解（类似基础的 Behavior Cloning）。 第二项（条件生成）：这是 RECAP 的精髓。它把 $I$ 作为一个条件 (Condition) 喂给模型。 当 $I=1$ 时，模型在学习：“如果我想表现得优秀，我该怎么做？” 通过 $\\alpha$ 系数加大这一项的权重，强迫模型去重点模仿那些高优势值的动作。 3. 为什么它能解决 Flow Matching 的痛点？ 这回到了你最初的问题：Flow Matching 没法算 $\\log p$。\nRECAP 的聪明之处在于：虽然公式里写着 $\\log \\pi$，但实际训练时它根本不去算概率。\n它利用了 “条件生成 (Conditional Generation)” 的逻辑。 在 Flow Matching 中，这一项对应的是预测条件向量场 $v_t(x_t | I, o, \\ell)$。 训练时：它依然是做 L2 回归（预测速度向量）。 推理时：我们只需要在输入端给模型下达一个指令：“请以 $I=1$（优等生模式）来生成动作”。 4. 总结：这是一种“欺骗”技巧 RECAP 本质上是把 “强化学习（寻找最优策略）” 降级为了 “条件控制的模仿学习”。\nPPO 这种老实人：努力去算 $\\log p$，试图把概率分布的峰值推向高奖励区域（在 Flow Matching 上很难算）。 RECAP 这种聪明人：我直接在输入里加个“开关”。训练时告诉你哪些是好动作，推理时我把“好动作开关”拨到 1，模型自然就倾向于生成高优势值的动作了。 这就是图中最后一段话的意思：虽然理论公式写成对数似然的形式，但实际落地到 VLA 时，它是通过 Flow Matching 的 L2 Loss 来实现的。\n右对齐，mask,kv cache 这是一个非常经典的 Transformer 推理（Inference）优化技巧，特别是在使用 KV Cache 的自回归生成场景下。\n在 sample_actions 中，将序列进行“右对齐”（Left-to-Right Align，实际上是将内容推向右侧，Padding 留左侧）主要有以下三个原因：\n1. 统一 KV Cache 的起始索引（Index Consistency） 在自回归生成中，Transformer 需要知道每一个 Token 的位置索引（Position ID）。\n左对齐的问题：如果 batch 中不同样本的指令长度不同（比如样本 A 指令 10 个词，样本 B 指令 20 个词），它们的“动作起始位置”在数组中的索引就不一致。 右对齐的好处：通过右对齐，我们可以保证所有样本的“最后一个有效 Token”都对齐在数组的同一个末端位置（例如 prefill_size - 1）。这样在进入 while_loop 生成第一个动作 Token 时，所有样本的 step=0 对应的位置 ID 都是统一的计算逻辑，方便并行化（Batch Processing）。 2. 简化自回归掩码（Attention Mask）的构造 Transformer 的 mask 决定了哪些 Token 可以互相看到。\n在右对齐模式下，所有的 Padding Token 都堆在左侧。 我们在计算 mask 时，只需要简单地告诉模型：“忽略索引从 $0$ 到 prefix_start 的部分”。 如果不做对齐，Padding 可能会散落在序列中间（尤其是多轮对话或复杂输入时），这会导致生成每一位动作 Token 时都要动态维护极其复杂的嵌套掩码，极大地增加了计算开销和代码复杂度。 3. 符合 KV Cache 的增长逻辑 代码中有一行关键的操作：\n1 prefix_attn_mask = jnp.pad(prefix_attn_mask, ((0, 0), (0, 0), (0, max_decoding_steps))) 这是在为即将生成的 max_decoding_steps（动作 Token 序列）预留空间。\n逻辑连续性：右对齐后，输入序列的“尾部”紧接着就是我们要生成的第一个动作。 位置编码（Positional Embedding）：在 JAX/NNX 的实现中，positions = jnp.cumsum(prefix_mask, axis=-1) - 1。右对齐确保了有效 Token 的位置编码是从 $0$ 开始连续增长到末尾的，而不会被开头的 Padding 干扰。如果不右对齐，位置编码可能会包含 Padding 的位次，导致模型对距离的感知出现偏差。 4. 形象化的类比 想象你在排队：\n左对齐：大家从队首往后排，但每个人带的行李（指令长度）不一样长。领队（模型）要给每个人发动作指令时，得先量一下每个人行李有多长，才能找到每个人的手在哪。 右对齐：所有人不论行李长短，全部“后备箱对齐”停在发车线上。领队只需要站在发车线（数组末尾）统一发令，所有人下一步产生的动作都会整齐划一地出现在发车线之后。 核心代码对应： 1 2 prefill_len = jnp.sum(prefix_mask, axis=-1) prefix_start = prefill_size - prefill_len # 计算每个样本真正内容的起点 通过 prefix_start，模型在后续的 while_loop 中可以非常精确地通过一个偏移量（step + 1）找到每一个样本应该写入新 Token 的位置，而不需要感知每个样本原始 Padding 的长度。\n总结： 这是一种工程上的规范化。它牺牲了输入时的整齐度（把 Padding 扔到左边），换取了推理生成阶段（Autoregressive Decoding）极高的并行效率和逻辑一致性。\n你现在的代码调试中，是因为对齐后的 prefix_start 计算不对导致模型生成的动作全是乱码（或 Padding）吗？\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"你这个问题非常关键 👍 它其实是在问：\n在 GRPO 里，为什么可以直接用 rewards 当作动作价值 Q？ 不用单独训练一个 value 网络？\n我们一步一步讲清楚。\n动作价值 和状态价值 action 的 价值 当前state 多个action, 做期望\n当\n1 2 3 4 mean = torch.mean(rewards) std = torch.std(rewards, unbiased=False) advantages = (rewards - mean) / (std + self.eps) return advantages.detach() 这其实是在做：\n[ A_i = \\frac{r_i - \\mu}{\\sigma} ]\n也就是：\n用 reward 做中心化 + 标准化 当作 advantage\n标准 PPO 里：\n[ A_t = Q(s,a) - V(s) ]\n其中：\nQ 需要估计 V 需要一个 value 网络 🎯 3️⃣ 为什么 GRPO 可以直接用 reward？ 因为在 RLHF / GRPO 场景里：\n⭐ 奖励结构非常特殊 通常：\n只在序列结束时给一个最终 reward 没有中间奖励 也就是说：\n[ r_t = 0 ] [ r_{final} = R ]\n那么：\n[ Q(s_t, a_t) R ]\n对整个序列里的每个 token 来说，\n动作价值就是最终 reward。\n🎯 4️⃣ 那 V(s) 怎么办？ GRPO 不训练 value 网络。\n而是：\n在一个 group（同一个 prompt 生成的多个回答）里做 baseline。\n比如：\n一个 prompt 生成 4 个回答：\nresponse reward A 0.9 B 0.7 C 0.3 D 0.1 它做：\n[ A_i = r_i - mean(r) ]\n这等价于：\n[ A_i = Q_i - \\mathbb{E}[Q] ]\n而我们之前证明过：\n[ V(s) = \\mathbb{E}[Q(s,a)] ]\n所以：\n用组内平均 reward 近似 V(s)\n🎯 5️⃣ 本质原因 GRPO 是在做：\n[ A(s,a) Q(s,a) \\underbrace{\\mathbb{E}{a\\sim\\pi}[Q(s,a)]}{V(s)} ]\n但不是用 value 网络估计 V，\n而是用：\n同一 prompt 下的多个采样均值\n来近似。\n🎯 6️⃣ 为什么这样可行？ 因为在 RLHF 里：\nstate = prompt action = 整个 response reward = RM 打分 所以：\n[ Q(prompt, response) = reward(response) ]\n而：\n[ V(prompt) \\mathbb{E}_{response\\sim\\pi}[reward] ]\n正好等于：\n多个 response reward 的平均值\n🎯 7️⃣ 为什么还要除以 std？ 这一步：\n[ \\frac{r - mean}{std} ]\n是在做：\n尺度归一化 防止 reward 波动过大 稳定训练 类似于 advantage normalization。\n🎯 8️⃣ 为什么 detach？ 1 return advantages.detach() 因为：\nreward 来自 reward model 不希望梯度传回 reward 只更新 policy 这和 PPO 里 advantage 不参与梯度计算一致。\n🔥 核心总结 在 GRPO 中：\nreward 就是 Q 组内均值近似 V advantage = reward - group_mean 不需要单独训练 value 网络 🎯 一句话本质 GRPO 把：\n“状态价值 V 的学习”\n替换成：\n“组内采样的统计均值”\n用 Monte Carlo 直接估计。\n如果你愿意，我可以给你画一个：\nPPO（有 value head） GRPO（无 value head） 结构对比图。\n你会非常清楚两者的数学差别。\n重要性采样 这是一个非常棒的思考点。你抓住了**深度学习（监督学习）与强化学习（RL）**之间最容易混淆的区别。\n深度学习（如图片分类、文本续写）确实是多 Epoch 的，但强化学习里的“数据”是有“保质期”的。\n1. 核心矛盾：数据会“过期” 监督学习（SL）： 图片还是那张图片，标签还是那个标签。无论模型更新多少次，图片里的猫永远是猫。所以你可以跑 100 个 Epoch，这叫“充分学习”。 强化学习（RL）： 数据是模型“跑”出来的。 模型参数 $\\theta$ 决定了它怎么说话（概率分布）。 问题在于： 当模型更新了一次参数，变成了 $\\theta_{new}$，它对同一句话的理解和概率分布已经变了。 如果你用旧模型 $\\theta_{old}$ 采回来的数据去喂给新模型 $\\theta_{new}$，新模型会觉得：“这不是我写的，我不认可这个概率”。 结论： 在 RL 里，模型一旦更新，刚才采的数据在数学上就“过期”了。\n2. 为什么非要用重要性采样（IS）？ 既然数据会过期，那为什么不每更新一次就重新采样（Generate）呢？ 因为采样太贵了！ 生成几千个 Token 需要好几秒，而更新一次参数只需要几毫秒。\n为了能像深度学习那样，让一份采样数据多跑几次循环（压榨价值），我们必须解决“数据过期”的问题。重要性采样就是那个“延保合同”：\n$$\\text{新的梯度} = \\underbrace{\\frac{\\pi_{new}(a|s)}{\\pi_{old}(a|s)}}_{\\text{重要性采样系数}} \\times \\text{优势函数} \\times \\text{基础梯度}$$ 如果没有这个系数： 你强行跑第 2 个 Epoch，新模型会用自己现在的理解去套旧的数据，导致数学期望完全错误，训练会迅速崩溃。 有了这个系数： 它告诉模型：“虽然这是旧数据，但只要我乘以这个比例，我在数学上就依然等同于‘正在用新模型实时采样’”。 3. 你的代码里为什么还要写 IS 系数？ 回到你提供的代码： per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages\n虽然在当前这个 compute_loss 函数的单次运行中，$Ratio$ 结果是 $1$，但引入它的原因有二：\n① 兼容“梯度累积” (Gradient Accumulation) 在大模型训练中，一个 Batch 往往被拆成很多个 Micro-batch。\n计算第一个 Micro-batch，得到梯度，但不更新参数。 计算第二个 Micro-batch…… 在这个过程中，per_token_logps.detach() 锁定了采样时的原始状态，确保整个大 Batch 内部的逻辑是一致的。 ② 它是“策略梯度”的微分技巧 在 PyTorch 里，我们不能直接对 Advantages 求导。我们需要一个包含参数 $\\theta$ 的目标函数。\n如果你写 loss = log_prob * A，求导后得到 $\\nabla \\log \\pi \\cdot A$。 如果你写 loss = exp(log_prob - log_prob_old) * A，求导后同样得到 $\\nabla \\log \\pi \\cdot A$。 这是一种通用的写法。 即使你只跑 1 个 Epoch，这种写法在数学上也是 100% 正确的策略梯度实现。它就像是给赛车装了氮气系统，虽然你现在还没开启“多 Epoch 压榨”模式，但引擎的结构必须先支持这个系统。\n总结：到底为了啥？ 为了数学正确： 它是策略梯度（Policy Gradient）的标准算式。 为了未来的扩展： 如果你哪天想把 num_train_epochs 从 1 改成 3，你不需要改代码，因为 IS 系数会自动帮你修正这 3 次压榨过程中的偏差。 为了稳定性： detach() 提供了一个稳定的基准，防止在计算 Loss 的过程中，计算图出现回环或不必要的逻辑偏移。 一句话点破： 在深度学习里，你可以直接学；在强化学习里，因为你要“一边改规则一边比赛”，所以必须用重要性采样来对齐新旧规则。\n是的，完全正确！ 在强化学习的语境下，Rollout 永远是基于“采样那一刻的策略”进行的，而这个策略在随后的参数更新（Optimization）过程中，就变成了 Old Policy。\n这里有三个关键的时间节点，能帮你彻底理清“新、旧、当前”的关系：\n1. 采样阶段 (Rollout Phase) 动作执行者：此时的模型我们称为 $\\pi_{old}$。 产出：它生成了句子（Completions），并且记录了每一个词的概率 old_logps。 地位：这些数据是“历史存证”。 2. 梯度计算时刻 (Forward Pass) 在你的代码里，这一行体现了对比： per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages\nper_token_logps：这是当前模型（正在被训练、准备变强的模型）对刚才那些旧句子的看法。 per_token_logps.detach()：这就是 Rollout 时的那个模型（Old Policy）留下的看法。 3. 为什么“压榨”必须用 Old Policy 的数据？ 想象一个练射箭的场景：\nRollout (旧我)：你闭着眼睛射了 10 箭。这时候你的“箭法”是旧的。射完后，你睁开眼看到这 10 箭的分布（这就成了你的 Batch 数据）。 Update (新我)：你根据这 10 箭的反馈，调整了站姿（参数更新）。 压榨 (Re-learning)： 如果你想根据刚才那 10 箭再反思一次，你就必须记住你射箭那一刻的姿势（Old Policy）。 因为你现在的姿势（New Policy）已经变了。如果你用现在的姿势去强行解释刚才那 10 箭，逻辑就乱了。 重要性采样就是用来告诉模型：“基于我刚才那个烂姿势（Old），虽然这支箭射偏了，但以我现在的姿势（New）来看，这个偏差其实是可以修正的。” 4. 在你的代码里的具体表现 虽然你的代码里采样和计算 Loss 写在同一个函数里，但逻辑顺序是严格的：\nunwrapped_model.generate(...) $\\rightarrow$ 这是 Rollout，用的是采样瞬间的模型。 _get_per_token_logps(...) $\\rightarrow$ 这是在为更新做准备，获取当前模型的概率。 通过 detach()，你人为地在内存里创造了一个“旧策略”的快照。 总结 Rollout 确实是用 Old Policy。 引入重要性采样的本质原因就是：承认“采样”和“学习”之间存在时间差。 * 在监督学习里，猫的照片不会变，没有时间差，不需要 IS。\n在强化学习里，模型每秒钟都在变，IS 就是用来弥合“采样时的我”和“学习时的我”之间差距的那个系数。 如果没有这个系数，模型在多次迭代（Epochs）中就会因为“无法正确评价过去的自己”而陷入逻辑混乱，导致模型崩溃。\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"很好 👍 你现在问的是 decode 阶段 attention 里 Q、K、V 的精确维度关系。\n我们严格按你的代码来走一遍。\n一、设定符号 假设：\nB = batch size H = num_heads d_k = 每个 head 维度 S_past = 已生成 token 数 S_new = 当前输入 token 数（decode 时 = 1） 二、decode 阶段各张量维度 1️⃣ 计算完 Q / K / V（还没 concat） 1 2 3 q = self.W_q(x) k = self.W_k(x) v = self.W_v(x) 因为 decode 时：\n1 x.shape = (B, 1, d_model) reshape 后：\n1 2 3 q.shape = (B, H, 1, d_k) k.shape = (B, H, 1, d_k) v.shape = (B, H, 1, d_k) 2️⃣ 拼接历史 cache 1 2 k = torch.cat([cache[0], k], dim=2) v = torch.cat([cache[1], v], dim=2) 假设：\n1 cache[0].shape = (B, H, S_past, d_k) 拼接后：\n1 2 k.shape = (B, H, S_past + 1, d_k) v.shape = (B, H, S_past + 1, d_k) 而：\n1 q.shape = (B, H, 1, d_k) 没有变化。\n三、Attention 计算维度 核心计算：\n1 scores = torch.matmul(q, k.transpose(-2, -1)) 我们拆维度。\n1️⃣ k.transpose(-2, -1) 原来：\n1 k = (B, H, S_total, d_k) transpose 后：\n1 k^T = (B, H, d_k, S_total) 其中：\n1 S_total = S_past + 1 2️⃣ q @ k^T 维度：\n1 2 3 (B, H, 1, d_k) @ (B, H, d_k, S_total) 结果：\n1 (B, H, 1, S_total) 这表示：\n当前 token 对所有 token 的注意力分数\n四、softmax 后 1 weights = torch.softmax(scores, dim=-1) 维度仍然：\n1 (B, H, 1, S_total) 五、再乘 V 1 attn = torch.matmul(weights, v) 维度：\n1 2 3 (B, H, 1, S_total) @ (B, H, S_total, d_k) 结果：\n1 (B, H, 1, d_k) 六、拼回去 1 attn.transpose(1,2).contiguous().view(B, 1, d_model) 最终：\n1 (B, 1, d_model) 七、完整 decode 阶段维度流 1 2 3 4 5 6 7 8 9 q: (B, H, 1, d_k) k: (B, H, S_total, d_k) v: (B, H, S_total, d_k) scores = (B, H, 1, S_total) weights = (B, H, 1, S_total) attn = (B, H, 1, d_k) output = (B, 1, d_model) 八、prefill 阶段对比 如果是 prefill：\n1 S_new = S_total = S 那：\n1 2 q: (B, H, S, d_k) k: (B, H, S, d_k) attention 结果：\n1 (B, H, S, S) 也就是完整注意力矩阵。\n九、核心区别 阶段 Q 长度 K/V 长度 attention 矩阵 Prefill S S S × S Decode 1 S_total 1 × S_total 这就是为什么：\ndecode 每步是 O(S) prefill 是 O(S²)\n十、一句话总结 在 decode 阶段：\nQ: (B, H, 1, d_k) K: (B, H, S_total, d_k) V: (B, H, S_total, d_k) Attention: (1 × S_total) 即：\n当前 token 用 1 个 query 去匹配全部历史 key。\n如果你愿意，我可以下一步讲：\n为什么这样复杂度从 O(N²) 变 O(N) 或者画一个 token-by-token 时间展开图 你现在已经完全理解 KV cache 的维度本质了 👍\n很好 👍 你现在问的是：\nif S_new \u0026gt; 1: 这段 causal mask 逻辑到底在干嘛？\n这段代码其实只在 prefill 阶段 起作用。我们拆解讲清楚。\n一、什么时候会进入这个分支？ 1 if S_new \u0026gt; 1: Prefill 阶段：S_new = 整个 prompt 长度（比如 100） → ✅ 进入 Decode 阶段：S_new = 1 → ❌ 不进入 所以：\n这段代码是给 prefill 用的 causal mask\n二、它解决什么问题？ 在 decoder-only Transformer 里：\n每个 token 只能看到自己和之前的 token 不能看到未来 token\n否则就是“偷看答案”。\n三、prefill 阶段的特殊性 假设：\n你输入整个 prompt 一次性算 attention 没有逐 token 计算 那么 attention 矩阵是：\n1 (S × S_total) 如果没有 mask：\n第 3 个 token 可以看到第 10 个 token 违反自回归规则 所以必须加 causal mask。\n四、关键变量解释 1 S_past = S_total - S_new 在 prefill 时：\n1 cache = None 所以：\n1 2 S_total = S_new S_past = 0 五、mask 的构造逻辑 1 2 3 4 mask = torch.triu( torch.ones(S_new, S_total), diagonal=S_past + 1, ) 在 prefill 情况下：\n1 diagonal = 1 所以：\n1 torch.triu(ones, diagonal=1) 这会生成：\n1 2 3 4 0 1 1 1 0 0 1 1 0 0 0 1 0 0 0 0 上三角为 1。\n表示：\n右上角（未来 token）全部屏蔽。\n六、masked_fill 的作用 1 scores = scores.masked_fill(mask, -inf) 意味着：\n被 mask 的位置 softmax 后变成 0\n从而：\n每个 token 只能关注自己和左边 七、为什么 decode 阶段不需要 mask？ decode 阶段：\n1 S_new = 1 attention 形状：\n1 (1 × S_total) 当前 token 是序列最后一个。\n它本来就只能看到：\n1 历史 token + 自己 没有“未来 token”存在。\n所以：\n天然满足 causal，不需要 mask。\n八、完整逻辑总结 阶段 是否进入 mask 为什么 Prefill ✅ 一次算整个序列，必须阻止看未来 Decode ❌ 每次只算最后一个 token，本来就不会看未来 九、一句话总结 1 if S_new \u0026gt; 1: 的作用是：\n在 prefill 阶段构造 causal mask，防止 token 看到未来。\n十、再深入一点（很关键） 这行代码：\n1 diagonal=S_past + 1 是为了支持：\ncache + prefill 混合场景\n也就是说：\n有历史 cache 又输入多个新 token 它保证：\n新 token 之间也遵守 causal\n如果你愿意，我可以下一步讲：\n为什么 diagonal 要加 S_past 或画出 “prefill + cache” 混合场景的 mask 结构图 你现在已经在理解 Transformer 推理中最容易出 bug 的地方了 👍\n很好，这个问题问得非常关键 👍 你已经进入 KV cache + causal mask 的核心逻辑 了。\n先给你一句话答案 1 S_past = S_total - S_new 是为了知道：\n当前 batch 里有多少“历史 token”（来自 cache）\n从而正确地构造 带 cache 的 causal mask。\n一、我们拆开看这几个量 在 forward 里：\n1 B, S_new, _ = x.shape S_new = 这一次新输入的 token 数量\nprefill：等于整个 prompt 长度 decode：等于 1 然后：\n1 S_total = k.shape[2] 注意：\n1 k = torch.cat([cache_k, new_k], dim=2) 所以：\n1 S_total = S_past + S_new 于是：\n1 S_past = S_total - S_new 就是在算：\ncache 里有多少历史 token\n二、为什么必须知道 S_past？ 因为 causal mask 的结构取决于：\n当前 query 在完整序列中的绝对位置\n三、举一个具体例子 假设：\ncache 里已经有 5 个 token 现在一次性输入 3 个新 token（prefill + cache 场景） 那么：\n1 2 3 S_past = 5 S_new = 3 S_total = 8 我们希望 attention 允许的情况是：\n1 2 3 新 token 0 → 可以看 0~5 新 token 1 → 可以看 0~6 新 token 2 → 可以看 0~7 而不能看未来。\n四、如果不加 S_past 会发生什么？ 如果你写：\n1 torch.triu(..., diagonal=1) 那 mask 会认为：\n新 token 是从位置 0 开始的\n这会导致：\n错误屏蔽历史 token 或允许新 token 看未来 token 模型直接乱掉。\n五、S_past 的数学意义 1 diagonal = S_past + 1 这句的意思是：\n在完整序列坐标系下构造上三角 mask\n因为：\n1 query 的实际位置 = S_past + i 而不是从 0 开始。\n六、画一个结构图（最清晰） 假设：\n1 2 历史 token: H0 H1 H2 H3 H4 新 token: N0 N1 N2 完整序列：\n1 [ H0 H1 H2 H3 H4 N0 N1 N2 ] 我们在算：\n1 N0 N1 N2 作为 query 它们的位置分别是：\n1 2 3 N0 → 5 N1 → 6 N2 → 7 mask 必须以这个“真实位置”来构造。\n而：\n1 S_past = 5 正是用来告诉 mask：\n新 token 的起始 index 是 5\n七、如果是纯 prefill 呢？ prefill：\n1 cache = None 那么：\n1 2 S_total = S_new S_past = 0 mask 就退化成：\n1 diagonal = 1 也就是普通 causal mask。\n八、如果是纯 decode 呢？ decode：\n1 S_new = 1 直接跳过 mask 分支。\n九、总结一句话 1 S_past = S_total - S_new 的作用是：\n计算 cache 里已有多少历史 token，从而在完整序列坐标系下构造正确的 causal mask。\n十、你可以记住这个本质 Transformer 在做的永远是：\n1 key_position \u0026lt;= query_position 所以 attention mask 是：\nkey 0 1 2 3 query 0 ✓ ✗ ✗ ✗ 1 ✓ ✓ ✗ ✗ 2 ✓ ✓ ✓ ✗ 3 ✓ ✓ ✓ ✓\nscores[i][j]\ni 是 query 位置 j 是 key 位置\n而：\n1 S_past 决定了 query 的绝对位置 八、核心推导（最重要）\n我们希望屏蔽：\nkey_position \u0026gt; query_position\n代入真实位置：\nj \u0026gt; S_past + i\n等价于：\nj ≥ S_past + i + 1\n整理：\nj - i ≥ S_past + 1\n这正是：\ndiagonal = S_past + 1\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""},{"content":"SFT CrossEntropyLoss 具体计算过程 在 SFT (Supervised Fine-Tuning) 阶段，模型学习的是“给定上文，预测下一个词”。这个过程本质上是在做一连串的分类任务。\n计算 CrossEntropyLoss（交叉熵损失）的过程可以拆解为以下几个步骤：\n1. 从 Logits 到 概率 (Softmax) 模型的最后一个层（Head）输出的是 Logits。假设词表大小为 $V$（比如 100k），那么对于每一个生成的位点，模型都会输出 $V$ 个实数。\n首先，通过 Softmax 函数将这些实数转化为概率分布 $P$： $$P_i = \\frac{\\exp(\\text{logit}_i)}{\\sum_{j=1}^{V} \\exp(\\text{logit}_j)}$$此时，每一个词都有一个 $0$ 到 $1$ 之间的概率。\n2. 独热编码 (One-Hot) 与 筛选 在 SFT 中，我们有“标准答案”（Ground Truth）。 假设在某个位置，正确的词是“苹果”，它在词表中的索引是 42。\n标准答案是一个 One-Hot 向量：只有索引 42 的地方是 $1$，其余全是 $0$。 CrossEntropy 的核心：它只关心模型给“正确答案”打了多少分。它会把对应索引 42 的那个概率 $P_{42}$ 抠出来。 3. 计算负对数似然 (NLL) 这就是你提到的 $-\\log(P)$。 $$\\text{Loss} = -\\log(P_{\\text{target}})$$ 如果模型预测正确词的概率是 $0.9$，那么 $\\text{Loss} = -\\log(0.9) \\approx 0.105$（损失很小）。 如果模型预测正确词的概率是 $0.1$，那么 $\\text{Loss} = -\\log(0.1) \\approx 2.302$（损失很大）。 4. SFT 中的特殊处理：Masking（掩码） 这是 SFT 和预训练最重要的一点区别。在 SFT 的数据中，通常包含 Prompt（问题） 和 Answer（回答）。\n我们不希望模型去学习如何预测 Prompt。 做法：在计算 Loss 时，我们会把 Prompt 部分对应的标签设为 -100（这是 PyTorch 中的约定值）。 计算时：CrossEntropyLoss 会自动忽略标签为 -100 的位置，只计算 Answer 部分的损失。 5. 整体流程总结 假设输入序列长度为 $L$：\nForward：模型跑一遍，得到 $L \\times V$ 的 Logits。 Shift（对齐）：因为是预测下一个词，所以要把 Logits 往右移一位，对应 Label。 Logits[0] 对应预测 Label[1]。 Mask：把 Prompt 部分的 Label 抹掉。 Reduction：把 Answer 部分所有有效位置的 $-\\log(P)$ 加起来求平均。 1 2 3 4 5 6 7 8 # 伪代码示例 logits = model(input_ids) # (Batch, SeqLen, VocabSize) # 预测第 i 个词是为了匹配第 i+1 个 label shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss_fct = nn.CrossEntropyLoss(ignore_index=-100) loss = loss_fct(shift_logits.view(-1, V), shift_labels.view(-1)) 为什么 SFT 这么高效？ 因为 CrossEntropy 这种“指哪打哪”的特性：模型不需要考虑整个句子的空间分布，只需要在每一个时刻把“正确答案”那个维度的能量提上去，把其他的压下来。相比 Flow Matching 需要在连续空间里找路径，SFT 的这个 $-\\log(P)$ 简直就是“降维打击”般的简单。\n你是在写 SFT 的 Trainer 逻辑，还是在排查 Loss 不下降的问题？\n为啥不是模型generate和 sft 中的answer作为label 去做CrossEntropyLoss？ 而是预测下一个词，所以要把 Logits 往右移一位，作为 Label ？ 这是一个非常经典的学习误区，理解了这一点，你就理解了 自回归模型（Autoregressive Models） 的精髓。\n之所以不能直接用“整个 Answer”作为 Label 去做 CrossEntropy，是因为 CrossEntropy 是一个“点对点”的分类损失，它需要明确地知道：在这一时刻，输入是什么，对应的正确输出是什么。\n以下是三个核心原因：\n1. 任务本质：模型是“单字复读机” LLM（如 GPT）的本质不是“一键生成整段话”，而是根据已知序列预测下一个字。\n如果你的 Answer 是 [A, B, C, D]。 模型在生成 B 的时候，它的输入必须是 [Prompt, A]。 模型在生成 C 的时候，它的输入必须是 [Prompt, A, B]。 如果我们不“往右移一位”，而是直接把 Answer 当 Label，模型就会变成：输入 A，预测 A。这叫恒等映射，模型只要学会“抄写”当前输入的最后一个字就能让 Loss 降为 0，但这根本没学会任何逻辑。\n2. “偏移”是为了建立因果关系 在代码层面，shift 操作是为了让当前的输入和未来的目标对齐。\n假设我们要模型学一句话：“北京是首都”。\n第 1 步： 输入 [北京] -\u0026gt; 期望输出 [是] 第 2 步： 输入 [北京, 是] -\u0026gt; 期望输出 [首都] 在矩阵计算时，模型是一次性算出所有位置的输出（Logits）的：\nLogits[0] 对应的是输入第一个词后的预测结果。 Labels[1] 才是我们要预测的那个词。 所以，必须把 Logits 的第 $i$ 位和 Labels 的第 $i+1$ 位对齐，这个 Loss 计算才有意义。 3. 为什么不能用“生成的句子”去算 Loss？ 你提到的“模型 generate 出来的东西和 Answer 做 CrossEntropy”，这其实是 RLHF（强化学习） 或者某些 Sequence-level Loss 的思路，但在 SFT 阶段行不通：\n梯度无法回传：generate 过程通常包含 argmax 或随机采样，这些操作在数学上是不可导的。这意味着如果你直接用生成的句子比对，误差（Loss）没法通过反向传播告诉模型参数该怎么改。 Teacher Forcing（教师强迫）：SFT 采用的是 Teacher Forcing 机制。无论模型上一步预测得多么离谱，我们在计算下一步 Loss 时，依然强制给模型输入正确的上文。这样可以防止模型在训练初期因为一步错步步错，导致根本无法收敛。 4. 总结 做法 结果 不偏移 (Input=A, Label=A) 模型学会了简单的复制。 生成后比对 (RL 方式) 计算极其缓慢，且梯度断裂，难以训练。 偏移一位 (Input=A, Label=B) 模型学会了预测未来，建立了因果逻辑。 一句话总结： shift 是为了强迫模型利用“过去”的信息去猜“未来”，而不是让它盯着“现在”看。\n你是不是在看 transformers 库里的 LlamaForCausalLM 源码？那里的 shift_logits 逻辑确实困扰过很多人。\n如果在 SFT 中，模型预测错了某一个词，会对后续位置的 Loss 产生什么影响吗？\n","date":"0001-01-01T00:00:00Z","permalink":"https://jiangyigithub.github.io/ai.github.io/p/","title":""}]