Efficient Resource-Constrained Training of Transformers via Subspace Optimization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Nguyen, Le-Trung, Tartaglione, Enzo, Nguyen, Van-Tam
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915829490122752
author Nguyen, Le-Trung
Tartaglione, Enzo
Nguyen, Van-Tam
author_facet Nguyen, Le-Trung
Tartaglione, Enzo
Nguyen, Van-Tam
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_09160
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Efficient Resource-Constrained Training of Transformers via Subspace Optimization
Nguyen, Le-Trung
Tartaglione, Enzo
Nguyen, Van-Tam
Machine Learning
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.
title Efficient Resource-Constrained Training of Transformers via Subspace Optimization
topic Machine Learning
url https://arxiv.org/abs/2510.09160