针对向量数据库检索的黑洞攻击

本仓库包含论文 Black-Hole Attack from Embedding Space Defects 配套的实现代码。黑洞攻击利用高维嵌入空间中的一种几何缺陷——中心性驱动的枢纽性（centrality-driven hubness）——来污染向量数据库：只需在簇质心附近注入少量恶意向量，这些向量便会成为大量用户查询的最近邻，而无需对查询内容有任何了解。

环境要求

• Python ≥ 3.10
• FAISS（推荐 GPU 版本：faiss-gpu；CPU 版本：faiss-cpu）
• NumPy、Pandas、PyArrow（用于 parquet 读写）
• scikit-learn、tqdm
• PyTorch 与 Transformers（仅 scripts/encode_*.py 需要；攻击流程本身不导入 torch）

安装核心依赖：

pip install numpy pandas pyarrow scikit-learn tqdm faiss-gpu

快速开始

1. 下载并准备数据

# Download BEIR datasets
bash scripts/download_data.sh

# Encode to vectors + build FAISS indexes
bash scripts/process_data.sh

在每个脚本中编辑 MODELS 和 DATASETS 数组，以选择特定的模型或数据集。

2. 运行基本攻击

# Full-Export Attack on HotpotQA with Contriever
python run.py --model contriever --src hotpotqa

这会运行完整流水线：加载向量 → 聚类 → 生成对抗向量 → 注入 → 构建投毒索引 → 评估 MO@10、ASR、FPR。

3. 关键参数

参数	默认值	描述
--model / --src	contriever / hotpotqa	嵌入模型和源数据集
--mode	default	default（同数据集）或 transfer（跨数据集）
--preprocess	default	default、query_trans 或 multi_query_transfer
--cluster	kmeans	kmeans、minibatch_kmeans、faiss_gpu
--num-copies / --epsilon	10 / 0.001	每个质心的扰动副本数量；扰动幅度
--eval-k / --sample-queries	10 / 3000	Top‑k 评估；采样的查询数量
--eval-index-types	FlatIP	要评估的索引类型：FlatIP、HNSW、IVF、IVFPQ
--eval-runners	evaluate	评估步骤：evaluate、evaluate_recall、evaluate_defense、evaluate_defense_performance_loss、evaluate_mitigation_defense、evaluate_mitigation_performance_loss
--result-subdir	自动	data/result/ 下用于存放输出 JSON 的子目录

攻击路径

论文中描述的四种攻击路径对应以下调用方式：

路径	命令
(1) 完整数据库导出	python run.py --mode default --preprocess default
(2) 部分数据库导出	在聚类前从外部导出部分向量；流水线其余部分与 (1) 相同
(3) 代理数据集攻击	python run.py --mode transfer --preprocess multi_query_transfer --victim <dataset>
(4) 投毒公开预嵌入数据集	修改预嵌入数据集的 .npy 文件并重新上传至 HuggingFace

对于 (3)，攻击者利用所有其他数据集的查询构建恶意向量，并将其注入受害者数据集。代理源数据集的查询 .npy 文件必须存在于 data/vector/ 中。

防御评估

要运行基于检测的防御：

python run.py --model contriever --src hotpotqa \
    --eval-runners evaluate_defense evaluate_defense_performance_loss

对于基于缓解的防御（一种或多种方法）：

python run.py --model contriever --src hotpotqa \
    --eval-runners evaluate_mitigation_defense evaluate_mitigation_performance_loss \
    --mitigation-methods cl2 zn nohub

演示

demo/ 包含一个自包含的最小示例，它从 BEIR 数据集构建一个小型知识库，对其进行编码，注入恶意向量，并测量检索层面的影响。它使用自己的 requirements.txt，不依赖 src/。CLI 详情请参见 demo/src/main.py。

数据集与模型

config.json 定义了支持的模型（contriever、bge、gte）和数据集（14 个 BEIR 基准）。要添加新模型，请扩展 config.json，并确保 .npy 向量文件存在于 data/vector/<model>_<dataset>.npy 下。

引用

@misc{li2026trustvectorsvectordatabase,
      title={Can You Trust the Vectors in Your Vector Database? Black-Hole Attack from Embedding Space Defects}, 
      author={Hanxi Li and Jianan Zhou and Jiale Lao and Yibo Wang and Zhengmao Ye and Yang Cao and Junfen Wang and Mingjie Tang},
      year={2026},
      eprint={2604.05480},
      archivePrefix={arXiv},
      primaryClass={cs.CR},
      url={https://arxiv.org/abs/2604.05480}, 
}