InsightEmb (Learning Action-Intent Embeddings for Agentic Insight Retrieval)

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.

✨ 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

pip install torch transformers>=4.53.0

1) Load the Model

Replace YOUR_HF_NAMESPACE/InsightEmb with the model repository ID.

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.

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.

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.

🎯 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:

@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}, 
}