ssToken: Self-modulated and Semantic-aware Token Selection for LLM Fine-tuning

An important design question is which layer's attention matrix should be used to compute the attention scores. In Appendix B, we conduct ablation studies comparing early, middle, and deep layers, and conclude that using deeper layers generally yields better results. This finding aligns with prior studies (Aljaafari et al., 2024; Zheng et al., 2024; Rocchetti & Ferrara, 2025), which suggest that semantic representations become increasingly abstract across layers: shallow layers (closer to the input) primarily capture syntax, local structures, and surface-level patterns such as positional relations, bracket matching, and syntactic cues, while deeper layers (closer to the output) focus more on semantic abstraction, high-level concepts, and task-relevant global information, which are typically more influential for instruction following. Moreover, prior works in other domains that rely on attention scores for analysis (Chen et al., 2024; Y e et al., 2025) are often incompatible with efficient attention implementations such as FlashAttention (Dao et al., 2022). To avoid this, we design a lightweight solution: during the forward pass, we use a hook to store the hidden states of the target layer and then perform a simple recomputation of that layer to retrieve its attention matrix. This design eliminates the need to output full attention matrices during the complete forward pass, thereby making our algorithm fully compatible with efficient attention mechanisms like FlashAttention and ensuring training efficiency.