TokenSeek: Memory Efficient Fine Tuning via Instance-Aware Token Ditching

TL;DR

TokenSeek tackles the memory bottleneck in fine-tuning large language models by introducing instance-aware token seeking and efficient token ditching, a universal plugin for Transformer-based models. By scoring tokens with context and gradient signals and updating only the selected ones, it achieves substantial activation-memory reductions while maintaining or improving performance, and it remains compatible with PEFT techniques like LoRa and QLoRA. The approach is architecture-agnostic, interpretable, and demonstrates strong generalization across model scales (0.5B–7B) and datasets, with notable gains on Llama3.2 and Qwen models. Overall, TokenSeek offers a practical path to memory-efficient LLM adaptation suitable for resource-constrained environments, contributing both to performance and to insights into token-level efficiency.

Abstract

Fine tuning has been regarded as a de facto approach for adapting large language models (LLMs) to downstream tasks, but the high training memory consumption inherited from LLMs makes this process inefficient. Among existing memory efficient approaches, activation-related optimization has proven particularly effective, as activations consistently dominate overall memory consumption. Although prior arts offer various activation optimization strategies, their data-agnostic nature ultimately results in ineffective and unstable fine tuning. In this paper, we propose TokenSeek, a universal plugin solution for various transformer-based models through instance-aware token seeking and ditching, achieving significant fine-tuning memory savings (e.g., requiring only 14.8% of the memory on Llama3.2 1B) with on-par or even better performance. Furthermore, our interpretable token seeking process reveals the underlying reasons for its effectiveness, offering valuable insights for future research on token efficiency. Homepage: https://runjia.tech/iclr_tokenseek/

Figures (10)

• Figure 1: Motivation behind TokenSeek and its preliminary comparison. (a) Breakdown of training memory under different batch side settings, revealing that activations are the primary bottleneck in training memory consumption. (b) Effective and efficient TokenSeek (ours) $vs.$ concurrent arts in performance and memory consumption on Llama3.2 1B (detailed results in Tab. \ref{['tab:1']})
• Figure 2: Overview of TokenSeek (ours) $vs.$TokenTune frameworks. (a) Instance-aware token seeking using context and gradient information (see §\ref{['sec:tokenseeking']} and Eq. \ref{['eq:all']}). (b) Efficient token ditching (see §\ref{['subsub:ETD']}). (c) TokenTune for random token selection (see analysis in Tab. \ref{['tab:1']} and §\ref{['sub:explain']}).
• Figure 3: The performance for the Llama3.2 1B with QLoRA setting, where the upper, lower and middle line indicate the maximum, minimum and the average results.
• Figure 4: Case study of a training instance. (a) Visualization of the top 50% selected tokens using context and gradient information, highlighted in red and blue, respectively. (b) Average attention map from the final layer. (c) Accumulated gradient map of activations in the penultimate layer. (d) Context importance scores obtained by column-wise accumulation. (e) Gradient importance scores obtained by summing across the hidden dimension. Additional visualizations are provided in §\ref{['appendix:visualization']}.
• Figure 5: Training loss curves of six settings on Qwen2.5 0.5B.Blue, red, and orange lines represent full parameter, LoHa, and QLoRA tuning, respectively. Lighter lines in each group indicate 10% token tuning, while darker lines indicate 50%. Detailed performance and memory usage results are provided in Tab. \ref{['tab:1']} and \ref{['tab:abla']}.