POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation

📄 arXiv: 2603.05500v1 📥 PDF

作者: Zeju Qiu, Lixin Liu, Adrian Weller, Han Shi, Weiyang Liu

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-03-05

备注: Technical report v1 (14 pages, 7 figures, project page: https://spherelab.ai/poetx/)

💡 一句话要点

POET-X：通过缩放正交变换实现内存高效的大语言模型训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 正交等价训练 内存效率 计算效率 模型训练 参数重整 GPU训练

📋 核心要点

  1. 现有大语言模型训练方法面临内存消耗高和计算开销大的挑战，限制了模型规模和训练效率。
  2. POET-X通过缩放正交变换，在保持训练稳定性的同时，显著降低了计算成本和内存占用。
  3. 实验表明，POET-X能够在单个GPU上预训练数十亿参数的LLM，优于传统优化器。

📝 摘要（中文）

高效且稳定的大语言模型（LLM）训练仍然是现代机器学习系统中的核心挑战。为了解决这个问题，参数重整正交等价训练（POET）被提出，它是一种谱保持框架，通过正交等价变换优化每个权重矩阵。虽然POET提供了强大的训练稳定性，但其原始实现由于密集的矩阵乘法而导致高内存消耗和计算开销。为了克服这些限制，我们引入了POET-X，这是一种可扩展且内存高效的变体，以显著降低的计算成本执行正交等价变换。POET-X保持了POET的泛化和稳定性优势，同时在吞吐量和内存效率方面实现了显著改进。在我们的实验中，POET-X能够在单个Nvidia H100 GPU上预训练数十亿参数的LLM，相比之下，诸如AdamW之类的标准优化器在相同设置下会耗尽内存。

🔬 方法详解

问题定义：论文旨在解决大语言模型（LLM）训练过程中内存消耗过高和计算开销过大的问题。现有的正交等价训练方法（如POET）虽然能保证训练的稳定性，但由于涉及大量的矩阵乘法，导致内存占用和计算复杂度显著增加，限制了其在大规模模型上的应用。

核心思路：POET-X的核心思路是通过缩放正交变换来降低正交等价训练的计算复杂度。具体来说，它不再直接进行完整的正交变换，而是通过一系列更小规模的变换来逼近，从而减少了矩阵乘法的计算量。这种方法在保证训练稳定性的前提下，显著降低了内存占用和计算开销。

技术框架：POET-X的技术框架仍然基于正交等价训练的思想，但其核心在于如何高效地实现正交变换。整体流程可以概括为：首先，对权重矩阵进行分解；然后，对分解后的子矩阵进行缩放的正交变换；最后，将变换后的子矩阵重新组合成更新后的权重矩阵。关键在于如何选择合适的分解方式和变换策略，以在计算效率和训练稳定性之间取得平衡。

关键创新：POET-X最重要的技术创新点在于其缩放的正交变换方法。与传统的正交等价训练方法相比，POET-X避免了大规模的矩阵乘法，而是通过一系列更小规模的变换来逼近，从而显著降低了计算复杂度。这种方法在保证训练稳定性的前提下，实现了内存效率和计算效率的提升。

关键设计：POET-X的关键设计包括：1) 权重矩阵的分解方式，例如可以使用SVD或QR分解；2) 缩放因子的选择，需要根据具体的模型和数据集进行调整，以保证训练的稳定性和收敛速度；3) 正交变换的具体实现方式，可以使用Givens旋转或Householder变换等方法。此外，损失函数和优化器的选择也需要与POET-X的特性相适应。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，POET-X能够在单个Nvidia H100 GPU上预训练数十亿参数的LLM，而传统的AdamW优化器在相同设置下会耗尽内存。这表明POET-X在内存效率方面具有显著优势。此外，POET-X在保持训练稳定性的同时，还实现了吞吐量的提升，进一步验证了其在计算效率方面的优势。

🎯 应用场景

POET-X具有广泛的应用前景，尤其是在资源受限的环境下训练大型语言模型。例如，可以在单张GPU上训练更大规模的模型，或者在边缘设备上进行模型微调。此外，POET-X还可以应用于其他需要稳定训练和高效计算的机器学习任务，例如图像识别、语音识别等。其降低内存占用和计算开销的特性，使得大规模模型的训练和部署变得更加可行。

📄 摘要（原文）

Efficient and stable training of large language models (LLMs) remains a core challenge in modern machine learning systems. To address this challenge, Reparameterized Orthogonal Equivalence Training (POET), a spectrum-preserving framework that optimizes each weight matrix through orthogonal equivalence transformation, has been proposed. Although POET provides strong training stability, its original implementation incurs high memory consumption and computational overhead due to intensive matrix multiplications. To overcome these limitations, we introduce POET-X, a scalable and memory-efficient variant that performs orthogonal equivalence transformations with significantly reduced computational cost. POET-X maintains the generalization and stability benefits of POET while achieving substantial improvements in throughput and memory efficiency. In our experiments, POET-X enables the pretraining of billion-parameter LLMs on a single Nvidia H100 GPU, and in contrast, standard optimizers such as AdamW run out of memory under the same settings.