KaLM-Embedding / KaLM-Reranker-V1-Small-R2

KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking

What is KaLM-Reranker-V1-R2?

KaLM-Reranker-V1-R2 is a substantially improved checkpoint release of KaLM-Reranker-V1. It keeps the same fast-but-not-late-interaction (FBNL) architecture and inference interface as the original release, while introducing a stronger multi-stage training recipe for three practical goals:

1. Stronger compression robustness. R2 supports Matryoshka embedding pooling (MEP) from 1x to 128x , extending the maximum validated compression ratio from 32x to 128x. Even at 128x compression, all three model sizes retain at least 92% of their average nDCG@10 at 2x compression on both BEIR and MIRACL.
2. Adjustable test-time compute scaling. A single KaLM-Reranker-V1-R2 checkpoint can trade compute for quality by changing the compression ratio. It can also use a coarse-to-fine cascade: cheaply screen all candidates with heavily compressed document representations, then spend more compute only on the most promising candidates.
3. Better multi-domain and multilingual ranking. At the default 4x compression ratio, every R2 model improves on BEIR, while MIRACL average nDCG@10 increases by +9.03 , +7.17 , and +4.85 points for Nano, Small, and Large, respectively.

KaLM-Reranker-V1-R2 vs. the original KaLM-Reranker-V1

All results below use the same default compression ratio, r=4 . KaLM-Reranker-V1 results are taken from the previous model cards; KaLM-Reranker-V1-R2 results are from the updated paper.

BEIR — Multi-domain Reranking

MIRACL — Multilingual Reranking

BEIR — Compression Robustness

MIRACL — Compression Robustness

KaLM-Reranker-V1-R2 consistently outperforms KaLM-Reranker-V1 at every shared compression ratio. The gains increase as compression becomes more aggressive, demonstrating substantially improved compression robustness.

Model family

The reported sizes are activated parameters . Nano, Small, and Large are initialized from the T5Gemma2 270M-270M, 1B-1B, and 4B-4B encoder-decoder families, respectively.

Prompt Template

 f"<Document>: {document}"

(
    f"<bos><start_of_turn>user\n"
    f"Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".\n\n"
    f"<Instruct>: {task_instruction}\n"
    f"<Query>: {query}<end_of_turn>\n"
    f"<start_of_turn>model\n\n\n\n"
)

Evaluation

BEIR

The KaLM-Reranker-V1 series matches or outperforms strong industrial rerankers from the Qwen, BGE, Jina, and Mxbai families.

MIRACL

KaLM-Reranker-V1 achieves competitive performance on MIRACL. Within its parameter group, KaLM-Reranker-V1-Large ranks first in 11 of the 18 languages (i.e., ar, bn, de, en, es, fi, fr, ja, ko, ru, and sw); six of these languages (i.e., en, es, de, ja, fr, and ru) have reported website usage rates totaling approximately 74.2% among websites with known content languages ( W3Techs, 2026 ).

LMEB

On LMEB-Dialogue, a compact embedding model paired with our Nano reranker, which has only 0.27B activated parameters, remains competitive with 7–12B embedding models.

Ablation on multi-stage training

Across all three model sizes and all seven compression ratios, performance on BEIR and MIRACL improves consistently from Stage 1 to Stage 3, demonstrating the effectiveness of our multi-stage training pipeline. Concretely, Stage 1 establishes a robust foundation for document reranking, distillation in Stage 2 substantially improves performance, and Stage 3 yields further modest gains. More importantly, robustness to compression generally improves across the three training stages. For example, from Stage 1 to Stage 3, the performance retention of KaLM-Reranker-V1-Nano at r = 128 relative to r = 2 increases from 92.88% to 93.80% on BEIR and from 90.93% to 92.15% on MIRACL.

Potential of test-time compute scaling

Results show that this multi-stage strategy largely preserves reranking effectiveness while reducing the estimated online computation cost by a factor of several. For example, on BEIR, KaLM-Reranker-V1-Small achieves an average nDCG@10 of 61.11 under setting (c), compared with 61.17 under setting (a), while reducing the estimated relative serving cost from 8.65x to 2.75x. Interestingly, settings (c) and (d) outperform setting (a) on some tasks, suggesting that reranking more candidates at a low compression ratio does not necessarily improve reranking quality.

Using Sentence Transformers

KaLM-Reranker-V1-Small-R2 can be loaded as a modular Sentence Transformers CrossEncoder . This integration requires sentence-transformers>=5.6,<6 , transformers>=5.3,<6 , and the PyTorch backend. The repository contains custom modeling code, so load only trusted revisions and pass trust_remote_code=True .

pip install "sentence-transformers>=5.6,<6" "transformers>=5.3,<6"

import torch
from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "KaLM-Embedding/KaLM-Reranker-V1-Small-R2",
    trust_remote_code=True,
    device="cuda",
    model_kwargs={"dtype": torch.bfloat16, "chunk_size": 4},
)

query = "What is the capital of China?"
documents = [
    "The capital of China is Beijing.",
    "Gravity attracts bodies toward one another.",
]
pairs = [(query, document) for document in documents]

# The default output is P(yes).
scores = model.predict(pairs)
rankings = model.rank(query, documents, return_documents=True)

# CrossEncoder prompts are interpreted as KaLM task instructions.
instruction = "Given a web search query, retrieve passages that answer the query."
custom_scores = model.predict(pairs, prompt=instruction)
custom_rankings = model.rank(query, documents, prompt=instruction)

# Use Identity to return yes_logit - no_logit instead of P(yes).
margins = model.predict(pairs, activation_fn=torch.nn.Identity())

print(f"scores: {scores}")
print(f"rankings: {rankings}")
print(f"custom_scores: {custom_scores}")
print(f"custom_rankings: {custom_rankings}")
print(f"margins: {margins}")

'''
scores: [9.95668530e-01 1.07208805e-04]
rankings: [{'corpus_id': 0, 'score': 0.99566853, 'text': 'The capital of China is Beijing.'}, {'corpus_id': 1, 'score': 0.000107208805, 'text': 'Gravity attracts bodies toward one another.'}]
custom_scores: [9.8967183e-01 2.0145691e-05]
custom_rankings: [{'corpus_id': 0, 'score': 0.9896718}, {'corpus_id': 1, 'score': 2.0145691e-05}]
margins: [ 5.4375   -9.140625]
'''

Inputs must be ordered as (query, document) . By default, queries are truncated to 512 tokens and documents to 1024 tokens. chunk_size=4 performs a mask-aware mean over each consecutive group of four encoder token states before passing the compressed encoder output to the decoder. Set chunk_size=None to disable compression, or change model[0].chunk_size after loading.

For CPU inference, use device="cpu" and model_kwargs={"dtype": torch.float32, "chunk_size": 4} . Only the PyTorch inference backend is currently supported; training, ONNX, and OpenVINO are not included in this release.

Using transformers

import argparse
from typing import Optional


def optional_positive_int(value: str) -> Optional[int]:
    if value.lower() == "none":
        return None
    try:
        parsed = int(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError(
            "must be a positive integer or 'none'"
        ) from error
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be a positive integer or 'none'")
    return parsed


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--model",
        default="KaLM-Embedding/KaLM-Reranker-V1-Small-R2",
        help="Hugging Face model ID or local checkpoint path.",
    )
    parser.add_argument(
        "--device",
        default=None,
        help="Inference device, such as 'cuda', 'cuda:0', or 'cpu'.",
    )
    parser.add_argument(
        "--dtype",
        default=None,
        choices=("bfloat16", "bf16", "float16", "fp16", "float32", "fp32"),
        help="Model parameter dtype. By default, use BF16 on CUDA and FP32 on CPU.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=32,
        help="Number of query-document pairs scored per inference batch.",
    )
    parser.add_argument(
        "--query-max-length",
        type=int,
        default=512,
        help=(
            "Maximum tokens in the raw query before it is inserted into the "
            "decoder prompt; prompt tokens are not included in this limit."
        ),
    )
    parser.add_argument(
        "--reranker-max-length",
        type=int,
        default=1024,
        help=(
            "Maximum encoder tokens for '<Document>: {passage}'. This is not a "
            "combined query-document context limit."
        ),
    )
    parser.add_argument(
        "--chunk-size",
        type=optional_positive_int,
        default=4,
        metavar="N|none",
        help=(
            "Number of encoder token hidden states per mean-pooled chunk; use "
            "'none' to disable encoder chunk pooling."
        ),
    )
    return parser


def main() -> None:
    args = build_parser().parse_args()

    from kalm_reranker import KaLMReranker

    reranker = KaLMReranker(
        args.model,
        device=args.device,
        dtype=args.dtype,
        batch_size=args.batch_size,
        query_max_length=args.query_max_length,
        max_length=args.reranker_max_length,
        chunk_size=args.chunk_size,
    )
    query = "What is the capital of China?"
    documents = [
        "The capital of China is Beijing.",
        "Gravity attracts bodies toward one another.",
    ]
    instruction = "Given a query, retrieve documents that answer the query."

    pairs = [(query, document) for document in documents]
    print("scores:", reranker.predict(pairs, instruction=instruction))
    print("rankings:", reranker.rank(query, documents, instruction=instruction))
    

if __name__ == "__main__":
    main()

'''
scores: [0.9956685304641724, 0.00010720881255110726]
rankings: [{'corpus_id': 0, 'score': 0.9956685304641724}, {'corpus_id': 1, 'score': 0.00010720881255110726}]
'''

Using vLLM

An experimental single-GPU adapter is available for offline LLM.classify() reranking and optional FastAPI serving. It reuses the original checkpoint without adding or modifying model weights.

The adapter has been validated with Python 3.12, vLLM 0.19.1, Transformers 5.6.2 and CUDA BF16:

conda create -n kalm-vllm python=3.12 -y
conda activate kalm-vllm
pip install "vllm==0.19.1" "transformers==5.6.2"

hf download KaLM-Embedding/KaLM-Reranker-V1-Small-R2 \
  --local-dir ./KaLM-Reranker-V1-Small-R2
pip install ./KaLM-Reranker-V1-Small-R2/vllm_support --no-deps
export VLLM_PLUGINS=kalm_t5gemma2

If you modify the vllm_support source code for some reason (e.g. changing the model path in vllm_support/src/kalm_t5gemma2_vllm_plugin/constants.py ), please run pip install ./KaLM-Reranker-V1-Small-R2/vllm_support --no-deps --force-reinstall to refresh the patch.

Offline Python:

from kalm_t5gemma2_vllm_plugin import KaLMVLLMReranker

query = "What is the capital of China?"
documents = [
    "The capital of China is Beijing.",
    "Gravity attracts bodies toward one another.",
]

with KaLMVLLMReranker(
    "KaLM-Embedding/KaLM-Reranker-V1-Small-R2",
    query_max_length=512,
    document_max_length=1024,
    encoder_chunk_size=4,
) as reranker:
    print(reranker.rank(query, documents))

Offline CLI:

kalm-vllm-rerank --return-margin

To deploy the online service, install the HTTP dependencies and keep the server running in the first terminal:

pip install "fastapi>=0.136,<0.137" "uvicorn>=0.46,<0.47"
export CUDA_VISIBLE_DEVICES=0
export VLLM_PLUGINS=kalm_t5gemma2

kalm-vllm-serve \
  --host 0.0.0.0 \
  --port 8000 \
  --model KaLM-Embedding/KaLM-Reranker-V1-Small-R2 \
  --query-max-length 512 \
  --document-max-length 1024 \
  --encoder-chunk-size 4 \
  --max-model-len 2048

In a second terminal, check the server:

conda activate kalm-vllm
kalm-vllm-client --base-url http://127.0.0.1:8000 --health

Use /rerank for one query and a list of documents. Results are sorted by score:

kalm-vllm-client \
  --base-url http://127.0.0.1:8000 \
  --endpoint rerank \
  --json-file ./KaLM-Reranker-V1-Small-R2/vllm_support/examples/rerank_request.json \
  --return-margin \
  --top-k 10

Use /score to score a batch of independent query-document pairs. Results preserve the input order and optional IDs:

kalm-vllm-client \
  --base-url http://127.0.0.1:8000 \
  --endpoint score \
  --json-file ./KaLM-Reranker-V1-Small-R2/vllm_support/examples/score_request.json \
  --return-margin

The default output is P(yes) . Set return_margin=true to also receive yes_logit - no_logit ; the client flag --return-margin applies the same setting to a JSON file request. The supported encoder chunk sizes are 1, 2, 4, 8, 16, 32 , with 4 as the default.

This adapter uses vLLM's plugin, scheduling and pooling interfaces while the T5Gemma2 semantic forward still runs through Transformers. It is not vLLM's native HTTP /score implementation or a complete vLLM-native kernel port. See the complete installation, API and troubleshooting guide .

Acknowledgements

We sincerely thank jina-reranker-v3 and Qwen3-Reranker for their valuable inspiration and contributions to the reranking community, from which we have learned a lot.

Citation

If you find this model useful, please consider citing our papers.

@misc{zhao2026kalmrerankerv1,
      title={KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking}, 
      author={Xinping Zhao and Jiaxin Xu and Ziqi Dai and Xin Zhang and Shouzheng Huang and Danyu Tang and Xinshuo Hu and Meishan Zhang and Baotian Hu and Min Zhang},
      year={2026},
      eprint={2606.22807},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2606.22807}, 
}

@misc{zhao2026kalmembeddingv2,
      title={KaLM-Embedding-V2: Superior Training Techniques and Data Inspire A Versatile Embedding Model}, 
      author={Xinping Zhao and Xinshuo Hu and Zifei Shan and Shouzheng Huang and Yao Zhou and Xin Zhang and Zetian Sun and Zhenyu Liu and Dongfang Li and Xinyuan Wei and Youcheng Pan and Yang Xiang and Meishan Zhang and Haofen Wang and Jun Yu and Baotian Hu and Min Zhang},
      year={2026},
      eprint={2506.20923},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2506.20923}, 
}

@misc{hu2025kalmembedding,
      title={KaLM-Embedding: Superior Training Data Brings A Stronger Embedding Model}, 
      author={Xinshuo Hu and Zifei Shan and Xinping Zhao and Zetian Sun and Zhenyu Liu and Dongfang Li and Shaolin Ye and Xinyuan Wei and Qian Chen and Baotian Hu and Haofen Wang and Jun Yu and Min Zhang},
      year={2025},
      eprint={2501.01028},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2501.01028}, 
}

Contact

If you encounter any issues, feel free to contact us via the email: [email protected]