通过子空间优化，让Transformer在设备端高效训练
Efficient Resource-Constrained Training of Transformers via Subspace Optimization
- 将训练限制在关键信息子空间，减少内存占用
- 内存使用降低62倍，计算量减少2倍，精度几乎不变
- 适合资源受限设备上的模型训练，如树莓派
随着人工智能日益融入日常生活，能源消耗与数据隐私成为迫切问题。在设备端进行学习可降低能耗并保护数据隐私，但现代神经网络规模的扩大给设备端训练带来重大挑战。尽管已有研究集中于轻量级卷积架构，本文则将基于子空间的训练方法应用于Transformer模型。受‘模型核心信息存在于固定子空间’这一思想启发，提出权重-激活子空间迭代（WASI）方法，通过限制训练仅在该子空间内进行，缓解反向传播的内存瓶颈，并提升推理效率。实验表明，WASI在保持与常规训练相当精度的同时，内存使用最高降低62倍，计算成本（FLOPs）最多降低2倍。在Raspberry Pi 5上，WASI实现约1.4倍于常规训练的训练与推理速度。代码已开源：https://github.com/Le-TrungNguyen/ICLR2026-WASI.git。
原文摘要 · Abstract (English)
As AI increasingly shapes daily life, energy consumption and data privacy have become pressing concerns. On-device learning trains models directly on edge devices, cutting energy consumption and safeguarding data privacy. However, the expanding scale of modern neural networks creates a major obstacle for on-device training. Although prior work has concentrated on compact convolutional architectures, we instead apply subspace-based training to transformer models. Motivated by the idea that a model's essential information lies in a fixed subspace, we introduce Weight-Activation Subspace Iteration (WASI), a method that mitigates the memory bottleneck of backpropagation and boosts inference efficiency in transformer models by restricting training to this subspace. Our results demonstrate that WASI maintains accuracy comparable to vanilla training while reducing memory usage by up to $62\times$ and computational cost (FLOPs) by up to $2\times$. On a Raspberry Pi 5, WASI achieves roughly $1.4\times$ faster training and inference than vanilla training. The code is available at https://github.com/Le-TrungNguyen/ICLR2026-WASI.git.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。