IMRNNs

Interpretable Modular Retrieval Neural Networks
Efficient, interpretable dense retrieval through dynamic embedding modulation.

Paper · PyPI · Model checkpoint · Project website · Author portfolio

Given a query and candidate documents, IMRNNs dynamically modulates both sides of their dense embeddings before ranking the documents by cosine similarity. A Query Adapter conditions each document representation on the query, while a Document Adapter uses feedback from the candidate set to adapt the query representation. The base encoder stays frozen.

IMRNNs operates inside the dense-retrieval stage. It does not add a downstream cross-encoder stage.

Choose your path

Installation

For text and embedding ranking with the released checkpoint:

python -m pip install imrnns

For training and BEIR evaluation:

python -m pip install "imrnns[train,eval]"

IMRNNs supports Python 3.10–3.12 and CPU inference. Set device="cuda" when a CUDA-enabled PyTorch installation and compatible GPU are available. The first text-ranking call downloads the pinned MiniLM encoder and the IMRNN checkpoint; later calls use the local Hugging Face cache.

Five-minute quick start

This walkthrough uses a genuine example from the BEIR SciFact corpus: claim 130 and three official document titles. SciFact marks document 27768226 as the relevant document in its test relevance judgments.

from imrnns import IMRNNAdapter

CLAIM = (
    "Articles published in open access format are more likely to be cited "
    "than traditional journals."
)

DOCUMENTS = {
    "27768226": "Open Access Increases Citation Rate",
    "38180456": "Short-term medical service trips: a systematic review of the evidence.",
    "16979690": (
        "Effect on the quality of peer review of blinding reviewers and asking "
        "them to sign their reports: a randomized controlled trial."
    ),
}

adapter = IMRNNAdapter.from_pretrained(
    encoder="minilm",
    dataset="scifact",
    device="cpu",
)

results = adapter.rank(
    query=CLAIM,
    documents=list(DOCUMENTS.values()),
    document_ids=list(DOCUMENTS),
    top_k=3,
)

for result in results:
    print(
        result.rank,
        result.document_id,
        f"base={result.base_score:.4f}",
        f"adapted={result.adapted_score:.4f}",
        f"delta={result.score_delta:+.4f}",
        sep=" | ",
    )

Observed with the released checkpoint on CPU:

1 | 27768226 | base=0.6616 | adapted=0.6380 | delta=-0.0236
2 | 38180456 | base=0.1311 | adapted=0.1115 | delta=-0.0196
3 | 16979690 | base=0.1331 | adapted=0.1096 | delta=-0.0235

The gold SciFact document is ranked first. Each RetrievalResult contains:

A negative score_delta does not mean the ranking became worse. It means that one pair's cosine score moved downward. Ranking quality depends on the relative ordering across all candidates; the corpus-level metrics later in this README measure that ordering directly.

Explain a retrieval decision

Use explain() with the same SciFact claim and its top document:

from pathlib import Path

explanation = adapter.explain(
    query=CLAIM,
    document=DOCUMENTS["27768226"],
    top_tokens=5,
)

print("query concepts:", explanation.query_tokens)
print("document concepts:", explanation.document_tokens)
print(
    f"base={explanation.base_score:.4f}",
    f"adapted={explanation.adapted_score:.4f}",
    f"delta={explanation.score_delta:+.4f}",
)

Path("imrnns-explanation.html").write_text(
    explanation.to_html(),
    encoding="utf-8",
)

Observed concept labels include bibliography, access, author, and citation for the query, and access and citation for the document. RetrievalExplanation also exposes the complete query and document modulation vectors through query_modulation and document_modulation.

The concepts are nearest encoder-vocabulary directions obtained through a Moore–Penrose back-projection. WordPiece fragments such as ##wall can appear, and these labels should be interpreted as inspection aids rather than causal natural-language rationales. explain() scores a single pair, whereas rank() can use feedback from several candidates, so their adapted scores can differ slightly.

Rank existing embeddings

rank_embeddings() accepts NumPy arrays, PyTorch tensors, or numeric Python sequences. The following example derives real embeddings from the same SciFact inputs instead of using random vectors:

from sentence_transformers import SentenceTransformer

BASE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
BASE_REVISION = "c9745ed1d9f207416be6d2e6f8de32d1f16199bf"

encoder = SentenceTransformer(
    BASE_MODEL,
    revision=BASE_REVISION,
    device="cpu",
)
query_embedding = encoder.encode(CLAIM, convert_to_numpy=True)
document_embeddings = encoder.encode(
    list(DOCUMENTS.values()),
    convert_to_numpy=True,
)

embedding_adapter = IMRNNAdapter.from_pretrained(
    encoder="minilm",
    dataset="scifact",
    device="cpu",
    load_encoder=False,
)

results = embedding_adapter.rank_embeddings(
    query_embedding=query_embedding,
    document_embeddings=document_embeddings,
    document_ids=list(DOCUMENTS),
    top_k=3,
)

This returns the same document order and scores as the text example because it uses the same pinned encoder and input strings. In a production retrieval system, replace these arrays with the vectors returned by your vector store. The query and document dimensions must match the checkpoint's input dimension.

Use a custom retrieval checkpoint

If you trained an adapter for another compatible encoder, load it without a text encoder and pass vectors directly:

custom_adapter = IMRNNAdapter.from_checkpoint(
    "./my-adapter.pt",
    encoder_model_name="my-organization/my-sentence-transformer",
    embedding_dim=768,
    load_encoder=False,
    device="cpu",
)

results = custom_adapter.rank_embeddings(
    query_embedding=query_vector,
    document_embeddings=document_vectors,
    document_ids=document_ids,
    top_k=10,
)

Checkpoint loading

The public release currently contains one ready-to-use adapter:

from_pretrained() is the simplest loader. To download the file explicitly:

imrnns download --encoder minilm --dataset scifact

The equivalent Python helper is:

from imrnns import download_checkpoint

downloaded = download_checkpoint(
    encoder="minilm",
    dataset="scifact",
)
print(downloaded.checkpoint_path)

Load a local copy directly:

adapter = IMRNNAdapter.from_checkpoint(
    "./imrnns-minilm-scifact.pt",
    encoder="minilm",
    device="cpu",
)

For reproducible deployments, pass a Hugging Face commit hash through revision. For offline execution after the files have been cached, pass local_files_only=True and, if needed, the same cache_dir used during the initial download.

Training and evaluation

Install the optional dependencies before running this workflow:

python -m pip install "imrnns[train,eval]"

One-command workflow

imrnns run downloads the BEIR dataset when necessary, builds the embedding cache, mines dense hard negatives, trains the adapter, and evaluates the saved checkpoint:

imrnns run \
  --encoder minilm \
  --dataset scifact \
  --datasets-dir ./datasets \
  --cache-dir ./cache/minilm-scifact \
  --output-dir ./checkpoints \
  --device cpu

The trained checkpoint is written to ./checkpoints/imrnns-minilm-scifact.pt, and the command prints a JSON report containing the training history, base metrics, adapted metrics, metric deltas, and the all-metrics pass indicator.

Run each stage separately

Build a reusable embedding and hard-negative cache:

imrnns cache \
  --encoder minilm \
  --dataset scifact \
  --datasets-dir ./datasets \
  --cache-dir ./cache/minilm-scifact \
  --device cpu

Train from that cache:

imrnns train \
  --encoder minilm \
  --dataset scifact \
  --datasets-dir ./datasets \
  --cache-dir ./cache/minilm-scifact \
  --output-dir ./checkpoints \
  --device cpu

Evaluate the saved adapter against the frozen base retriever:

imrnns evaluate \
  --encoder minilm \
  --dataset scifact \
  --datasets-dir ./datasets \
  --cache-dir ./cache/minilm-scifact \
  --checkpoint ./checkpoints/imrnns-minilm-scifact.pt \
  --device cpu

The evaluation output includes base_metrics, adapted metrics, metric_delta, and beats_base_all_metrics. The final field is true only when every adapted metric is strictly greater than its corresponding base metric.

Default training configuration

Use imrnns train --help or imrnns run --help to inspect and override the available hyperparameters.

Custom datasets

IMRNNs accepts a local dataset in the standard BEIR layout:

my_dataset/
├── corpus.jsonl
├── queries.jsonl
└── qrels/
    ├── train.tsv
    └── test.tsv

Minimal records using the same SciFact example look like this:

{"_id":"27768226","title":"Open Access Increases Citation Rate","text":"Scientific abstract text goes here."}

{"_id":"130","text":"Articles published in open access format are more likely to be cited than traditional journals."}

query-id	corpus-id	score
130	27768226	1

Load and inspect a split in Python:

from imrnns import IMRNNDataset

dataset = IMRNNDataset.from_beir_directory(
    "./my_dataset",
    split="test",
)
print(len(dataset.corpus), len(dataset.queries), len(dataset.qrels))

Use --dataset-path ./my_dataset instead of --dataset scifact with the training commands. When official train and test qrels exist, IMRNNs reserves 15% of the training queries for validation and keeps the test set untouched. A dataset with one qrels split uses a deterministic 70/15/15 partition.

Encoders

The built-in encoder aliases can be used when caching and training:

Only the MiniLM–SciFact adapter is distributed as a ready-to-use checkpoint. The other aliases are available for creating your own caches and checkpoints.

For another SentenceTransformers encoder, provide its model name and dimension:

imrnns run \
  --encoder-model-name my-organization/my-model \
  --embedding-dim 768 \
  --encoder-revision COMMIT_HASH \
  --dataset-path ./my_dataset \
  --cache-dir ./cache/custom-model \
  --output-dir ./checkpoints

Use --query-prefix and --passage-prefix when required by the selected encoder. Reuse the identical model, revision, prefixes, and embedding dimension when loading its checkpoint.

Published checkpoint evaluation

The released MiniLM–SciFact checkpoint was evaluated on the complete official 300-query SciFact test set. Base and adapted retrieval use the same top-100 candidate sets.

The training study contains dataset splits, configuration, additional datasets, and metric deltas.

Reference

Command-line interface

Run imrnns <command> --help for every available option.

Public Python API

The package also exports IMRNN, ModelConfig, EncoderSpec, RetrievalResult, RetrievalExplanation, and TokenAttribution for advanced integration.

Troubleshooting

Development

git clone https://github.com/YashSaxena21/IMRNNs.git
cd IMRNNs
python -m pip install -e ".[dev]"
pytest
ruff check src tests scripts examples
python -m build
twine check dist/*

Bug reports, feature proposals, and focused pull requests are welcome through GitHub Issues and GitHub Pull Requests.

Citation

If IMRNNs supports your research, please cite the EACL 2026 paper:

@inproceedings{saxena-etal-2026-imrnns,
  title = {{IMRNN}s: An Efficient Method for Interpretable Dense Retrieval via Embedding Modulation},
  author = {Saxena, Yash and Padia, Ankur and Gunaratna, Kalpa and Gaur, Manas},
  booktitle = {Findings of the Association for Computational Linguistics: EACL 2026},
  year = {2026},
  pages = {6324--6337},
  doi = {10.18653/v1/2026.findings-eacl.333},
  url = {https://aclanthology.org/2026.findings-eacl.333/}
}

Authors and acknowledgments

IMRNNs is authored by Yash Saxena, Ankur Padia, Kalpa Gunaratna, and Manas Gaur. Visit Yash Saxena's portfolio for additional projects and publications.

This work is associated with the University of Maryland, Baltimore County and the KAI² Lab, and was published in the Findings of EACL 2026.

License

The code, scripts, documentation, and checkpoints are available under the Creative Commons Attribution 4.0 International License. Attribution requirements are documented in ATTRIBUTION.md.