AegisGuard: RL-Guided Adapter Tuning for TEE-Based Efficient & Secure On-Device Inference

We very thank the reviewer PYU5 for the valuable feedback on our paper. Our responses to the comments are below.

Question 1: Is there latency overhead for the sparse metrics operations after structural pruning?

Answer 1: No, there is no significant latency overhead introduced by sparse matrix operations in our method. This is because we apply structural pruning, specifically by removing attention heads, rather than performing unstructured pruning at the level of individual elements. As a result, the remaining weights can still be represented and computed using dense matrix operations, without introducing sparsity. Therefore, our method preserves the efficiency of dense computation and does not suffer from the irregular memory access patterns or hardware inefficiencies commonly associated with sparse operations.

Question 2: What is the tradeoff and benefit from the compression?

Answer 2: The tradeoff of compression is that it requires longer finetuning time before model deployment, it's about 1.5 $\times$ fine tuning time costs than normal fine tuning without compression. However, this process only occurs once and does not introduce repetive overhead after deployment. It does not introduce overhead to the online inference process.

The benefit is a lower computation workload in TEE in the online phase. We provide a breakdown of inference latency components (GPU, TEE, and Transfer) in Table 1 of Section 4.2. For example, in the case of ViT-Base, the GPU processing time is 5.94 ms, TEE execution takes 4.31 ms, and data transfer costs 4.89 ms per batch. These results confirm that compared to Shield-LoRA, AegisGuard achieves 2-3 $\times$ reductions in TEE execution and transfer overhead. Additionally, we also added the TEE computation costs reduction of pruning is presented below separately. The results are conducted on the ViT-Base model due to rebuttal time constraints. This experiment measures the TEE execution time (averaging 10 batches) under different pruning ratios, demonstrating how pruning directly reduces the secure enclave’s computation cost. We will put this ablation study into appendix in the next version.

We thank the reviewer RwdN for the valuable feedback on our paper. Our responses to the comments are below.

Question 1: The robustness of the layer sensitivity estimation.

Answer 1: Our layer sensitivity estimation method is robust to the choice of input samples, as shown by its convergence behavior, stability across random seeds, performance when experimented on different domain datasets, and its design based on parameter-space perturbations. We support this claim in the following:

• The entire RL-guided fine-tuning is a convergence process that runs over many random batches, not a single sample. The overall objective of our method is to fine-tune a large model that would converge to an optimal loss with a fixed number of layers. Note that during our designed RL-guided fine-tuning process, the model is not just exposed to a single random batch, but multiple random batches throughout. Additionally, our weighted layer selection sampling mechanism is designed to differentiate layer wise sensitivity with more sensitive layers (those with higher sensitivity scores) being fine-tuned more frequently. Therefore, if our estimation were highly sensitive to the batch composition, the selected sensitive layers would fluctuate throughout training, leading to instability or divergence. In contrast, we observe that the final adapter selection converges, resulting in successful downstream task performance. This confirms that sample-specific variance is smoothed out through repeated exposure and weighted selection during RL optimization.

• Layer sensitivity is stable or converges across different datasets. The results presented in Appendix D.3, Figure 6, also demonstrate this point. The sensitive layer selection process is stable and shows convergence in experiments conducted on domain-specific tasks across different random seeds. This shows that our method is robust to different task domains.

• Our perturbation is injected into LoRA weights rather than inputs, and is larger than the perturbations injected into inputs in previous works. Since the core design concept of our method is to disrupt the impact of specific LoRA layer weights rather than subtly perturbing inputs to enhance training robustness, we can inject large noise into it without considering model performance. This observation is drawn from the analysis of weight magnitude distribution presented in Figure 5 of the Appendix. LoRA is designed to introduce small, low-rank updates to the base model, and its parameter magnitudes are significantly smaller. This ensures that our perturbation can be large to effectively diminish the contribution of the fine-tuned LoRA layer while preserving the functionality of the original base model. Thus, the influence of the perturbed layer can be isolated and evaluated robustly, without being influenced by the chosen samples.

Question 2: In Table 2, why ViT-Base C100 Shield-LoRA (19.25) was not the best (colored rectangle) but AegisGuard (22.67)? Same questions for LLaMA-7B.

Answer 2 In Table 2, Shield-LoRA is not a baseline defense solution so we don't highlight Shield-LoRA in all settings. It is only a reference point (security upper-bound of black-box setting) for comparison. The colored rectangle is only for the best-performing defense among the baselines and the proposed solution.

Limitation 3: The limited scenario of adapter-based LoRA fine-tuning.

Answer 3 We focus on LoRA as it is the most commonly used parameter efficient fine-tuning method currently. Especially in the privacy field, many works[a][b] design their privacy methods based on LoRA-based fine-tuning scenarios. Additionally, many industrial on-device practices such as Apple[d] and Nvidia[e] have also adopted LoRA as a mainstream fine-tuning method. For example, Apple developed a new framework using LoRA adapters that incorporates a mixed 2-bit and 4-bit configuration strategy to maintain model quality. Nvidia developed the NeMo framework that mainly supports dynamic multi-LoRA inference, enabling simultaneous inference requests with different LoRA models.

Besides, compared with normal adapters, LoRA contains less trainable parameters and are thus more compatible with TEE contraints (memory and computation constraints[c]). When we adopt TEE to provide black-box level security for LLM fine-tuned with private or sensitive data, efficiency is also a key objective. LoRA achieves this by significantly reducing the number of trainable parameters through low-rank matrix decomposition, compared with traditional adapter-based techniques. It not only improves computational efficiency but also achieves strong performance in modern applications, making it well-suited for TEE-based secure fine-tuning. As a result, in our setting (where resource efficiency and security are critical), LoRA is a practical choice.

[a].Sun, Youbang, et al. "Improving LoRA in Privacy-preserving Federated Learning." The Twelfth International Conference on Learning Representations(ICLR2024).

[b].Luo, Zihao, et al. "Privacy-Preserving Low-Rank Adaptation Against Membership Inference Attacks for Latent Diffusion Models." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 39. No. 6. 2025.

[c].Sun, Zhichuang, et al. "Shadownet: A secure and efficient on-device model inference system for convolutional neural networks." 2023 IEEE Symposium on Security and Privacy (SP). IEEE, 2023.

[d].Apple. "Introducing Apple’s On-Device and Server Foundation Models"

[e].NVIDIA. "Parameter-Efficient Fine-Tuning with NVIDIA NIM for LLMs"

Thanks for providing valuable feedback on our paper. Our responses to the comments are below.

Question 1: How is the impact of pruning on accuracy compensated by the next step of fine-tuning? If it's because pruning is losing very little, then where is the limit? How much more pruning can be done to actually affect downstream accuracy?

Answer 1: The accuracy drop of pruning can be compensated by the next step of fine tuning is because we adopt a mixed strategy of pruning and fine-tuning. Pruning is periodically applied during the fine-tuning process, and the pruned weights are permanently removed (i.e., not reconnected). After each pruning step, the remaining weights continue to be updated, progressively learning more task-relevant information during subsequent training. This allows the model to gradually adapt to the reduced parameter space and partially recover accuracy.

In contrast to traditional structural pruning methods that prune uniformly across all layers [a], our method adopts a dynamic and selective pruning strategy, focusing primarily on the shielded layers, those would execute inside the TEE. Consequently, the total number of pruned parameters is relatively small compared to conventional pruning methods, which helps preserve downstream task accuracy.

To better understand the pruning limits, we conduct additional experiments under various pruning ratios (e.g., 20%, 30%, and 50%) on ViT-Base and report the corresponding interplay and downstream accuracy.

As shown above, AegisGuard's performance remains stable at 20% and 30% compared with baseline, but it has a obvious degrade (around 4% drop) at 50% pruning ratio. These results help delineate the boundary of effective pruning in our setting.

Weakness 2： Ablation study of latency overhead.

Answer 2: We provide a breakdown of inference latency components (GPU, TEE, and Transfer) in Table 1 of Section 4.2. For example, in the case of ViT-Base, the GPU processing time is 5.94 ms, TEE execution takes 4.31 ms, and data transfer costs 4.89 ms per batch. These results confirm that compared to Shield-LoRA, AegisGuard achieves 2-3 $\times$ reductions in TEE execution and transfer overhead.

To further quantify the effect of pruning within the TEE, we include an ablation study as shown below to illustrate the execution cost of ViT-Base. This experiment measures the TEE execution time (averaging 10 batches) under different pruning ratios, demonstrating how pruning directly reduces the secure enclave’s computation cost. We will put this ablation study into appendix in the next version.

Weakness 3: The discussion of the design philosophy of the chosen approach compared to other methods.

Answer 3: One representative approach is gradient-based methods[b][c], but they can not be applied to our scenario to estimate layer/adapter sensitivity. It is because these methods estimate the sensitivity based on layer/adapter gradients and generally assume that all layers are updated simultaneously during training.

In our setting, we adopt a dynamic fine-tuning strategy, where only a subset of layers/adapters are activated and updated at each training step, that the rest are frozen. This strategy is to identify which layers/adapters are more sensitive, and can better absorb task-specific features. As a result, gradient-based sensitivity estimation becomes infeasible, since gradients are not available for the frozen layers at any estimation step.

To address this limitation, we propose a perturbation-based sensitivity estimation method. This approach estimates a layer's influence by observing the loss change when a perturbation is applied to its parameters. It is straightforward to implement in the dynamic fine-tuning setting, computationally efficient. Additionally, to mitigate variance from single batche samples, we perform sensitivity estimation at fixed intervals (e.g., every 20 steps), ensuring more stable and robust evaluations. Overall, our design philosophy prioritizes and can better compatibility with our dynamic fine-tuning strategy.

Weakness 4: Format issues

Answer 4: Thank you for pointing out the format issues. We will revise the formatting of fonts, as well as improve the layout and flow of figures and tables to enhance the readability and presentation quality in the next version.

[a]. Ma, Xinyin, Gongfan Fang, and Xinchao Wang. "Llm-pruner: On the structural pruning of large language models." Advances in neural information processing systems 36 (2023): 21702-21720

[b]. Sahoo, Sabyasachi, et al. "A layer selection approach to test time adaptation." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 39. No. 19. 2025.

[c]. Lee, Yoonho, et al. "Surgical Fine-Tuning Improves Adaptation to Distribution Shifts." The Eleventh International Conference on Learning Representations.