Agent-UIBK-v1

A 179M-parameter retriever for deep-research agents

UIBK-Embedding-v1 | BrowseComp-Plus | PyLate | FastPLAID

---

About Agent-UIBK-v1

A deep-research agent does not search with a user's question. It searches with short queries it writes itself, in the middle of a long chain of reasoning, and the reasoning says what the agent is actually looking for. Agent-UIBK-v1 is a retriever for that setting. It reads the agent's current reasoning together with its search query, in the format introduced by AgentIR:

Reasoning: <the agent's reasoning in the current turn>

Query: <the search query>

Agent-UIBK-v1 is UIBK-Embedding-v1 fine-tuned on agent search data. It keeps that model's design: one encoder holds a late-interaction (ColBERT) retriever and a dense retriever, and both are stored in one list of 128-dimensional vectors, so a plain MaxSim adds the two scores. It loads as a PyLate ColBERT model and works with any PLAID index.

Notably it:

• Reaches 72.41% accuracy on BrowseComp-Plus as the search tool of gpt-oss-120b (reasoning effort high) with search and get_document. That is on par with Agent-ModernColBERT (72.53, one question of 830 apart), the strongest gpt-oss-120b entry on the leaderboard.
• Reaches 64.10% with the search tool only, above Agent-ModernColBERT (63.86) and behind only AgentIR-4B (66.99), a model 22 times larger.
• Leads the retrieval-only BrowseComp-Plus leaderboard on 6 of its 8 metrics, with evidence nDCG@10 of 24.78 against 20.26 for Qwen3-Embedding-8B.
• Reads long inputs: queries (reasoning plus query) and documents up to 2048 tokens.

Results

BrowseComp-Plus: deep-research agent (gpt-oss-120b, reasoning effort high)

BrowseComp-Plus asks 830 hard questions over a fixed corpus of 100,195 web documents. The agent searches the corpus with the retriever (top 5 results, 512-token snippets) and, in the get_document setting, can also read a full document by its id. Accuracy is judged by Qwen3-32B.

The numbers of the other models are those of the BrowseComp-Plus leaderboard (October 2026). The UIBK rows are our own runs; the UIBK-Embedding-v1 row uses the tool-call repairs below but plain queries, as that model was not trained on reasoning.

How to read the comparison.

• Each row is a single run, and agent runs vary. Between two of our own fine-tunes with nearly equal quality, about 160 of the 830 questions were solved by one run and not by the other. Differences of one to three points between rows are therefore within run-to-run variation. The difference between Agent-UIBK-v1 and Agent-ModernColBERT with get_document is one question.

• Our runs use the benchmark's own agent client with two changes:

  • Reasoning-aware queries. The retriever receives "Reasoning: ...\n\nQuery: ..." as in AgentIR. This is the input the model was trained on.
  • Tool-call repairs. gpt-oss under vLLM 0.30 sometimes formats a tool call in a way the server does not parse. The client then executes such a message as the tool call, removes stray format tokens from tool names, and counts reasoning-only responses towards the iteration limit.

  The prompt, the tools, top-k, the snippet length and the judge are the benchmark's. AgentIR's own client patches similar vLLM issues. The repairs are not specific to this model: with UIBK-Embedding-v1 they moved accuracy from 56.57 (829 of 830 queries finished) to 59.04.

• Being able to read full documents matters more than any retriever difference here: get_document adds 8.3 points to Agent-UIBK-v1.

BrowseComp-Plus: retrieval only

Here the model ranks the corpus for the 830 plain questions, with no agent and no reasoning. Evidence documents are those needed to answer a question; gold documents also contain the answer.

Documents are truncated to 2048 tokens. Agent-UIBK-v1 leads at the top of the ranking, and Qwen3-Embedding-8B (45 times larger) leads at depth 1000.

General retrieval: NanoBEIR (13 datasets, NDCG@10)

The model is specialised for agent search and gives up some general retrieval quality. For ordinary queries, use UIBK-Embedding-v1.

Both rows come from our development evaluation, which uses exact MaxSim search with no index. Most of the loss is on ArguAna, whose queries are whole arguments. Two things changed there:

• Agent-UIBK-v1 lets its token head read every query token, because the search query comes after a long reasoning. UIBK-Embedding-v1 uses only the first 64 query tokens for token matching.
• The fine-tuning data contains no queries of that kind.

We have not separated these two effects.

How It Works

The architecture is that of UIBK-Embedding-v1:

• One trunk, two tops. The ModernBERT embeddings and the first 18 layers are shared. The token (ColBERT) head and the dense head each have their own copy of the last 4 layers.
• Token head. The ColBERT projection (768 → 128), one vector per token.
• Dense head. Latent-attention pooling gives one 768-dimensional vector. It is rotated and cut into 3 slices, which take the first three positions of the output list.
• One MaxSim. Dimension 0 of every vector is a tag, so a query token only matches document tokens and a slot only matches slots. Plain MaxSim is then, up to a constant per query, token MaxSim + 0.75 × dense cosine.

The one difference at inference is the query: up to 2048 tokens, and the token head reads all of them (UIBK-Embedding-v1: 256 tokens, of which the token head reads the first 64).

Model Details

Model Description

• Model Type: PyLate ColBERT model with an additional dense head
• Base model: UIBK-Embedding-v1 (from ModernBERT-base)
• Parameters: 179M (shared trunk and token top 149.0M, dense top 20.1M, output layer 9.9M)
• Query Length: 2048 tokens (both heads read all of them)
• Document Length: 2048 tokens
• Output Dimensionality: 128 per vector; one vector per token, of which the first 3 carry the dense embedding
• Similarity Function: MaxSim
• Prompts: query: and document: (required)
• Query format: Reasoning: <reasoning>\n\nQuery: <query> (Reasoning: Empty when there is no reasoning)
• Language: English
• License: Apache 2.0

Model Sources

• Model code: branched_transformer.py (the two-top encoder) and hybrid_tokens.py (the output layer) in this repository, the same code as in UIBK-Embedding-v1
• Documentation: PyLate Documentation

Usage

The model is used with PyLate. Three things differ from an ordinary ColBERT model:

• it needs trust_remote_code=True, because two small modules in this repository define the two-top encoder and the output layer;
• the prompts must be passed when encoding: prompt_name="query" for queries and prompt_name="document" for documents;
• queries should carry the agent's reasoning: Reasoning: <reasoning>\n\nQuery: <query>.

pip install -U pylate "transformers>=5.3" "sentence-transformers>=5.3"

Quick start

from pylate import models, rank

model = models.ColBERT("DataScience-UIBK/Agent-uibk-v1", trust_remote_code=True)


def agent_query(query: str, reasoning: str | None = None) -> str:
    """The input format the model was trained on (AgentIR)."""
    return f"Reasoning: {reasoning or 'Empty'}\n\nQuery: {query}"


query = agent_query(
    "planet known as the Red Planet",
    reasoning="The question describes a planet with a reddish surface caused by iron oxide. "
              "I should find which planet is called the Red Planet before checking its moons.",
)
documents = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
    "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
]

queries_embeddings = model.encode([query], is_query=True, prompt_name="query")
documents_embeddings = model.encode(documents, is_query=False, prompt_name="document")
print(queries_embeddings[0].shape, documents_embeddings[0].shape)
# (49, 128) (18, 128)

ranking = rank.rerank(
    documents_ids=[[1, 2, 3, 4]],
    queries_embeddings=queries_embeddings,
    documents_embeddings=[documents_embeddings],
)
print([(d["id"], round(d["score"], 2)) for d in ranking[0]])
# [(2, 48.42), (4, 48.35), (3, 48.32), (1, 48.2)]

MaxSim sums over query tokens, so scores grow with the length of the query. Compare scores only within one query.

Using it as an agent's search tool

In an agent loop, build the retriever input from the reasoning the LLM produced in the turn that issued the search call, and keep the plain query in the conversation. With gpt-oss on the Responses API, this is the text of the reasoning items of the response, as in AgentIR's client:

from pylate import indexes, models, retrieve

model = models.ColBERT("DataScience-UIBK/Agent-uibk-v1", trust_remote_code=True)
index = indexes.PLAID(index_folder="pylate-index", index_name="corpus", override=True, batch_size="auto")
index.add_documents(
    documents_ids=corpus_ids,
    documents_embeddings=model.encode(corpus_texts, batch_size=32, is_query=False, prompt_name="document"),
)
retriever = retrieve.ColBERT(index=index)


def search(query: str, reasoning: str | None, k: int = 5):
    text = f"Reasoning: {reasoning or 'Empty'}\n\nQuery: {query}"
    embeddings = model.encode([text], is_query=True, prompt_name="query")
    return retriever.retrieve(queries_embeddings=embeddings, k=k)[0]

For large collections of long documents, keep batch_size="auto" in indexes.PLAID(...): with 2048-token documents the fixed default scoring batch can run out of GPU memory. For indexing, reranking and loading an existing index, the examples in the UIBK-Embedding-v1 card apply unchanged.

Training Details

Training Data

• Tevatron/AgentIR-data (Apache 2.0): 5,238 search turns of a deep-research agent on 249 WebShaper questions. Each turn has the agent's reasoning and query, one positive and seven hard negatives. This is the data Agent-ModernColBERT was trained on.
• LRAT-Train (LRAT, SIGIR 2026): retrieval pairs mined from Tongyi-DeepResearch trajectories on InfoSeek questions over a Wikipedia corpus. The positive is the browsed document the judge accepted; the negatives are unbrowsed results of the same search.
  • We used the 95,811 of its 96,504 pairs that have at least 7 negatives, with the first 7 negatives of each.
  • LRAT's pairs hold only the search query. We rebuilt the AgentIR input from LRAT's released trajectories: the reasoning step right before the search call that issued the query. It was found for all pairs but one, which got "Empty".
• Overlap with BrowseComp-Plus.
  • No BrowseComp-Plus question shares a 10-word sequence with an AgentIR query.
  • No BrowseComp-Plus evidence document is among the AgentIR training documents.
  • The only 10-word sequence the InfoSeek questions share with BrowseComp-Plus questions is the stock phrase "was added to the National Register of Historic Places in".

Training Procedure

• Start: the training checkpoint of UIBK-Embedding-v1.
• Loss: a contrastive loss on each head, the token head's MaxSim (temperature 0.02) and the dense head's cosine (scale 50), with 7 hard negatives per query and the other queries' documents on the same GPU (8 queries) as in-batch negatives. There is no teacher.
• Batches: each batch comes from one of the two datasets, drawn with multinomial sampling (alpha 0.5). AgentIR makes up about 19% of the batches, so each AgentIR row is seen about 3.7 times.
• Hyperparameters:
  • one epoch, 3,157 steps of batch 32;
  • learning rate 3e-6, linear decay, no warm-up;
  • queries and documents up to 2048 tokens;
  • prompts query: and document: ;
  • bf16, 23.5 minutes on 4 H100 GPUs.
• Packaging: the output layer keeps UIBK-Embedding-v1's settings (slot weight 0.225, that is token MaxSim + 0.75 × dense cosine), without the 64-token limit on query tokens.

The learning rate, the 2048-token lengths and the AgentIR input format follow the recipe of Agent-ModernColBERT.

Framework Versions

• Python: 3.12
• Sentence Transformers: 5.3.0
• PyLate: 1.6.0
• Transformers: 5.3.0
• PyTorch: 2.9.0+cu128
• Accelerate: 1.15.0

The BrowseComp-Plus runs served gpt-oss-120b and the Qwen3-32B judge with vLLM 0.30.0.

Limitations

• English only.
• Specialised for agent search queries with reasoning. On ordinary short queries and on long argument-style queries it is weaker than UIBK-Embedding-v1 (see NanoBEIR above).
• The BrowseComp-Plus numbers are single runs with one LLM (gpt-oss-120b), and differences of a few points between runs are common.
• Documents were indexed up to 2048 tokens. Longer documents are truncated.
• Multi-vector output: a ColBERT index is larger than a single-vector index.
• Needs trust_remote_code=True, transformers >= 5.3 and the two prompts.

Citation

BibTeX

Agent-UIBK-v1

@misc{uibk2026agent,
  title={Agent-UIBK-v1: A Retriever for Deep-Research Agents},
  author={{Data Science Group, University of Innsbruck}},
  year={2026},
  howpublished={\url{https://huggingface.co/DataScience-UIBK/Agent-uibk-v1}},
}

UIBK-Embedding-v1

@misc{uibk2026embedding,
  title={UIBK-Embedding-v1: Unifying late Interaction and Bi-encoders with Knowledge distillation},
  author={{Data Science Group, University of Innsbruck}},
  year={2026},
  howpublished={\url{https://huggingface.co/DataScience-UIBK/uibk-embedding-v1}},
}

AgentIR

@article{chen2026AgentIR,
  title={AgentIR: Reasoning-Aware Retrieval for Deep Research Agents},
  author={Zijian Chen and Xueguang Ma and Shengyao Zhuang and Jimmy Lin and Akari Asai and Victor Zhong},
  year={2026},
  journal={arXiv preprint arXiv:2603.04384}
}

LRAT

@inproceedings{zhou2026lrat,
  title={Learning to Retrieve from Agent Trajectories},
  author={Zhou, Yuqi and Dai, Sunhao and Qu, Changle and Pang, Liang and Xu, Jun and Wen, Ji-Rong},
  booktitle={Proceedings of the 49th International ACM SIGIR Conference on Research and Development in Information Retrieval},
  year={2026}
}

BrowseComp-Plus

@article{chen2025BrowseCompPlus,
  title={BrowseComp-Plus: A More Fair and Transparent Evaluation Benchmark of Deep-Research Agent},
  author={Chen, Zijian and Ma, Xueguang and Zhuang, Shengyao and Nie, Ping and Zou, Kai and Liu, Andrew and Green, Joshua and Patel, Kshama and Meng, Ruoxi and Su, Mingyi and Sharifymoghaddam, Sahel and Li, Yanxi and Hong, Haoran and Shi, Xinyu and Liu, Xuye and Thakur, Nandan and Zhang, Crystina and Gao, Luyu and Chen, Wenhu and Lin, Jimmy},
  year={2025},
  journal={arXiv preprint arXiv:2508.06600}
}

PyLate

@inproceedings{DBLP:conf/cikm/ChaffinS25,
  author       = {Antoine Chaffin and Rapha{\"{e}}l Sourty},
  title        = {PyLate: Flexible Training and Retrieval for Late Interaction Models},
  booktitle    = {Proceedings of the 34th {ACM} International Conference on Information and Knowledge Management, {CIKM} 2025},
  pages        = {6334--6339},
  publisher    = {{ACM}},
  year         = {2025},
  url          = {https://github.com/lightonai/pylate},
  doi          = {10.1145/3746252.3761608},
}

Acknowledgements

We thank the authors of AgentIR and of LRAT for releasing their agent search data and trajectories; LightOn for PyLate, FastPLAID and Agent-ModernColBERT, whose recipe we followed; the BrowseComp-Plus team for a fair and reproducible benchmark; and OpenAI and Qwen for the open models used as the agent and the judge.