SMAAT: Scalable Manifold-Aware Adversarial Training for Large Language Models

We thank the reviewer for their feedback and bringing to our attention the adaptive attack work.

lack of evidence for scalability: does the hypothesis of intrinsic dimensionality still hold for larger language models?

Response: We conducted additional experiments on the fine-tuned version of the quantized LLAMA-2 7B model, employing PEFT. Our results indicate that the intrinsic dimensionality of the model displays a non-monotonic pattern. Consistent with our hypothesis, integrating SMAAT into this model results in choosing the first layer to perform adversarial training, which does not provide additional scalability benefits. We nevertheless performed AT at the last layer and, as expected, this did not yield improved robustness.

ablation study of adversarial training on different layers: since this paper proposes to only fine-tune the last layer, what will happen if we include more layers for adversarial training? do we gain better robustness since more model parameters are included for adversarial training?

Response: We evaluated the robustness and generalization impact of adversarial training in the intermediate layers. Our findings are presented in the newly added Figure 4 of the revised paper. In alignment with our main finding, when the adversarially trained layer is positioned closer to the input layer robustness of the model decreases as higher intrinsic dimensionality makes it more difficult to generate off-manifold examples.

lack of adaptive attack: I strongly recommend the authors to design the adaptive attack [1] to approximate the lower bound of empirical adversarial robustness under their defense. This paper only evaluates model robustness under a simple 5-step PGD attack, which is not enough.

Response: To assess the robustness of our method under more targeted attacks, we considered the following three strategies: 1. enhancing the strength of PGD attacks by increasing the steps to 50 and targeting the most vulnerable class, thereby generating attack samples more potent than those utilized during training; 2. implementing targeted PGD attacks focused on the internal representations of an intermediate layer (“feature-level attack” introduced in [1]), rather than targeting the output label, to more effectively exploit our approach which primarily bolsters robustness at the final layer. The robustness results presented in Table 3 of the revised manuscript are the averaged results of these two attacks instead of the 5-step PGD attack applied earlier.

References

[1] Sara Sabour, Yanshuai Cao, Fartash Faghri, and David J Fleet (2015). Adversarial manipulation of deep representations arXiv preprint arXiv:1511.05122, 2015

We appreciate your valuable feedback and constructive comments.

The three observations mentioned in the abstract appear not to be discussed in detail in the paper. How do these observations motivate the methodology?

Response: In our revision, we have further clarified how the three fundamental observations motivated our approach. For this, we expanded Section 3.0 to also include our intuitions that drove our problem formulation.

How effective are SMAAT generated adversarial examples compared to standard adversarial examples in attacking LLMs?

Response: AEs generated in the last layer are expected to be more potent compared to AEs generated at the initial layer as attack strength can be more freely adjusted. To test this, we deployed the AEs generated by SMAAT at the last layer by feeding them to the last layer of the BERT model. Our results obtained on three datasets show that the robustness of the standard BERT model reduced to 0 in all cases. This suggests that AEs generated with SMAAT are more effective than those created in the initial layer.

Are the baseline and adversarial training models being compared in the paper using the same augmentation or training strengths? The strength and budget of the baselines appear not to be presented in the paper.

Response: The strength of AEs in each model varies among methods. PGD-based AT methods that rely on initial layer training are subject to low epsilon values. Whereas our approach, which generates AEs in the last layer, is not subject to such constraint. In contrast, for methods like TMD, SAFER, and RSMI the strength of AE cannot be defined in the same manner. Therefore, in our experiments, we adhered to the hyperparameters recommended in their respective original publications. All models were trained for an equivalent number of epochs.

We believe overall training complexity of these methods can best be assessed in terms of their training run-times. In this regard, SMAAT brings only a marginal run-time overhead to the clean models, which took only a couple of second for BERT and RoBERTa.

It is unclear about the generalizability of the method. Will the proposed method be effective against other types of attacks such as typos?

Response: To more comprehensively evaluate the effectiveness of our method against a wider array of attacks, we carried out tests using the datasets included in the AdvGLUE benchmark. This benchmark encompasses AEs generated from 17 distinct attack types, typo based TextBugger included. The results of these tests are shown in Table 3 of the revised manuscript. The performance of our method mirrors the trends observed in our preliminary assessments. SMAAT exhibits average robustness improvements of 5.6% and 2.6% on the two new benchmarks for BERT, and 12.1% and 2.5% for RoBERTa, in comparison to standard and FreeLB++ models, respectively, while maintaining comparable generalization.

We thank the reviewer for their feedback on our work.

I am confused about the definition of ‘off-manifold’ and ‘on-manifold’. It would be better for the authors to provide more clear definitions and high-level explanations to help understand.

Response: The manifold hypothesis stands as one of the most compelling explanations for the susceptibility of deep neural networks to adversarial samples. This hypothesis fundamentally posits that data resides on a low-dimensional manifold within a high-dimensional representation space, and that a network, during training, learns to approximate this manifold. Consequently, an off-manifold sample, deviating from this foundational manifold, leads to undefined behavior in the network. Accordingly, an off-manifold sample is one that diverges from the underlying manifold and the network’s behavior is undefined. For better clarity, we have incorporated this explanation into Section II (last paragraph of page 3) of the revised paper.

I am confused about the choice of $l^*=n$. In the objection function in Eq. (10), it seems that any $l^*$ satisfies when $ID(i-1) , ID(i), \forall i < l^*$. Based on my observation of Figure 3, $l^*$ can be chosen any number smaller than 13 since ID is monotonically decreasing. Therefore, based on Eq. (10), I cannot understand why the choice of $l^*=n$ is optimal.

Response: In our definition, the optimal layer is the one that minimizes the length of the back-propagation path. Therefore, it is preferable to perform adversarial training at output layers as opposed to input layers, which will enhance scalability. The main finding of our research is the demonstration that when the intrinsic dimensionality of the feature manifold decreases monotonically, adversarial training can be effectively performed at the layer with the lowest intrinsic dimension. As evidenced by the data in Figure 4, the final layer—the one with the highest index—exhibits the lowest intrinsic dimension. Therefore, implementing adversarial training at this layer not only optimizes generalization and robustness but also accomplishes these goals with minimal computational cost.

Empirical results are limited. The evaluation should be conducted on various datasets in the GLUE benchmark.

Response: We further evaluated our method on four datasets in the GLUE benchmark (including SST2, QQP , QNLI, RTE). Rather than performing adversarial attacks separately, as in Table 1, we preferred to use the AdvGLUE benchmark which includes the adversarial counterparts of the same for four datasets. The outcomes of these evaluations are presented in Table 3 of the revised paper. Notably, the patterns identified in other benchmarks are consistently observed here as well. Overall, SMAAT exhibits average robustness improvements of 5.6% and 2.6% on the two new benchmarks for BERT, and 12.1% and 2.5% for RoBERTa, in comparison to standard and FreeLB++ models, respectively, while maintaining comparable generalization.

The title seems to overclaim the contribution of this paper. All the results in this paper are shown on two language models, i.e., BERT and RoBERTa. However, I did not see the results of ‘large language models (LLMs)’ such as Llama2-70b. I am wondering whether the proposed method can be scalable to LLMs. If so, please show the empirical justification.

Response: We extended our experiments to include a fine-tuned version of the quantized LLAMA-2 7B model, utilizing PEFT. In this process, we discovered that the intrinsic dimensionality of the representations generated by this model is non-decreasing, mirroring the trend observed in vision models. Additional information about these experiments is presented in Appendix Section H.

Thank you for carefully reading our paper and for the suggestions. We corrected the mathematical representation and revised the proof of Theorem 3.1 accordingly.

As mentioned in the paper, the basis $U_l$ obtained through SVD is an orthonormal basis. Thus, $U_lU_L^T = I$, which makes the projection error as defined in the paper always equal to zero.

Response: We compute the projection error using the eigenvectors corresponding to the top-k eigenvalues, where k represents the ID of the respective layer. Therefore, samples aligned with the data manifold exhibit lower projection errors, while others deviate with higher errors in the k-dimensional data manifold. Our notation is modified to include the index $k$, representing the top-k eigenvectors, which translates into the effective ID of the layer.

Can the authors clarify how to derive equation 7 from equation 6 ? It is critical as it links the search for the optimal layer $l^*$ to theorem 3.1.

Response: To enhance the clarity of the derivation, we introduced an intermediate (Eq. 7). We also note that Eq. 8 (old Eq. 7) sets a sufficient condition to satisfy Eqs. 6 and 7.

In theorem 3.1, the proof in Appendix is given for the opposite side: If $\|(I-U_{(i-1)}U_{(i-1)}^T)\delta_{(i-1)}\| < \|(I-U_iU_i^T)\delta_(i)\|$ then rank(i-1) < rank(i). There is also the equality case that is included in the theorem but not in the proof. Maybe the inequalities should be reversed in the formulation of the theorem, which would make a proof by contraposition ? This would also make sense with the remaining of the paper, since we observe decreasing ID and not increasing ID.

Response: We appreciate your insightful observation. We have carefully revised Theorem 3.1 and its proof. The notation is now coherent and the direction of the inequalities are correct.

Similarly, for the objectives of SMAAT (eq. 9 and 10), the inequalities on the ID do not make sense since we observe decreasing ID throughout the layers. I think the inequality should be reversed as well, i.e. $ID(i) < ID(i-1)$

Response: In accordance with our initial description in the text, the direction of inequalities are corrected.

What is k and $U_l^k$ in the computation of ID (equation 11) ? I assume it is the k first rows and columns of $U_l$, which would make ID close to the rank of the representation. This should be clarified, as well as the link between ID and the rank, since the theorem is stated using the rank operator.

Response: This is correct: the variable $k$ represents the number of top eigenvalues needed to obtain a projection similar to the original samples and also used as the ID of features. We revised the paper for better clarity.

What is $\lambda_{max}$? Is it the maximum eigenvalue? If so, I'm unsure about the assumption used in the proof of theorem 3.1 about the ratio of $\lambda_{max}$. There is an additional layer in the network from which the jacobian is computed in the numerator, thus the Lipschitz constant of the numerator is greater than the one of the denominator, making the ratio greater than 1. Does the remaining of the proof holds with this?

Response: We have updated the mathematical formulation of the theorem to be consistent with the revised assumptions. In the new version of the theorem, we employed a proof by contraposition to establish its validity. We believe these modifications greatly contribute to the clarity and accuracy of our work.

The proposed approach is interestingly the opposite of YOPO [1], can the authors develop on the link and differences regarding this method ?

Response: We believe this seeming contradiction is due to differences between vision and text domains. In Section G (Figure 6) of the revised paper, we present the measured intrinsic dimensionality of the ViT model which is noticeably different from that of text models. In fact, our approach suggests that the optimal layer $l^*$ for the ViT model would be the initial layer due to the increasing ID. We have incorporated this clarification into the paper to provide a more nuanced understanding of the relationship between layer-wise behavior and robustness.

Typos

Response: We corrected typos. However, in Eqs. 4 and 10 (old Eq. 9), we define an AE to be off-manifold in any layer, and not in all of the layers. Therefore, we kept $\exists$ notation instead of $\forall$.

Thank you for your time and suggestions.

Argument about the manifolds between the layers is not very clear to the reader.

Response: We revised our statement in Section I (next to last paragraph on page 2) to better clarify the manifold relation between layers.

Results on Table 1 are difficult to interpret. I would suggest just to boldify the best result for every attack in each dataset without having any underlined results.

Response: The best results in Table 1 are in bold.

I believe that the method is probably not very novel or of high contribution.

Response: We would like to emphasize that the novelty of our method lies in (i) the discovery that the intrinsic dimensionality of internal representations in fine-tuned decoder models monotonically decreases and that, in this case, (ii) applying AT in the higher layers covers all the AEs from the previous layers and offers better robustness than applying AT in the initial layers. This justifies (iii) applying PGD-based adversarial training to only to the last layer. Our method not only improves robustness and generalization capabilities but also achieves these improvements at a significantly lower computational cost compared to state-of-the-art methods. In our revision, we have extended our experimental analysis to include more tasks/datasets, thereby reinforcing the efficacy of our method.

Have you conducted any experiments where you use adversarial examples generated from other layers instead of the only the last one. Sometimes theorems and experiments can provide very different outcomes.

Response: As part of our revision, we conducted additional experiments focusing on the application of adversarial training to intermediate layers. The outcomes of these experiments are presented in Figure 4 of the revised manuscript. In agreement with our hypothesis, the results demonstrate a reduction in robustness when the adversarially trained layer is closer to the input layer, and the best robustness is achieved in the last layer.

I read the limitations at the end of the paper. Although this is not in the scope of this work I would suggest to still try this method with image data.

Response: We extended the application of our method to (i) vision and (ii) LLM domains. To achieve this, we evaluated the intrinsic dimensionality of the ViT model and the quantized LLAMA-2 7B model. Comprehensive details about these models and the outcomes of these experiments are presented in Appendices G and H. Our findings reveal that both models have non-monotonic variations in intrinsic dimensionality. We further applied SMAAT to both models to test our hypothesis. We did not observe enhanced robustness, as expected by our formulation.