Unsupervised Layer-Wise Dynamic Test Time Adaptation for LLMs

Test-time adaptation (TTA) for large language models (LLMs) updates model parameters at inference time using signals available at deployment. This paper focuses on a common yet under-explored regime: unsupervised, sample-specific TTA, where the model adapts independently for each prompt using only the prompt itself, without gold answers or external supervision. Although appealing, naive unsupervised TTA with a fixed, handcrafted learning rate can be unstable: updates may overfit to prompt-specific statistics, drift from the desired answer distribution, and ultimately degrade generation quality. This failure mode is not surprising, as in this case TTA must adapt to a single prompt within only a few gradient steps, unlike standard training that averages updates over large datasets and long optimization horizons. Therefore, we propose layer-wise dynamic test-time adaptation, a framework which explicitly modulates TTA strength as a function of prompt representation, LLM structure and adaptation step. In our setting, TTA updates only LoRA parameters, and a lightweight hypernetwork predicts per-layer, per-step learning-rate multipliers, enabling fine-grained control. Experiments across various datasets and LLMs consistently show that our method substantially strengthens TTA by learning effective scaling patterns over adaptation steps and transformer layer projections, improving stability while delivering better performance.

💡 Research Summary

The paper tackles the instability of unsupervised, sample‑specific test‑time adaptation (TTA) for large language models (LLMs). In this regime, each incoming prompt is processed independently: the model performs a few gradient steps using only the prompt text, without any gold answer, and then generates a response before resetting its parameters. A fixed global learning rate, as used in prior work, often leads to either negligible updates or catastrophic over‑fitting because gradient magnitudes vary dramatically across transformer layers and across adaptation steps.

To address this, the authors propose a two‑fold solution. First, they restrict the adaptation to Low‑Rank Adaptation (LoRA) parameters, updating only the query and value projection matrices of each transformer layer while keeping the backbone frozen. This drastically reduces the number of trainable parameters and the computational overhead, making per‑prompt adaptation feasible.

Second, they introduce a lightweight hypernetwork named SCALE‑NET that predicts per‑layer, per‑step learning‑rate multipliers. Given a prompt x, the current adaptation step k (out of a total of K steps), and the total number of steps K, SCALE‑NET outputs a vector s(k)∈ℝ^L_≥0 where L is the number of transformer layers. The actual update for LoRA parameters ϕ_ℓ at layer ℓ becomes:
ϕ_ℓ^{(k+1)} = ϕ_ℓ^{(k)} – η·s(k)ℓ·∇{ϕ_ℓ}