Zhijie Nie

Author directory

2026

Retrieval-Augmented Generation (RAG) can fail under knowledge conflicts, where a model’s parametric prior competes with retrieved evidence and leads to unfaithful outputs. Existing mitigation methods typically operate on the final decoding distribution or rely on weight editing with access to ground-truth targets. Through layer-wise analysis, we identify a **parametric suppression** pattern: in many conflict cases, evidence-supported answers become salient in intermediate representations but are weakened by later layer outputs. To address this, we propose CoRect (**Co**ntext-Aware Logit Contrast for Hidden State **Rect**ification), a training-free inference-time method. By contrasting contextualized and context-free forward passes, CoRect first selects an evidence-supported token direction without using gold labels. It then rectifies hidden states in layers whose outputs suppress this direction. Across Question Answering (QA) and summarization benchmarks, CoRect improves average faithfulness and reduces hallucinations compared with strong decoding baselines.
Soft compression stores a text span as k continuous compressed vectors for a frozen decoder, but existing compressors are usually tied to the decoder used during training. We study whether those vectors can be retargeted to a different frozen decoder through its input-embedding interface under a controlled reconstruction protocol. Our benchmark covers seven target decoders, three vector budgets, and directed source–target pairs, comparing Random, shared-embedding least squares (LS), Converter-only, Enc-Conv, and target-native Origin. Converter-only maps vectors from a frozen source compressor, whereas Enc-Conv also adapts the source-initialized vector generator. Transfer is feasible but decoder-conditioned: trainable retargeting can approach Origin BLEU on decoder-compatible targets, while LS is often close to Random despite using shared token anchors. Enc-Conv consistently matches or outperforms Converter-only, with its largest gains on difficult transfers. Selected-pair multi-seed runs and a small SQuAD probe provide complementary evidence, although downstream utility remains target-dependent. Perturbation and geometry probes are consistent with decoder-specific readable regions. Finally, pair-level readability scores provide weak-to-moderate pre-transfer screening signals.
Sparse retrieval underpins modern search systems, from web search to retrieval-augmented generation. Existing work has introduced Learned Sparse Retrieval (LSR) to push beyond exact lexical matching toward richer semantics. Yet LSR has so far remained tied to encoder-style bidirectional architectures, and its extension to multimodal settings still relies heavily on auxiliary cross-modal modules. To address these limitations, we introduce UEmbed (Unified Embedding), a decoder-only multimodal embedding model that produces both sparse lexical and dense representations in one causal forward pass. UEmbed appends N learnable special tokens to the input and partitions the vocabulary into N disjoint subsets. Each token’s causal hidden state predicts sparse weights over its assigned subset, and the N subsets are concatenated into the full sparse vector. Trained on public data, we release UEmbed at 2B, 4B, and 9B scales. UEmbed-9B reaches 71.8 (dense) and 71.0 (sparse) on MMEB-v2, outperforming multimodal embedding models trained on publicly available data (e.g., RzenEmbed). On BEIR, UEmbed also remains competitive with strong dense and sparse baselines. Furthermore, we demonstrate the practical utility of UEmbed across three dimensions: effectiveness, efficiency, and agentic applications. Overall, UEmbed offers a new paradigm: it unifies dense and sparse embeddings in one model, while further extending sparse retrieval to unify text and multimodal inputs.
Legal judgments may inadvertently contain errors, yet human appellate review faces immense pressure to improve efficiency. While current legal AI focuses on legal judgment prediction or document generation, judgment review is fundamentally a retrospective diagnostic task. In this paper, we introduce Appellate Review to assess large language models’ (LLMs) legal diagnostic reasoning underlying appellate review. We specifically target Chinese criminal cases, under the assumption that their highly structured logic, which maps explicit facts to strict statutory penalties, provides a verifiable ground truth. We decompose Appellate Review into three subtasks: error detection, error classification, and error correction. To support this, we construct AR-Bench, a controlled benchmark comprising 8,700 judgments with counterfactually introduced errors spanning charges, prison terms, and fines. Evaluations of 16 LLM variants reveal substantial limitations: models classify known errors relatively well but struggle with unguided detection and precise correction. Further analyses of a representative model show that it identifies explicit textual contradictions but fundamentally fails at holistic deductive reasoning and rule-based arithmetic.
Dense retrieval represents queries and documents as high-dimensional embeddings, but these representations can be redundant at the query level: for a given information need, only a subset of dimensions is consistently helpful for ranking. Prior work addresses this via pseudo-relevance feedback (PRF) based dimension importance estimation, which can produce query-aware masks without labeled data but often relies on noisy pseudo signals and heuristic test-time procedures. In contrast, supervised adapter methods leverage relevance labels to improve embedding quality, yet they learn global transformations shared across queries and do not explicitly model query-aware dimension importance. We propose a Query-Aware Adaptive Dimension Selection framework that learns to predict per-dimension importance directly from query embedding. We first construct oracle dimension importance distributions over embedding dimensions using supervised relevance labels, and then train a predictor to map a query embedding to these label-distilled importance scores. At inference, the predictor selects a query-aware subset of dimensions for similarity computation based solely on the query embedding, without pseudo-relevance feedback. Experiments across multiple dense retrievers and benchmarks show that our learned dimension selector improves retrieval effectiveness over the full-dimensional baseline as well as PRF-based masking and supervised adapter baselines.

2025

Multi-task prompt tuning utilizes multiple high-resource source tasks to improve performance on low-source target tasks. Existing approaches transfer the soft prompt trained by combining all source tasks or a single “high-similar” source task one-time-only. However, we find that the optimal transfer performance often comes from a combination of source tasks, which is neither one nor all. Further, we find that the similarity between source and target tasks also changes dynamically during fine-tuning after transfering, making similarity calculation in the initiation stage inadequate. To address these issues, we propose a method called Dynamic Task Vector Grouping (DTVG), whose core ideas contain (1) measuring the task similarity with task vectors instead of soft prompt, (2) grouping the optimal source task combination based on two metrics: target similarity and knowledge consistency; (3) dynamically updating the combination in each iteration step. Extensive experiments on the 26 NLP datasets under different settings demonstrate that DTVG effectively groups similar source tasks while reducing negative transfer, achieving the start-of-art performance.
Text embeddings from large language models (LLMs) have achieved excellent results in tasks such as information retrieval, semantic textual similarity, etc. In this work, we show an interesting finding: when feeding a text into the LLM-based embedder, the obtained text embedding will be able to be aligned with the key tokens in the input text. We first fully analyze this phenomenon on eight LLM-based embedders and show that this phenomenon is universal and is not affected by model architecture, training strategy, and embedding method. With a deeper analysis, we find that the main change in embedding space between these embedders and their LLM backbones is in the first principal component. By adjusting the first principal component, we can align text embedding with the key tokens. Finally, we give several examples to demonstrate the vast application potential of this finding: (1) we propose a simple and practical sparse retrieval method based on the aligned tokens, which can achieve 80% of the dense retrieval effect of the same model while reducing the computation significantly; (2) we show that our findings provide a novel perspective to help understand novel technologies (e.g., instruction-following embedding) and fuzzy concepts (e.g., semantic relatedness vs. similarity) in this field.

2024

Sentence Representation Learning (SRL) is a crucial task in Natural Language Processing (NLP), where contrastive Self-Supervised Learning (SSL) is currently a mainstream approach. However, the reasons behind its remarkable effectiveness remain unclear. Specifically, many studies have investigated the similarities between contrastive and non-contrastive SSL from a theoretical perspective. Such similarities can be verified in classification tasks, where the two approaches achieve comparable performance. But in ranking tasks (i.e., Semantic Textual Similarity (STS) in SRL), contrastive SSL significantly outperforms non-contrastive SSL. Therefore, two questions arise: First, what commonalities enable various contrastive losses to achieve superior performance in STS? Second, how can we make non-contrastive SSL also effective in STS? To address these questions, we start from the perspective of gradients and discover that four effective contrastive losses can be integrated into a unified paradigm, which depends on three components: the Gradient Dissipation, the Weight, and the Ratio. Then, we conduct an in-depth analysis of the roles these components play in optimization and experimentally demonstrate their significance for model performance. Finally, by adjusting these components, we enable non-contrastive SSL to achieve outstanding performance in STS.