castorini/slideon-onnx

SlideOn ONNX

A graph-optimized ONNX export of LightOn's SlideOn, an English sparse retrieval model, prepared for Anserini integration. This repository distributes an ONNX conversion, not a new upstream LightOn checkpoint.

Base checkpoint: 370c600. The original Apache-2.0 license is included in LICENSE.

Setup

The model was exported and optimized on Linux x86_64 (kernel 6.8.0), using Python 3.12.3 and CPU execution (no CUDA).

ONNX export

Package	Version
PyTorch	2.14.1+cpu
Transformers	5.21.0
ONNX	1.23.2
ONNX Runtime (validation)	1.31.1

The export uses torch.onnx.export with opset 18, dynamo=False, eager attention, and FP32 model parameters. SlideOn's native FP16 roundtrips are preserved. Batch size is fixed at 1; sequence length is dynamic.

Graph optimization

Optimization was performed in a separate environment using ONNX Runtime 1.17.0, NumPy 1.26.4, and CPUExecutionProvider, matching the runtime version used by the Anserini integration.

The exported graph was optimized with ORT_ENABLE_EXTENDED and saved as slideon-optimized.onnx. No additional FP16 or INT8 conversion was applied.

Runtime pruning

The graph now requires a fourth input, top_k (scalar INT64), in addition to input_ids, attention_mask, and is_document. Set top_k=0 to disable pruning, or a positive value to keep at most that many features (clamped to 32768). Use nonnegative values. The setting applies to both query and document modes.

Top-k selection occurs before the final FP16 rounding. Output weights remain FP32; scaling and integer quantization belong to the caller. Equal pre-rounding weights are resolved by lower latent ID; exact PyTorch tie parity is not guaranteed.

This is an input-interface change. Direct ONNX clients must download the new model and supply top_k; old clients without that input cannot run this graph.

Use an Anserini build containing runtime-top-k support. Both the hosted artifact and the local cache file are named slideon-optimized.onnx. Existing caches are not automatically refreshed: delete any old cached model once after upgrading (default: ~/.cache/pyserini/encoders/slideon-optimized.onnx; use your custom cache directory if configured). Keep the tokenizer and config files. Anserini downloads the new model on the next run. New users need no cleanup or Python conversion step.

Anserini defaults to query k=64. To override it, put the JVM property -Danserini.slideon.topK=128 before the Java main class (or use 0 for no pruning). Changing query k does not change an existing document index.

The runtime-pruning graph passed four synthetic ONNX regression tests and eleven targeted Java tests, including local-model inference at k=0, 1, 32, and 128. These are functional checks, not a claim of full native-PyTorch numerical parity.