Jianxin Zhang

Author directory

2026

Implicit In-Context Learning compresses demonstration examples into a single context vector and injects it into the model’s activation space, achieving few-shot-level performance at near zero-shot inference cost. However, existing approaches typically aggregate demonstrations from all classes into a shared, task-level context vector, capturing global task information but without explicitly preserving fine-grained, class-conditional semantic distinctions. In this work, we propose Class-Conditional Context Vectors (C3V), a simple yet effective extension to implicit in-context learning that explicitly models class-specific contextual information by constructing separate context vectors for each class, without relying on explicit prompts. These class-conditional context vectors are additively injected into the model’s residual streams in a single forward pass, enabling contextual contributions to be modulated in a class-aware manner while keeping the backbone frozen. We evaluate C3V on multiple text classification benchmarks across several families of large language models. Experimental results demonstrate that C3V consistently outperforms task-level context vector baselines, and achieves higher average accuracy than conventional few-shot learning on most models.
Large language models (LLMs) have demonstrated strong capabilities in processing long contexts, enabling them to tackle tasks involving long textual inputs such as multi-turn conversations, legal documents, or retrieved documents in Retrieval-Augmented Generation (RAG) systems. However, despite their ability to handle long sequences, the resulting decoding latency and memory overhead remain substantial, posing challenges for real-world deployment. Recent advances in KV Cache reuse have shown potential to mitigate these costs, but these methods still suffer from notable performance degradation. To address this issue, we conduct an in-depth investigation of recomputation-based reuse methods and observe that the recomputed tokens often fail to align with the context segments most relevant to the question. This misalignment hinders proper updates to the critical contextual representations. Therefore, we propose the Attention-Aware Accurate KV Cache Fusion algorithm (A3), which precomputes and selectively fuses the KV Cache of text chunks based on their relevance to the question, achieving accurate integration with minimal computational overhead. Extensive experiments on various benchmarks and LLMs demonstrate that A3 achieves competitive task performance compared to five baselines while reducing the time-to-first-token (TTFT) by 2×.

2025

Large language models (LLMs) have achieved remarkable performance across various reasoning tasks. However, many LLMs still encounter challenges in reasoning, especially for LLMs with fewer parameters or insufficient pre-training data. Through our experiments, we identify that noise accumulation across layers often leads to unstable token predictions during reasoning. We find that contrasting the probability distributions across layers effectively mitigates this interference. Building on this insight, we propose Adaptive Layer-Wise contrastive decoding (ALW), a novel framework that enhances reasoning ability by dynamically disentangling noise in shallow layers from critical signals in deep layers. Extensive experiments on several reasoning benchmarks demonstrate that ALW consistently improves answer accuracy across multiple LLMs while maintaining inference efficiency. For example, we achieve a 48% improvement on the Gsm8k using the LLaMA-7B model and an absolute accuracy increase of 5.2 points on the BBH evaluation benchmark with the LLaMA-65B model.