InstEmb: Instruction-Following Embeddings that Anticipate Future Model Outputs

Problem Motivation

In instruction‑following retrieval, a query such as “What features does this tent have for outdoor use?” is semantically close to candidate items, but the user cares about specific dimensions (e.g., durability vs. portability). Existing LLM embedding pipelines use only the hidden state of the last input token, which captures input‑intrinsic semantics but misses the semantics that will appear in the model’s future output.

InstEmb Overview

InstEmb appends a set of learnable look‑ahead tokens after the input sequence. During training, a frozen teacher receives the true output while the student receives the look‑ahead tokens. The student’s look‑ahead token states are aligned with the teacher’s output‑conditioned hidden states, providing output‑aware semantics without requiring decoding at inference time. In inference the model performs a single prefilling pass and produces an embedding that fuses input semantics with anticipated output semantics.

Student input: [instruction + input + <eos> + look-ahead tokens]
Teacher input: [instruction + input + <eos> + truncated gold output]

Representation Self‑Distillation

Two distillation objectives are used. The first aligns hidden states with a mean‑squared‑error (MSE) loss, directly regressing the student’s look‑ahead token representations to the teacher’s output hidden states. The second aligns output probability distributions with a KL‑divergence loss applied to the language‑model head. Experiments show MSE excels on fine‑grained instruction tasks, while KL is more robust on generic embedding benchmarks.

Multi‑View Contrastive Learning

To preserve stable input‑intrinsic semantics, InstEmb applies multi‑view contrastive learning on the final input token. Four views are constructed for each sample: (1) the student’s first encoding of the input, (2) the student’s second encoding with a different dropout mask (SimCSE‑style), (3) the frozen teacher’s encoding of the input, and (4) the student’s encoding of the target output sequence. These views form a positive set, while other samples serve as negatives. An InfoNCE loss encourages the model to keep input semantics stable while integrating teacher and output views.

Dual‑Anchor Alignment Pooling (DAAP)

Pooling combines two anchors: Input‑Intrinsic Anchor (hidden state of the last input token) and Output‑Aware Anchor (average hidden state of the look‑ahead tokens). The final embedding is the average of these two anchors, ensuring that both training objectives are reflected during inference.

Experimental Setup

InstEmb is built on LLaMA‑3‑8B‑Instruct and trained on ~200 k abstractive QA examples from 11 datasets for one epoch (Adam, lr=5e‑6). During inference, eight look‑ahead tokens are used.

Evaluation covers three task families:

Instruction‑following retrieval (FollowIR, InfoSearch)

Instruction‑following embedding (Inst.STSb, IntentEmotion, NYTCluster)

General sentence‑embedding benchmarks (AskUbuntu, 20News, SciDocsRR, StackOverflowDup)

Results

InstEmb‑MSE achieves an average score of 28.5 and a +15.6 p‑MRR improvement on FollowIR, surpassing strong baselines such as FollowIR‑7B and PromptRetriever, and also leads on InfoSearch. On instruction‑following embedding tasks, InstEmb‑MSE‑DAAP reaches 67.08 average, well above the Inbedder reimplementation (59.90). On generic benchmarks, InstEmb‑KL‑DAAP attains 63.39, demonstrating that the method improves instruction capability without sacrificing general semantic quality.

Ablation Studies

Removing any of the four contrastive views (student first encoding, student second encoding, teacher encoding, or student output encoding) causes notable drops, confirming the importance of SimCSE‑style augmentation and teacher guidance. Varying look‑ahead token length shows that even a single token yields substantial gains, though benefits plateau for tasks with short outputs. Pooling ablations reveal that DAAP outperforms simple last‑token or mean‑pooling, highlighting the need to treat the two anchors distinctly.

Interpretability Analysis

Attention visualizations show that the original LLaMA‑3 model concentrates attention at the sequence start, whereas InstEmb shifts focus to the system‑prompt and instruction endings. Hidden‑state cosine similarity analysis demonstrates low similarity between the last input token and later positions (preserving input semantics) and high similarity between look‑ahead tokens and gold output tokens (validating output‑aware learning).

Conclusion

InstEmb provides a practical pathway for instruction‑following embeddings by jointly preserving input‑intrinsic semantics and injecting output‑aware semantics via learnable look‑ahead tokens, self‑distillation, multi‑view contrastive learning, and DAAP. Crucially, it achieves these gains without any decode‑then‑encode step at inference, offering both efficiency and performance improvements for downstream retrieval and embedding tasks.