Beyond Low-rank Decomposition: A Shortcut Approach for Efficient On-Device Learning

Question 1: How our rank selection strategy might scale to transformer-based models or LLMs?

Our rank selection strategy is fully applicable to transformer-based models and LLMs with billions of parameters, and it incurs only a one-time cost.

The main idea of our strategy involves performing a brute-force search over a 2D search space, where one dimension represents the number of layers to fine-tune and the other represents the number of explained variance values considered. Increasing or decreasing either of these factors directly impacts the search time for the optimal solution. The total number of parameters or model architecture does not directly affect our rank selection strategy. As long as backpropagation is required for training, it is still applicable.

Furthermore, the choice of rank selection algorithm is not the key point of ASI (see our response to Weakness 2 of Reviewer dju2).

Question 2, Weakness 2 and Suggestion 3: Apply ASI to transformer models and LLMs.

We have conducted additional experiments as follows:

• Transformer models: We applied the same training strategy with ASI implemented on the linear layers within the MLP blocks:

  • Swin Transformer for classification tasks: We used Swin Transformer pretrained on ImageNet-1K from the PyTorch library and evaluated it on five different downstream datasets: CIFAR-10, CIFAR-100, Pets, Flowers, and CUB.
  • TinyLlama on BoolQ dataset: We extended our technique to TinyLlama, a large language model (LLM) with 1.1B parameters, using the BoolQ dataset, which consists of yes/no questions. Since TinyLlama is a very large model, applying $\text{HOSVD}_\varepsilon$ directly would be computationally infeasible given our available resources, making it impossible to construct a budget for comparison. Therefore, instead of using the proposed rank selection strategy, we set a fixed expected rank of 20 for ASI and compared it against vanilla training.

• Image segmentation: We applied the same training policy and downstream datasets as used in (Nguyen et al., 2024) and (Yang et al., 2023).

These results are available at the following anonymized link: https://imgur.com/a/qOsfYU5. Overall, ASI gives similar results to those in our submission. It outperforms other methods in both computational complexity and activation memory consumption. We will add these new results to the camera-ready version.

Please note that in addition to MCUNet, we also conducted numerous experiments with MobileNetV2, ResNet18, and ResNet34. The results of these experiments are presented in Tables 1 and 2 of our submission.

Question 3, Weakness 1, Suggestion 1 and Suggestion 2: Writing style and how ASI works.

• Writing style: Thank you for your feedback. We will carefully consider it and make appropriate revisions.

• How ASI works: Regarding the steps in our method, as described in Fig. 1:

  • Step 1: Before the training begins, we measure the perplexity $\mathcal{P}$ (defined in Eq. (7)) of the pretrained model for each predefined explained variance by feeding a minibatch of pretrained data. We then save $\mathcal{P}$ to a file.
  • Step 2: Using the saved file, we run a brute-force search algorithm to find the optimal ranks for each fine-tuned layer. The optimal rank is the one that satisfies Eq. (8) and Eq. (9), i.e., the rank that minimizes the total perplexity while ensuring the memory does not exceed the given budget $\mathcal{B}$.
  • Step 3: We use ASI to "compress" the activation map of each fine-tuned layer into a subspace with the rank corresponding to the one found in Step 2. These ranks remain fixed throughout the training process.

Question 4: Challenges in applying ASI to other architectures.

We do not foresee any issues applying ASI to models architecturally different than those we have tested. As long as the model requires backpropagation and includes convolutional/linear layers, ASI can still be applied. Would you suggest some architectures where you think ASI might face challenges? We would be happy to hear your thoughts.

We appreciate your review, below is the answer to your only question.

Question: How do we calculate training FLOPs?

Currently, we measure training FLOPs based on theoretical calculations, which consist of the sum of the FLOPs required for both the forward and backward passes. The necessary formulas for a convolutional layer (with similar derivations for linear layers) are derived in equations (13)–(17) of our submission:

• Vanilla training:

  • Forward FLOPs: Defined in Eq. (17).
  • Backward FLOPs: Defined in Eq. (16).

• ASI:

  • Forward FLOPs: $O_{\text{vanilla}} + O_{\text{ASI}}$, where $O_{\text{ASI}}$ is defined in Eq. (14).
  • Backward FLOPs: Defined in Eq. (15).

• $\text{HOSVD}_\varepsilon$:

  • Forward FLOPs: $O_{\text{vanilla}} + O_{\text{HOSVD}\varepsilon}$, where $O_{\text{HOSVD}\varepsilon}$ is defined in Eq. (13).
  • Backward FLOPs: Defined in Eq. (15), which is identical to ASI.

The key reason ASI achieves lower training FLOPs compared to vanilla training is that its low-rank gradient calculation significantly reduces computational costs. This reduction more than compensates for the minor overhead introduced by mapping activation maps to subspaces during the forward pass (see Fig. 5).

Weakness 1: Use of \citet.

Thank you for this note. We will revise it in the camera-ready version.

Weakness 2: Other rank search algorithms besides brute-force?

Yes, there are certainly alternative methods—for example, using dynamic programming, we might reduce the computational complexity from exponential to linear, at the cost of employing more memory.

However, in the context of on-device learning, the models used are typically small with a limited number of layers, and we are heavily constrained by memory. As a result, brute-force search remains computationally feasible while ensuring an optimal solution.

That said, rank selection itself is not the core contribution of ASI. It would be interesting to explore smarter methods for rank-selection, relying on the learning dynamics of Deep Models (eg. leveraging something like the Information bottleneck); however, this still remains an open research quest. We leave further exploration of rank search algorithms for future work.

Weakness 3: Apply ASI to transformer models and LLMs.

We appreciate your feedback. To address this, we expanded our experiments to Swin Transformer, image segmentation, and LLM task using TinyLlama with 1.1B parameters. For further details, please refer to Question 2 of Reviewer N3Rd.

Weakness 4: The problem of accuracy loss.

We acknowledge that accuracy drops due to compression—this is the tradeoff. As such, this phenomenon also occurs with all other related state-of-the-art techniques, including $\text{HOSVD}_\varepsilon$, Gradient Filter, LoRA, and its variants. The key point is that ASI achieves a superior Pareto curve compared to other methods (Fig. 4).

Weakness 5: Stability of activation map and combination with GaLore.

Thank you for suggesting GaLore. Our response is as follows:

• Stability of activation maps:

  • ASI is specifically designed for fine-tuning, where large learning rates are typically not used. As a result, the input to the activation function does not change drastically across training iterations, i.e., it remains stable within a small number of iterations.
  • Moreover, (Virmaux & Scaman, 2018) indicated that most commonly used activation functions have a Lipschitz constant of 1, meaning that if the input remains stable, the output (activation map) remains stable as well.

• Effect of GaLore:

  • Based on our understanding, GaLore does not directly affect the stability of activation maps.
  • GaLore is a gradient compression technique, while ASI specifically works on activation maps. In principle, both methods could be combined to maximize efficiency.
  • However, GaLore is particularly useful for Adam, where optimizer state becomes a concern. In contrast, for on-device learning, which is our primary focus, SGD without momentum is preferable since it eliminates the need to store optimizer state. Consequently, GaLore offers little benefit in this setting.

We acknowledge that GaLore is an interesting approach, and leveraging it could lead to promising results-and will be added in the related works section. However, we have not observed how it influences the stability of activation maps in fine-tuning. Could you please clarify your thoughts on this?