Unlike conventional dense retrieval, which mainly captures semantic similarity, InsightEmb learns progress-oriented action-intent matching. The model is trained on mathematical reasoning data and transfers to interactive agent environments without environment-specific retriever fine-tuning.
For the complete training pipeline, insight-generation scripts, and end-to-end agent evaluation, see the
official GitHub repository.
-
Progress-Oriented Retrieval
Retrieves insights according to whether they can resolve the current bottleneck and advance the agent toward its goal.
-
Action-Intent Matching
Connects concrete agent states with abstract, procedurally useful rules rather than relying only on topical overlap.
-
Math-Only Retriever Training
Learns transferable retrieval geometry from mathematical problems and reasoning trajectories, without target-environment retriever optimization.
-
Two-Stage Contrastive Curriculum
- Situation-to-Insight Matching: aligns problems and intermediate reasoning states with abstract heuristic rules.
- Situation-to-Experience Matching: groups different-looking situations that require the same underlying strategy.
-
Partial-Trajectory Supervision
Uses truncated reasoning traces to train retrieval from intermediate states, improving sensitivity to evolving procedural bottlenecks.
-
Cross-Domain Transfer
Evaluated on ALFWorld, WebShop, ScienceWorld, and SRA-Bench, with consistent gains over the base embedder and strong reasoning-oriented retrievers.
1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4
5MODEL_ID = "YOUR_HF_NAMESPACE/InsightEmb"
6DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
7
8# InsightEmb is initialized from Qwen3-Embedding-4B.
9tokenizer = AutoTokenizer.from_pretrained(
10 "Qwen/Qwen3-Embedding-4B",
11 trust_remote_code=True,
12 padding_side="left",
13)
14
15model = AutoModel.from_pretrained(
16 MODEL_ID,
17 trust_remote_code=True,
18 torch_dtype=torch.bfloat16,
19 attn_implementation="flash_attention_2",
20).to(DEVICE).eval()
InsightEmb uses last-token pooling followed by L2 normalization.
1def last_token_pool(last_hidden_states, attention_mask):
2 left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
3 if left_padding:
4 return last_hidden_states[:, -1]
5
6 sequence_lengths = attention_mask.sum(dim=1) - 1
7 batch_size = last_hidden_states.shape[0]
8 return last_hidden_states[
9 torch.arange(batch_size, device=last_hidden_states.device),
10 sequence_lengths,
11 ]
12
13
14@torch.inference_mode()
15def encode(texts, max_length=8192):
16 inputs = tokenizer(
17 texts,
18 padding=True,
19 truncation=True,
20 max_length=max_length,
21 return_tensors="pt",
22 ).to(DEVICE)
23
24 outputs = model(**inputs)
25 embeddings = last_token_pool(
26 outputs.last_hidden_state,
27 inputs["attention_mask"],
28 )
29 return F.normalize(embeddings, p=2, dim=1)
Queries should include a retrieval instruction. Candidate insights are encoded directly.
1retrieval_instruction = (
2 "Given the current agent state, retrieve insights that resolve the "
3 "current bottleneck and help the agent make progress toward its goal"
4)
5
6agent_state = """
7Goal: buy a light-grey dining set under $250.
8History: searched twice and browsed several result pages.
9Observation: many partially matching products are visible, but none selected.
10"""
11
12query = f"Instruct: {retrieval_instruction}\nQuery:{agent_state}"
13
14insights = [
15 "Keep refining the query until an exact product appears.",
16 "Include all critical attributes in the search, then select the required variant before purchasing.",
17 "Browse more result pages before committing to a product.",
18]
19
20query_embedding = encode([query])
21insight_embeddings = encode(insights)
22
23scores = (query_embedding @ insight_embeddings.T).squeeze(0)
24top_indices = torch.topk(scores, k=2).indices.tolist()
25
26for rank, index in enumerate(top_indices, start=1):
27 print(f"Insight {rank}: {insights[index]}")
For dynamic retrieval after every environment step, see the
ALFWorld, WebShop, and ScienceWorld implementations.
It can be used as a bi-encoder retriever or as the dense component of a hybrid retrieval pipeline.
1@misc{chung2026insightemblearningactionintentembeddings,
2 title={InsightEmb: Learning Action-Intent Embeddings for Agentic Insight Retrieval},
3 author={Tsz Ting Chung and Jiangnan Li and Jie Zhou and Mo Yu},
4 year={2026},
5 eprint={2608.04761},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2608.04761},
9}
