Joint Product Quantization

Joint Product Quantization (JPQ) first learns product quantization (PQ) centroids over document embeddings from an existing index, then jointly optimizes these centroids and the query encoder.

Training

JPQ training requires:

  • A BiEncoder model

  • A FlexIndex containing document embeddings produced by the BiEncoder

  • A list of training examples, where each item is a dict containing

    • query (str): the input query text

    • doc_id_a (str): document identifier for a relevant (positive) document

    • doc_id_b (str): document identifier for a non-relevant (negative) document

pyterrier_dr.jpq.JPQTrainer implements the training code. Call fit() to train JPQ. After training, call jpq_index() to save the JPQIndex. See example below:

Example for training JPQ with E5
from pyterrier_dr import FlexIndex, E5
from pyterrier_dr.jpq import JPQTrainer

index = FlexIndex("path/to/e5_index")
model = E5()
trainer = JPQTrainer(model, index, M=96, nbits=8, pq_impl="faiss2opq")

training_docpairs = [
    {
        "query": "chemical reactions in water",
        "doc_id_a": "doc_1",
        "doc_id_b": "doc_2",
    },
    # ...
]
save_path = "e5_jpq"
trainer.fit(training_docpairs=training_docpairs)
jpq_index = trainer.jpq_index(save_path)
trainer.query_encoder.model.save_pretrained(save_path)

Retrieval

JPQ retrieval requires:

  • The fine-tuned query encoder

  • The corresponding JPQIndex

With a JPQIndex, you can create a PQ retriever by calling retriever_pq(). An example is provided below:

Example for retrieval using JPQ (E5)
from pyterrier_dr import E5
from pyterrier_dr.jpq import JPQIndex
from sentence_transformers import SentenceTransformer

path = "e5_jpq"

model = E5()
model.model = SentenceTransformer(path, device="cuda")
jpq_index = JPQIndex(path)

pipeline = model.query_encoder() >> jpq_index.retriever_pq()
results = pipeline.search("chemical reactions in water")

API Reference