PaperlayerView exact revision on arXiv

Source-linked AI summary

ULME: Predictive Coupling for Mutually Reinforcing Embedding and Language Modeling

Yang Liu, Zhiqiang Yuan, Shu-Dong Huang, Jiancheng Lv, Liwei Yang, Qiwei Yan, Kailin Lv, Jinchao Zhang

arXiv:2610.30705v1cs.CLcs.AI

TL;DR

Embedding and generation are often handled by separate models, while existing unified models may not let the two objectives shape each other. ULME couples them in one causal decoder using learnable slots and answer-side predictive supervision. Across two model families and eight scales, it improves MTEB over the strongest same-backbone baseline while remaining competitive on generation benchmarks.

  • Problem

    Embedding and generation are commonly handled by separately trained models, and existing unified models may use different computations or parallel objectives that do not directly supervise one another.

    Source: Introduction · pp. 1–2

  • Method

    ULME trains a causal decoder with learnable slots, using contrastive embedding and language-modeling losses plus stop-gradient alignment from a pooled answer summary to a query-side slot.

    Source: Introduction · p. 2

  • Results

    3.3–4.1 points on MTEB over the strongest same-backbone baseline, with ULME competitive on embedding and generation across Qwen2.5 and Qwen3 at eight scales.

    Source: Introduction · p. 2

  • Takeaways & Limitations

    A single ULME checkpoint beats specialized embedding baselines on MTEB and BRIGHT while remaining on par with a generation-only fine-tune on standard accuracy benchmarks.

    Source: Conclusion · p. 8

    ULME training partially degrades chat-style multi-turn dialogue, strict output-format compliance, and some instruction-following behaviors relative to the untouched Instruct or non-thinking checkpoint.

    Source: F. Limitations and Future Work · p. 13

Abstract

from arXiv · show

Text embedding and text generation are two of the most useful capabilities of modern large language models, yet they are typically provided by separate specialized models (an embedder and a generator) trained with incompatible objectives and deployed side by side in retrieval-augmented generation and agent systems. Recent unified embedding-generation models such as GritLM share parameters between the two capabilities, but they either switch attention regimes between modes or optimize embedding and language modeling as parallel objectives that never supervise each other. We argue that these designs are only partially unified: the two capabilities coexist on the same weights, but do not shape each other. We propose ULME, a post-training framework that unifies embedding and language modeling on a causal decoder-only LLM through predictive coupling. ULME inserts a chain of learnable slots between the query and the answer, reads a semantic embedding and a reasoning embedding from two designated slots, and trains three objectives on a shared causal backbone: a contrastive loss anchoring the semantic slot to document representations, a language-modeling loss supervising answer generation, and a stop-gradient predictive loss that aligns the reasoning slot with a pooled summary of the answer, so that generation actively shapes what the embedding encodes. Across two model families (Qwen2.5, Qwen3) and eight scales spanning 0.5B-8B parameters, a single ULME checkpoint improves MTEB by 3.3-4.1 points over the strongest open-source baseline, tightens BRIGHT and AdvBench-IR simultaneously, and remains on par with a generation-only fine-tune on standard accuracy benchmarks (MMLU, GSM8K, MMLU-Pro, HumanEval), demonstrating that genuine unification yields mutually reinforcing rather than merely coexisting capabilities.

Introduction

ULME targets the separation between query retrieval and answer generation by coupling their representations within one causal decoder. Its learnable slots provide distinct readouts, while answer-side representations supervise a query-side latent.

Introduction

  • Retrieval and generation often process the same query in separate models, leaving their query representations disconnected.The paper frames this as treating the semantic understanding needed for retrieval as independent from that needed to generate an answer.Source: Introduction · p. 1
  • A query’s final hidden state can collide as a globally comparable retrieval embedding and a position-dependent predictor of the first answer token.The paper motivates distinct readouts that retain shared semantics while avoiding this representational collision.Source: Introduction · p. 1
  • Prior unified models may use different attention computations or optimize embedding and generation in parallel without direct supervision between them.ULME instead uses answer-side hidden states from the same forward pass as a stop-gradient target for an inference-time query-side latent.Source: Introduction · pp. 1–2
  • ULME uses a learnable slot chain in one causal forward pass: a semantic slot aligns to documents, generation is supervised over the answer, and a reasoning slot aligns to a pooled answer summary.The predictive alignment uses stop-gradient, and all three losses share the model parameters and hidden states.Source: Introduction · p. 2
  • 3.3–4.1 points on MTEB over the strongest same-backbone baseline, with ULME competitive on embedding and generation across two model families and eight scales.The reported families are Qwen2.5 and Qwen3, at scales from 0.5B to 8B.Source: Introduction · p. 2

Related Work

Prior work uses decoder-only LLMs both for generation and for dense embeddings, but embedding-only fine-tuning can give up generation, while existing unified models leave gaps in how the tasks interact. ULME instead couples them in one causal computation.

Large Language Models

  • Decoder-only LLMs use causal attention and next-token prediction, while retrieval-augmented generation and agents typically pair them with separately trained retrievers.These systems use retrieved context to support generation or agent tasks.Source: Related Work · p. 2

LLM-based Embedding Models

  • Contrastive fine-tuning of decoder-only LLMs can produce competitive embedders, but those models are no longer usable as generators.E5-Mistral and NV-Embed pool a hidden state as the embedding; other approaches modify context or remove the causal mask.Source: Related Work · p. 2

Generative Retrieval

  • Generative retrieval can retrieve by generating relevant dialogue fragments, yet practical retrieval stacks still keep an embedder beside a generator.The cited F2RVLM work also describes a two-stage system that indexes semantic fragments with a dedicated embedding model.Source: Related Work · pp. 2–3

Unified Embedding-Generation Models

  • Figure 1 depicts ULME’s shared slot chain with semantic query–document alignment, answer-token prediction, and stop-gradient query–answer alignment.The figure labels the total loss as the sum of the semantic, generation, and predictive losses.Source: Figure 1 · p. 3
  • Existing unified models share a backbone, but may switch attention regimes or train embedding and generation in parallel without either objective directly shaping the other.ULME’s stated departure is one causal computation with objectives coupled through a shared latent state.Source: Related Work · p. 3

Unified Language Modeling and Embedding Model Architecture

ULME uses a chain of learnable slots between the query and answer to expose separate semantic and reasoning readouts within one causal decoder pass. Retrieval, answer generation, and predictive coupling supervise these readouts and the answer.

Unified Language Modeling and Embedding Model Architecture

  • ULME inserts learnable slots between the query and answer, keeping both readouts and answer generation on one causal forward pass without changing masking or adding an auxiliary branch.Each slot attends to the query and earlier slots; answer tokens attend to all slots.Source: Model Architecture · p. 3
  • Slot 0 is the semantic embedding, slot N−1 the reasoning embedding, and answer-span hidden states serve as both generation target and teacher for the latter.The document is encoded analogously with a slot for semantic retrieval.Source: Model Architecture · p. 3, Retrieval and Generation Objectives · p. 4, Predictive Coupling · p. 4
  • A contrastive loss supervises the semantic slot, while next-token cross-entropy supervises answer generation over the answer span.The language-modeling loss reaches slot states through the answer's attention to the slot chain.Source: Retrieval and Generation Objectives · pp. 3–4
  • Predictive coupling directly aligns the terminal slot with a pooled summary of the answer, adding sequence-level supervision beyond language modeling's indirect token-level signal.The answer summary is pooled from answer-side hidden states; the predictive loss uses one-sided in-batch InfoNCE.Source: Predictive Coupling · p. 4
  • Stop-gradient keeps predictive learning one-sided: the less-informed query-side slot moves toward the answer-side representation without sending predictive gradients back into it.The answer-side representation is separately supervised by language modeling.Source: Predictive Coupling · p. 4
  • Together, language modeling populates the slot chain indirectly through answer-token dependencies, predictive coupling directly supervises its terminal slot, and semantic retrieval anchors slot 0 to document space.The full objective combines the semantic, generative, and predictive losses with weights.Source: Predictive Coupling · p. 4

Experiments

ULME is evaluated as a single checkpoint for embedding and generation across two model families and eight scales. It improves retrieval benchmarks over embedding baselines while remaining competitive with generation-only fine-tuning, and ablations examine how its slots and losses contribute.

Evaluation setup

  • Each configuration is evaluated from one ULME checkpoint without isolated head fine-tuning or capability toggles, across Qwen2.5 and Qwen3 models spanning 0.5B–8B parameters.Default settings use N=4 slots and train on E5 contrastive data for embedding and OpenThoughts-114k for generation.Source: Experiments · p. 4, Main Results · p. 5

MTEB

  • 66.0 on Qwen3-8B MTEB(eng, v2) is 9.2 points above LLM2Vec and 4.1 above LLM2VEC-GEN on the same backbone.Against LLM2VEC-GEN, gains persist across Qwen3-{1.7B, 4B, 8B}: +3.3, +3.5, and +4.1 points.Source: Table 1 · p. 5, Main Results · p. 5

Generation

  • ULME matches or exceeds Gen-only FT on all but one knowledge and math cell, and leads it at every scale on HumanEval and HumanEval+.At Qwen3-8B, HumanEval+ is 67.7 vs 62.8; both ULME and Gen-only FT occasionally trail the backbone, a pattern attributed to continued-training distribution shift.Source: Table 2 · p. 6, Main Results · p. 6

BRIGHT and AdvBench-IR

  • On BRIGHT, ULME gains +0.7 / +0.8 / +1.2 nDCG@10 points over LLM2VEC-GEN across Qwen3-{1.7B, 4B, 8B}, while AdvBench-IR also improves.Against LLM2Vec, AdvBench-IR falls by 25.3% on Qwen3-1.7B and 20.3% on Qwen3-8B; Table 3 reports both benchmarks.Source: Main Results · p. 5, Table 3 · p. 6

Slot design

  • Removing the slot chain drops MTEB-lite by 5.2 points and MMLU by 2.0, while both axes peak at N=4 before eroding slightly at N=8.The ablation uses Qwen2.5-7B-Instruct and evaluates embedding on 10 MTEB tasks and generation on MMLU and GSM8K.Source: Table 4 · p. 6, Ablation Studies · pp. 6–7

Loss decomposition

  • Removing Lpred drops MTEB-lite by 3.6 points while leaving generation flat; removing Lgen costs embedding 2.8 points despite the embedding loss remaining intact.Removing Lsem drops MTEB-lite from 63.8 to 31.2 while MMLU 76.6 and GSM8K 90.1 remain essentially unchanged.Source: Table 5 · p. 6, Ablation Studies · p. 7

Predictive target

  • Mean answer-span pooling beats last-token and length-weighted predictive targets on every embedding benchmark, with a 1.9-point MTEB-lite gap; generation is essentially unchanged.The authors attribute the last-token target's weaker embedding results to its exposure to next-token prediction pressure.Source: Table 6 · p. 7, Ablation Studies · p. 7

Slot specialization

  • Across 41 MTEB tasks, slot 0 leads on lexical-overlap retrieval, while slot N−1 leads on tasks requiring reasoning about document claims, summaries, or expression.The reported examples include slot 0 gains of 4.84 on SprintDup and slot N−1 gains of 3.60 on ClimateFEVER.Source: Ablation Studies · p. 7, Figure 2 · p. 7

Conclusion

ULME couples embedding and generation in a single decoder-only model, improving retrieval performance across Qwen2.5 and Qwen3 scales while remaining on par with generation-only fine-tuning on standard accuracy benchmarks.

Conclusion

  • Across Qwen2.5 and Qwen3 at eight scales from 0.5B to 8B, ULME beats specialized embedding baselines on MTEB and BRIGHT while remaining on par with generation-only fine-tuning on standard accuracy benchmarks.The benchmarks include MMLU, GSM8K, MMLU-Pro, and HumanEval.Source: Conclusion · p. 8
  • ULME uses one causal forward pass and learnable slots, with contrastive, language-modeling, and stop-gradient predictive losses coupling embedding and generation.The predictive loss aligns the query-side slot with a pooled summary of the answer.Source: Conclusion · p. 8
  • The results show that the coupled objective is not a compromise between embedding and generation capabilities.The paper characterizes the coupling as allowing one capability to benefit the other.Source: Conclusion · p. 8

A. Additional Implementation and Evaluation Details

The appendix describes the model configurations and data, evaluation protocols, and baseline sources used for the reported experiments.

Training details

  • All ULME and Gen-only FT runs use eight NVIDIA H20 96 GB GPUs, and Table 7 lists training hyperparameters across the eight configurations.Only the backbone identifier differs between configurations.Source: A. Additional Implementation and… · p. 10

Backbones and training data

  • ULME uses officially post-trained Qwen2.5 Instruct and Qwen3 non-thinking backbones, E5-format contrastive data, and OpenThoughts-114k prompt-and-answer pairs.The embedding data includes one mined hard negative per query; the generation data supplies reasoning-heavy prompts and reference chains of thought.Source: A. Additional Implementation and… · p. 10

Evaluation protocols

  • Embedding evaluation pools the semantic slot in one causal pass; AdvBench-IR uses 520 harmful instructions and reports hit rate, where lower is better.Generation evaluations use the same checkpoint without a post-hoc head swap or mode toggle.Source: A. Additional Implementation and… · p. 10

Baselines

  • Baseline values come from LLM2VEC-GEN under a uniform protocol, and unreported backbone-baseline combinations are omitted.The authors attribute observed ULME gaps to the training objective rather than protocol drift or independent reproduction.Source: A. Additional Implementation and… · p. 10

B. Full Results across All Eight Backbones

Across all eight backbones, ULME improves embedding results over LLM2Vec and is at least on par with Gen-only FT on knowledge and math, while outperforming it on code.

Full results

Training and baseline setup

  • Table 7's hyperparameters are shared across all eight ULME configurations except for the backbone identifier.Gen-only FT uses the same backbone, OpenThoughts-114k data, and optimizer schedule as ULME, differing only in its objective.Source: Table 7 · p. 10, B. Full Results across All Eight… · p. 11

C. MTEB Evaluation Prompts and Analysis

MTEB embeddings use short task-specific instructions and semantic-slot pooling, with instruction routing varying by task type. The analysis links these instructions to task-aware similarity and the division of labor between semantic and reasoning slots.

Prompt design

  • At inference, ULME prepends a short task-specific instruction and pools the semantic slot from a single causal forward pass.The instruction describes the input type and target property, such as topic clustering or retrieving refuting arguments.Source: C. MTEB Evaluation Prompts and Analysis · p. 11
  • Prompts are single instruction lines followed immediately by raw input, typically 8–12 words and without role markers, inline examples, or chat templates.The trailing colon-space precedes the raw input.Source: C. MTEB Evaluation Prompts and Analysis · p. 11

Symmetric vs. asymmetric routing

  • MTEB routes symmetric matching with shared instructions, single-side embedding with a query-only instruction, and asymmetric retrieval with distinct query and document instructions.For asymmetric retrieval, each instruction names its side's input type and target property.Source: C. MTEB Evaluation Prompts and Analysis · pp. 11–12

Effect on embedding quality

  • Task-aware instructions specify which similarity axis to follow when surface overlap and semantic relevance diverge, while matching the E5-style instruction format used during training.The passage connects this task-aware prompting to the semantic and reasoning slots' division of labor and says no test-time prompt engineering is applied.Source: C. MTEB Evaluation Prompts and Analysis · p. 12

F. Limitations and Future Work

ULME's generation-side training is limited to a small public reasoning corpus, and the paper notes partial degradation in some instruction and dialogue behaviors. Future work includes scaling the experiments and extending the framework to multimodal backbones.

Limitations

  • Generation-side training uses a small public reasoning corpus, and preliminary checks find partial degradation in chat, output-format compliance, and some instruction-following behaviors.The training uses OpenThoughts-114k in particular and omits the proprietary conversational, safety, and format-following data used to align the original backbones.Source: F. Limitations and Future Work · pp. 12–13

Future work

  • The main experiments cap at 8B parameters and the OpenThoughts-114k generation corpus, leaving the scaling behavior of predictive-coupling gains an open question.The paper asks whether gains grow, saturate, or invert with larger backbones and more diverse, larger-scale generation data.Source: F. Limitations and Future Work · p. 13
  • The authors propose extending the token-level slot-chain approach to vision-language and audio-language backbones for retrieval and generation across modalities.They point to visual RAG and agentic visual evidence management as settings where retrieval and reasoning are interleaved.Source: F. Limitations and Future Work · p. 13

Related papers

Loading 2610.30705v1…