--- license: apache-2.0 language: - en - zh base_model: - Qwen/Qwen3-Embedding-4B pipeline_tag: feature-extraction library_name: transformers tags: - embedding - retriever - RAG - Agentic - Skills --- # InsightEmb (Learning Action-Intent Embeddings for Agentic Insight Retrieval) [![Paper](https://img.shields.io/badge/Paper-arXiv%3A2608.04761-red)](https://arxiv.org/pdf/2608.04761) [![Code](https://img.shields.io/badge/GitHub-InsightEmb-black?logo=github)](https://github.com/ttchungc/InsightEmb) [![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://ttchungc.github.io/projects/insightemb/) **InsightEmb** is an embedding model for **agentic insight retrieval**. It retrieves abstract rules and strategies that resolve an agent's current procedural bottleneck and enable the next useful action. 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](https://github.com/ttchungc/InsightEmb). --- ## ✨ Key Features - **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. --- ## 🚀 Quick Start ### Installation ```bash pip install torch transformers>=4.53.0 ``` ### 1) Load the Model Replace `YOUR_HF_NAMESPACE/InsightEmb` with the model repository ID. ```python import torch import torch.nn.functional as F from transformers import AutoModel, AutoTokenizer MODEL_ID = "YOUR_HF_NAMESPACE/InsightEmb" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # InsightEmb is initialized from Qwen3-Embedding-4B. tokenizer = AutoTokenizer.from_pretrained( "Qwen/Qwen3-Embedding-4B", trust_remote_code=True, padding_side="left", ) model = AutoModel.from_pretrained( MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", ).to(DEVICE).eval() ``` ### 2) Encode Text InsightEmb uses last-token pooling followed by L2 normalization. ```python def last_token_pool(last_hidden_states, attention_mask): left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] if left_padding: return last_hidden_states[:, -1] sequence_lengths = attention_mask.sum(dim=1) - 1 batch_size = last_hidden_states.shape[0] return last_hidden_states[ torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths, ] @torch.inference_mode() def encode(texts, max_length=8192): inputs = tokenizer( texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt", ).to(DEVICE) outputs = model(**inputs) embeddings = last_token_pool( outputs.last_hidden_state, inputs["attention_mask"], ) return F.normalize(embeddings, p=2, dim=1) ``` ### 3) Retrieve Insights for an Agent State Queries should include a retrieval instruction. Candidate insights are encoded directly. ```python retrieval_instruction = ( "Given the current agent state, retrieve insights that resolve the " "current bottleneck and help the agent make progress toward its goal" ) agent_state = """ Goal: buy a light-grey dining set under $250. History: searched twice and browsed several result pages. Observation: many partially matching products are visible, but none selected. """ query = f"Instruct: {retrieval_instruction}\nQuery:{agent_state}" insights = [ "Keep refining the query until an exact product appears.", "Include all critical attributes in the search, then select the required variant before purchasing.", "Browse more result pages before committing to a product.", ] query_embedding = encode([query]) insight_embeddings = encode(insights) scores = (query_embedding @ insight_embeddings.T).squeeze(0) top_indices = torch.topk(scores, k=2).indices.tolist() for rank, index in enumerate(top_indices, start=1): print(f"Insight {rank}: {insights[index]}") ``` For dynamic retrieval after every environment step, see the [ALFWorld, WebShop, and ScienceWorld implementations](https://github.com/ttchungc/InsightEmb). --- ## 🎯 Intended Use InsightEmb is designed for retrieval settings where relevance is primarily **procedural or structural**, including: - agent memory and insight retrieval - state-aware retrieval for interactive agents - skill and strategy retrieval - reasoning-intensive retrieval - retrieval of abstract rules with limited lexical overlap It can be used as a bi-encoder retriever or as the dense component of a hybrid retrieval pipeline. --- ## 📜 Citation If you find this work useful, please cite: ```bibtex @misc{chung2026insightemblearningactionintentembeddings, title={InsightEmb: Learning Action-Intent Embeddings for Agentic Insight Retrieval}, author={Tsz Ting Chung and Jiangnan Li and Jie Zhou and Mo Yu}, year={2026}, eprint={2608.04761}, archivePrefix={arXiv}, primaryClass={cs.CL}, url={https://arxiv.org/abs/2608.04761}, } ```