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

AI summary

Overview

• Research area: Natural Language Processing, specifically memory-efficient fine-tuning of large language models (LLMs).
• Technical level: Advanced. The paper assumes familiarity with Transformer internals, attention maps, backpropagation and the chain rule, activation memory accounting, and parameter-efficient fine-tuning (PEFT) methods such as LoRA, LoHa, and QLoRA.
• Scope in one sentence: The paper introduces TokenSeek, a drop-in plugin that scores individual training tokens using attention and gradient signals and then backpropagates only through the highest-scoring 10% of tokens, cutting fine-tuning memory while matching or exceeding full-token tuning accuracy across Qwen2.5 0.5B and Llama3.2 1B/3B.

What This Paper Is About

Fine-tuning LLMs is memory-hungry, and the paper states that intermediate activations, not weights or optimizer states, are the dominant cost (it reports activations at 87% of memory in Llama3 8B and 60GB of activations for GPT-2 1.5B). Existing activation-saving methods apply one uniform, data-agnostic strategy to every training instance, which the authors argue makes fine-tuning both ineffective and unstable. TokenSeek addresses this by judging token importance separately for each instance and then dropping the gradients, and therefore the cached activations, of the tokens judged least useful.

Key Contributions

1. Instance-aware token seeking. A scoring function that combines context information (column-wise accumulated attention weights) with gradient information (gradient magnitude of the loss with respect to penultimate-layer activations), merged as I(t_j) = α·log[I₁(t_j)] + β·Norm[I₂(t_j)].
2. Efficient token ditching. A mechanism that backpropagates only through selected tokens, so only the selected activations a_t^(l) need to be cached instead of the full activation a^(l).
3. Demonstrated memory savings with on-par or better accuracy. The paper reports 65.7% maximum memory reduction on Llama3.2 1B at 10% tokens while scoring 41.13 versus 40.82 for full-token tuning, and 14.8% memory consumption under QLoRA with 52.61 versus 40.82.
4. Interpretability analysis. A case study showing that context scores concentrate on earlier positions (long-tail distribution driven by the causal mask and the attention sink) while gradient scores concentrate on later "Response" positions, and that the two signals are complementary.

Main Findings

• Activations dominate training memory. The paper reports activations at 87% of memory consumption in Llama3 8B and 60GB of activations for GPT-2 1.5B, motivating activation-targeted optimization.
• Peak memory reduction on Llama3.2 3B. Peak memory is reduced by 60.7% with TokenSeek alone, and down to 11.1% when combined with QLoRA, which the authors say enables training on a single A100 GPU without OOM. Average memory usage drops by 26.9% and 86.7% in those two settings.
• Llama3.2 1B improves in every setting. TokenSeek alone scores 41.13, with LoHa 52.58, and with QLoRA 52.61, versus 40.82 for full-token tuning. The abstract reports 2.8 GB, 14.8% of original memory on Llama3.2 1B.
• Qwen2.5 0.5B with QLoRA marginally beats the baseline. 48.45 versus 48.43 average score at 19.2% average and 13.4% maximum memory.
• Random token selection is clearly worse. TokenTune (random) underperforms TokenSeek in the reported settings, for example 49.22 versus 52.61 on Llama3.2 1B with QLoRA, and 41.06 versus 42.28 on Qwen2.5 0.5B with LoHa.
• Stability improves. Comparing runs with tunable-token ratios from 10% to 50% on Llama3.2 1B with QLoRA, TokenSeek shows a narrower variance band than the random baseline.
• Memory scales with token ratio. On Llama3.2 1B with QLoRA, average memory falls from 32.6% at 50% tokens to 14.8% at 10% tokens, while the average score stays between 51.80 and 52.75.
• GPU budget. Peak memory with QLoRA plus TokenSeek ranges from 14.2 GiB down to 5.5 GiB, compared to 38.8 GiB for full token tuning, against reference GPU capacities of 40GB on A100 and 24GB on RTX 4090.
• Weights for the two signals matter but are not fragile. In the ablation on Qwen2.5 0.5B with QLoRA, α=5, β=5 gives 48.65 average, α=7, β=3 gives 48.53, α=1, β=0 gives 48.45, α=3, β=7 gives 48.28, and α=0, β=1 gives 46.39.
• PEFT pairs better with token ditching than full tuning. The paper observes full parameter tuning reaches lower training loss (possible overfitting), while QLoRA reports 47.66 versus 48.85 for QLoRA plus TokenSeek.
• A scale-dependent caveat. TokenSeek shows performance degradation on Qwen under plain (non-PEFT) settings while Llama does not, which the authors attribute to the limited representational capacity of smaller models.
• Interpretability patterns. Context scores prefer earlier tokens, keeping semantically meaningful content such as tokens related to mathematical learning while filtering definite articles and prepositions; gradient scores emphasize later "Response" positions.

Methodology in Plain English

During backpropagation, a Transformer must keep the intermediate activations from the forward pass so it can compute gradients later, and these cached activations are the largest memory consumer in fine-tuning. TokenSeek asks a simple question per training instance: which tokens actually matter for this particular example?

It answers using two signals. The first is attention. Because attention columns show how much every token attends to a given token, summing down a column gives a context-importance score; the paper applies a log transform because the distribution has a long tail. The second is gradients. The authors take the gradient of the loss with respect to the activations of the penultimate layer, sum the magnitudes across the hidden dimension, and min-max normalize the result. The two scores are combined with weights α and β into one importance score per token.

With scores in hand, the model backpropagates only through the top-scoring tokens (10% by default) and treats the rest as if their gradient contribution were zero. Because the unselected activations no longer need to be multiplied by anything in the backward pass, they do not need to be cached at all. The authors note that scoring requires only a forward pass and a partial backward pass, achieved by freezing all layers except the output head and the final decoder block, and they estimate that tuning 10% of tokens theoretically requires only about 1% of the activation memory. The whole scheme is architecture-agnostic, so it can be layered on top of existing PEFT methods such as LoHa and QLoRA.

Experiments fine-tune Qwen2.5 0.5B, Llama3.2 1B, and Llama3.2 3B on the Open-Platypus dataset using 21

Authors’ 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/