跳转至

xLSTM Scaling Laws: Competitive Performance with Linear Time-Complexity

会议: ICLR2026
arXiv: 2510.02228
代码: NX-AI/xlstm_scaling_laws
领域: llm_efficiency
关键词: scaling laws, xLSTM, 线性复杂度, Transformer对比, 推理效率

一句话总结

系统对比 xLSTM 与 Transformer 的 scaling law,证明 xLSTM 在训练损失-算力 Pareto 前沿、过训练 regime 和推理速度上全面优于同规模 Transformer,且优势随上下文长度增大而增长。

背景与动机

  • Scaling law 是 LLM 设计的核心指导工具(Kaplan 2020, Chinchilla 2022),但现有研究几乎全部聚焦在 Transformer 架构
  • xLSTM 等线性复杂度架构已在十亿参数级别展现竞争力(xLSTM 7B),但缺乏与 Transformer 的系统性 scaling 对比
  • 传统 FLOP 近似 \(C(N,D)=6ND\) 忽略了注意力机制的计算量,无法公平比较线性/二次复杂度模型
  • 对推理效率(TTFT、step time)与上下文长度的交互影响也缺乏系统分析

核心问题

  1. 训练效率: 给定算力预算,xLSTM 与 Transformer 谁的 loss 更低?
  2. 过训练 regime: xLSTM 在高 token/parameter 比下是否仍保持稳定的 power-law 指数?
  3. 上下文长度: 线性 vs 二次复杂度如何影响 compute-optimal 模型大小?
  4. 推理: 在 TTFT 和 step time 上两种架构如何随上下文长度缩放?

方法详解

实验框架

  • 模型: Llama-2 风格 dense multi-head Transformer vs xLSTM 7B 架构(纯 mLSTM 层 + MLP)
  • 规模范围: 80M–7B 参数,2B–2T tokens,总计 672 次训练(292 Transformer + 380 xLSTM)
  • 总算力: \(3.2 \times 10^{23}\) FLOPs
  • 训练数据: DCLM-Baseline(高质量过滤网页文档),GPT-NeoX tokenizer,默认序列长度 8192

精确 FLOP 计算

  • 抛弃简化的 \(6ND\) 近似,采用精确 FLOP 公式,区分注意力计算(二次项)和前馈计算
  • 对 xLSTM 的递归更新、mLSTM 矩阵运算也做精确 FLOP 统计

Scaling Law 拟合

  • 参数化拟合: \(\hat{L}(N,D) = E + (A N^{-\alpha} + B D^{-\beta})^{\gamma}\),引入 \(\gamma\) 自由参数提升拟合质量
  • IsoFLOP 方法: 固定算力预算,变化 \(N\)\(D\),拟合二阶多项式找最优 \(N^*(H)\)\(D^*(H)\)
  • Power-law 外推: \(\hat{N}^*(H) = A' \cdot H^a\), \(\hat{D}^*(H) = B' \cdot H^b\)

推理建模

  • 将推理时间建模为 \(\tau = \text{FLOPs}_{\text{algo}} / \alpha_{\text{eff}} + \epsilon\)\(\tau = \text{Bytes}_{\text{mem}} / \beta_{\text{eff}} + \epsilon\)
  • 基于 roofline model 判断是计算密集还是内存密集
  • 分 prefill 和 generation 两阶段分别分析

实验关键数据

训练 Scaling

发现 细节
Pareto 支配 xLSTM 在近 5 个数量级算力范围内严格 Pareto-dominate Transformer
过训练指数 xLSTM 的 power-law 指数 \(\eta\)\(M=22\)\(M=2200\) 范围内保持恒定,与 Transformer 一致
Compute-optimal 大小 相同算力下,xLSTM 最优模型更大(线性运算更便宜→更多参数分配给深度/宽度)
上下文长度影响 Transformer 在 2048→16384 时最优模型大小显著下降;xLSTM 保持稳定

推理性能

指标 16k prefill 结果
TTFT xLSTM 比同尺寸 Transformer 低 30–50%
Step time xLSTM 与 prefill 长度无关(常数);Transformer 线性增长
极端对比 16k prefill 下最大 xLSTM 的 step time < 最小 Transformer 的 step time

通用规律

  • Compute-optimal 模型的"loss vs 模型大小"关系在 xLSTM 和 Transformer 间近似落在同一条线上——暗示性能与模型大小存在跨架构的普适关系

亮点

  • 全面系统性: 672 次训练覆盖 5 个数量级算力,同时考量训练 + 推理 + 上下文长度
  • 精确 FLOP 计算: 告别 \(6ND\) 近似,为线性/二次架构对比提供公平基准
  • 实用指导: 证明 xLSTM 在过训练 regime 下指数稳定,支持"小模型 + 大数据"的实际部署策略
  • 推理建模: 基于 roofline 的理论模型与实测高度吻合

局限性 / 可改进方向

  • 仅考虑 cross-entropy loss,未评估下游任务(推理、代码、多语言等)
  • 未涉及 MoE 或 Attention+xLSTM 混合架构
  • 推理实验限于单 GPU,未考量多 GPU 分布式推理场景
  • 训练数据仅用 DCLM-Baseline,未验证数据分布变化的影响
  • 未探讨 xLSTM 在超长上下文(>16k)下的实际质量表现(如 recall 能力)
  • 未与 Mamba、RWKV 等其他线性架构做横向对比

与相关工作的对比

  • Chinchilla (Hoffmann 2022): 本文复现了 Transformer 的 compute-optimal 指数,并扩展到 xLSTM
  • Gadre 2024 / Sardana 2024: 本文在过训练 regime 分析上与之一致,但增加了跨架构维度
  • Shen 2024: 展示线性模型与 Transformer "on par",本文更进一步证明 xLSTM "优于" Transformer
  • Poli 2024: 混合架构优于纯 Transformer;本文证明纯线性架构也能胜出
  • Porian 2024: 本文复现了其 Transformer power-law 指数 \(a\)

启发与关联

  • xLSTM 的 Pareto 支配性意味着在同等算力下可获得更好的预训练模型,对资源受限场景特别有价值
  • 上下文长度对 compute-optimal 模型大小的影响是一个被广泛忽略的维度,值得在其他架构(Mamba、RWKV 等)中验证
  • 推理优势随上下文增长而扩大,暗示在长上下文推理(如 CoT、文档理解)中线性架构潜力巨大
  • 跨架构的"模型大小 vs loss"普适关系是一个值得深入研究的理论问题
  • 精确 FLOP 计算方法论可直接复用于评估 Mamba、RWKV、RetNet 等其他线性架构的 scaling 行为
  • 过训练 regime 指数恒定这一发现为 "小模型多数据" 部署策略提供了理论保障

评分

  • 新颖性: ⭐⭐⭐⭐ — 首个系统性线性复杂度 vs Transformer scaling law 对比
  • 实验充分度: ⭐⭐⭐⭐⭐ — 672 次训练、多维度分析、理论+实测推理建模
  • 写作质量: ⭐⭐⭐⭐ — 结构清晰,图表专业
  • 价值: ⭐⭐⭐⭐ — 为线性复杂度架构的工程部署提供了重要的 scaling 指导
  • 综合: ⭐⭐⭐⭐ — 实验扎实、结论清晰,对架构选型有直接参考价值