SHAQ

Reference implementation of SHAQ (shadow query generation), a defense against embedding inversion attacks (EIAs) on cloud-hosted vector databases.

EIAs such as vec2text reconstruct a document from its stored embedding by exploiting the tight coupling between an embedding and its text. SHAQ breaks that coupling. Instead of perturbing document embeddings, it replaces them:

1. Generation. A generative LM writes n_g shadow queries per document — plausible user queries covering different semantic facets. K-Means clustering keeps n_k diverse representatives. This is semantic decomposition.
2. Indexing. The shadow queries are encoded, and their embeddings replace the document embeddings in the vector database, scattered in random order with a mapping back to document IDs. This is embedding decoupling.

A breached database yields embeddings of queries about documents, never the documents themselves — while retrieval still succeeds whenever a user query matches any one facet.

Install

conda create -n shaq python=3.12.3 && conda activate shaq
pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/cu121
pip install -e ".[quant]"

[quant] adds bitsandbytes for 8-bit loading of the 32B generator; skip it if you only need evaluation. For development, pip install -e ".[dev]" adds pytest and ruff.

Shadow query generation uses a gated Hugging Face model. Authenticate through the environment — never hardcode a token:

export HF_TOKEN=hf_...

Quick start

shaq prepare-data --datasets scifact          # download, subsample, truncate, encode
shaq generate     --datasets scifact          # Algorithm 1 (needs a GPU)
shaq evaluate     --experiment main           # retrieval utility
shaq attack       --experiment main           # defense efficacy
shaq analyze      --datasets scifact          # decomposition / decoupling
shaq make-tables  --experiment main           # render results as markdown or LaTeX

Every command takes --help. Artifacts land under data/, results under results/; both are gitignored and regenerable. See docs/DATA.md.

To reproduce a specific table from the paper, see docs/REPRODUCE.md.

Layout

configs/          datasets, models, and one YAML per experiment
src/shaq/
  data/           BEIR loading, subsampling, preprocessing, artifact stores
  generation/     Algorithm 1: prompts, LM wrapper, parsing, K-Means selection
  index/          Algorithm 2: VectorStore, scattering, FAISS search
  defenses/       none | noise | secret_scaling | shaq, behind one interface
  attacks/        vec2text standard and adaptive attacks
  metrics/        retrieval utility, defense efficacy, decomposition analysis
  cli/            the `shaq` command
tests/            runs without network, GPU or model downloads

The organising idea is the defense interface (src/shaq/defenses/base.py). Every defense answers the same three questions — what goes in the vector database and what each row maps to, what the client sends at query time, and what an adversary observes — so retrieval, attack and evaluation each have one code path covering all four defenses.

Configuration

Experiments are YAML, not command-line archaeology:

name: main
seed: 42
embedding_model: gtr-t5-base
truncate: true
datasets: [scifact, nq, ...]
defenses:
  - {kind: none}
  - {kind: noise, sigma: 0.01}
  - {kind: shaq, n_k: 10}

Run any of them with shaq evaluate --experiment <name>. Ad-hoc runs still work via flags (--dataset scifact --defense shaq --n-k 10).

Tests

pytest

No network, GPU or model downloads: a deterministic hashing encoder stands in for the real one, so the invariants under test stay separable from model behaviour. The suite that matters most is tests/test_index_mapping.py, which pins the property everything else depends on — a stored vector always resolves to the document whose shadow query produced it, under scattering, uneven group sizes and corpus reordering.

Environment notes

The paper's experiments ran on Rocky Linux with NVIDIA H100s, Python 3.12.3, PyTorch 2.3.0 and transformers 4.44.2.

macOS: the torch and faiss-cpu wheels each bundle their own libomp, and two OpenMP runtimes in one process make multi-threaded FAISS search segfault. Run FAISS single-threaded:

export KMP_DUPLICATE_LIB_OK=TRUE
export SHAQ_FAISS_THREADS=1

Citation

@article{feng2026shaq,
  title  = {Shadow Queries for Private Retrieval in Vector Databases},
  author = {Feng, Xinguo and Ma, Zhongkui and Wang, Zihan and Yan, Chuan and
            Yang, Guowei and Abuadbba, Alsharif and Bai, Guangdong},
  year   = {2026}
}

License

MIT — see LICENSE.