arXiv subjects

Zhongqiao Lin

Publications and source records attributed to Zhongqiao Lin.

3 recordsLinked to original sources

arXiv2025

SimLens for Early Exit in Large Language Models: Eliciting Accurate Latent Predictions with One More Token

Intermediate-layer predictions in large language models (LLMs) are informative but hard to decode accurately, especially at early layers. Existing lens-style methods typically rely on direct linear readout, which is simple but often drifts away from the model's eventual prediction. We proposeSimLens, a simple training-free decoder for single-token decision tasks that keeps only the start token and a candidate answer token ([s] and [a]) and performs one lightweight continuation through the remaining upper layers. This surprisingly small modification recovers much more accurate latent predictions than direct linear decoding. We further introduce Linear SimLens, a lightweight linear approximation for entropy-based confidence estimation, and combine the two in SimExit, a hybrid early-exit mechanism. On ARC, BoolQ, and HeadQA with LLaMA-7B and Vicuna-7B, SimLens improves Iso-Compute accuracy in all six settings, with an average gain of +0.43 even when fair compute includes the extra two-token post-forward overhead. SimExit yields an average 1.15$\times$ speedup at the best-accuracy operating points and 1.40$\times$ when allowing up to a 1 percentage-point accuracy drop. Ablations show that [s] and [a] play distinct roles as global condition and semantic anchor, respectively.

cs.CL↗

arXiv2024

Label Words as Local Task Vectors in In-Context Learning

Large Language Models (LLMs) have demonstrated remarkable abilities, one of the most important being in-context learning (ICL). With ICL, LLMs can derive the underlying rule from a few demonstrations and provide answers that comply with the rule. Previous work hypothesized that the network creates a task vector in specific positions during ICL. The task vector can be computed by averaging across the dataset. It conveys the overall task information and can thus be considered global. Patching the global task vector allows LLMs to achieve zero-shot performance with dummy inputs comparable to few-shot learning. However, we find that such a global task vector does not exist in all tasks, especially in tasks that rely on rules that can only be inferred from multiple demonstrations, such as categorization tasks. Instead, the information provided by each demonstration is first transmitted to its answer position and forms a local task vector associated with the demonstration. In some tasks but not in categorization tasks, all demonstrations' local task vectors converge in later layers, forming the global task vector. We further show that local task vectors encode a high-level abstraction of rules extracted from the demonstrations. Our study provides novel insights into the mechanism underlying ICL in LLMs, demonstrating how ICL may be achieved through an information aggregation mechanism.

cs.CL↗

arXiv2024

From Transformer to Biology: A Hierarchical Model for Attention in Complex Problem-Solving

Attention is fundamental to cognition, yet it remains a challenge to understand attention in tasks approaching real-world complexity. Here, we approached this problem by modeling gaze patterns of monkeys playing Pac-Man. We first show a transformer network trained to reproduce their gameplay developed internal attention patterns closely matching the monkeys' eye movements. By dissecting the network's attention, we revealed a hierarchical structure comprising two components: a value-based layer encoding fixed object salience, coupled with a dynamic interaction layer tracking relational information between game elements. We further developed a condensed model in which reward-driven attention serves as a gain modulator and is integrated with spatial attention maps, predicting attention as well as the transformer. Together, our study pioneers the use of AI architectures as analytical tools and bridges mechanistic interpretability with cognitive neuroscience to yield novel, testable insights into how the brain coordinates reward, spatial cognition, and attention in complex environments.

q-bio.NC↗