Skip to content

Commit 5690b41

Browse files
driPyfdriPyf
committed
simplify README
1 parent 27fe19c commit 5690b41

2 files changed

Lines changed: 25 additions & 274 deletions

File tree

‎README.md‎

Lines changed: 15 additions & 215 deletions
Original file line numberDiff line numberDiff line change
@@ -1,239 +1,39 @@
1-
# OMEGA: One-Model Efficient Generalized Approximate Nearest Neighbor Search
1+
# OMEGA
22

3-
OMEGA is an adaptive approximate nearest neighbor (ANN) search system that dynamically adjusts search effort per query with minimal preprocessing cost. It achieves significant speedups by using machine learning to predict when sufficient search has been performed to find top-K results with high confidence.
3+
Adaptive ANN search that learns when to stop searching. Train once on K=1, works for any K.
44

5-
---
5+
Built on HNSW + LightGBM. Adapts search effort per query instead of using fixed parameters for everything.
66

7-
## What is OMEGA?
7+
## Installation
88

9-
OMEGA is a K-generalizable learned vector search system built on top of the HNSW (Hierarchical Navigable Small World) algorithm. The key innovation is that **a single model trained on K=1** can generalize to any K value through iterative refinement and trajectory-based features.
10-
11-
It combines:
12-
13-
- **Python Frontend**: Built on the [big-ann-benchmarks](./README_Big_ANN_Benchmarks.md) framework for evaluation and orchestration
14-
- **C++ Backend**: Modified [hnswlib](./hnswlib/) with adaptive search capabilities
15-
- **Machine Learning**: LightGBM models to predict optimal stopping points during search
16-
17-
The system learns from query patterns to adaptively determine when to stop searching, eliminating the need for manual parameter tuning while maintaining high recall.
18-
19-
---
20-
21-
## What Problem Does It Solve?
22-
23-
Traditional ANN search algorithms and learned search systems face two key challenges:
24-
25-
### 1. Fixed Parameters Problem
26-
Traditional algorithms use **fixed parameters** (like `ef` in HNSW) for all queries:
27-
- **Easy queries** are over-searched, wasting computation
28-
- **Hard queries** are under-searched, potentially missing relevant results
29-
- **Manual tuning** is required for each dataset and recall target
30-
31-
### 2. Multi-K Preprocessing Cost Problem
32-
Existing learned search systems (like DARTH, LAET) require:
33-
- **Training separate models for each K value** (K=1, K=10, K=100, etc.)
34-
- **Expensive preprocessing** that can exceed 22% of serving cost in production
35-
- **Poor generalization** across different K values
36-
37-
OMEGA solves both problems by:
38-
- **Adapting search effort per query** based on learned patterns
39-
- **Using ONE model trained on K=1** that generalizes to all K values
40-
- **Minimal preprocessing cost** (only train once for K=1)
41-
- **Iterative refinement** that reduces top-K to K sequential top-1 problems
42-
43-
---
44-
45-
## How Does It Work?
46-
47-
OMEGA uses a **K-generalizable learning approach** with three key innovations:
48-
49-
### 1. Trajectory-Based Features
50-
Instead of query-specific features, OMEGA uses **distance variation patterns** during graph traversal:
51-
- Distance to current candidate
52-
- Distance changes over search trajectory
53-
- Statistical properties of the candidate pool
54-
- These features naturally generalize across K values through **masking**
55-
56-
### 2. Iterative Refinement via Masking
57-
OMEGA reduces top-K search to **K sequential top-1 problems**:
58-
- Train a single binary classifier for K=1: "Has top-1 been found?"
59-
- For K>1: Mask already-found results and repeat K times
60-
- Each iteration uses the same K=1 model
61-
- No need to train separate models for different K values
62-
63-
### 3. Statistical Forecasting
64-
To reduce model invocation overhead:
65-
- Pre-profile probability tables during training
66-
- Map recall targets to prediction intervals
67-
- Use lookup tables for fast inference
68-
- Only invoke model at strategic checkpoints
69-
70-
### Training Phase
71-
72-
1. **Collect Statistics**: Run sample queries (typically 4,000) with K=1 and collect:
73-
- Distance to entry point
74-
- Distances to nearest neighbors found so far
75-
- Number of graph hops and distance comparisons
76-
- Traversal window statistics
77-
78-
2. **Train Model**: Use LightGBM to train a binary classifier:
79-
- Input: 12 trajectory-based features
80-
- Output: Probability that top-1 result has been found
81-
- Model is trained only for K=1
82-
83-
3. **Generate Lookup Tables**:
84-
- **Threshold table**: Maps model confidence scores to recall levels
85-
- **Interval table**: Maps recall targets to prediction intervals
86-
87-
### Inference Phase
88-
89-
1. **Start Search**: Begin traversing the HNSW graph
90-
2. **For each rank k from 1 to K**:
91-
- Collect trajectory features
92-
- Use K=1 model to predict if top-1 is found
93-
- Stop when predicted recall target is reached
94-
- Mask the found result
95-
- Continue for next rank
96-
3. **Return top-K results**
97-
98-
---
99-
100-
## Prerequisites and Setup
101-
102-
### System Requirements
103-
104-
- **Operating System**: Linux
105-
- **Conda**: Miniconda or Anaconda
106-
- **C/C++ Compiler**: GCC/G++ with C++17 support
107-
- **Sufficient disk space** for datasets
108-
109-
### Installation
110-
111-
#### 1) Create and Activate the Conda Environment
112-
113-
Create the environment from `environment.yml`:
9+
**Requirements**: Linux, Conda, GCC with C++17
11410

11511
```bash
12+
# Setup environment
11613
conda env create -f environment.yml
11714
conda activate OMEGA
118-
```
119-
120-
#### 2) Build Dependencies (HNSWLIB + OMEGA)
121-
122-
Build required libraries and bindings:
12315

124-
```bash
16+
# Build
12517
bash ./scripts/build_hnswlib.sh
126-
```
127-
128-
#### 3) Download and Prepare Datasets
129-
130-
Download and preprocess a dataset:
13118

132-
```bash
19+
# Download dataset
13320
python create_dataset.py --dataset deep-100M
13421
```
13522

136-
Replace `deep-100M` with other supported dataset names (e.g., `bigann-100M`, `gist-1M`).
137-
138-
---
139-
140-
## How to Run OMEGA
141-
142-
### Runbooks
143-
144-
Pre-generated runbooks are located at `neurips23/runbooks/`. You can also generate runbooks manually:
145-
146-
```bash
147-
python bench_runbook_gen2.py
148-
```
149-
150-
### End-to-End Evaluation
151-
152-
Run the complete end-to-end evaluation:
23+
## Running Experiments
15324

15425
```bash
26+
# Run full evaluation (OMEGA vs LAET/DARTH/Fixed-ef baselines)
15527
bash ./scripts/eval_end2end.sh
156-
```
157-
158-
This script evaluates **LAET**, **DARTH**, **OMEGA**, and **Fixed efSearch** baseline on the configured dataset/runbook(s), and produces metrics such as latency and recall.
159-
160-
After `./scripts/eval_end2end.sh` finishes, you will see a per-dataset JSON result under `./plot/eval_end2end_results/` (e.g., `./plot/eval_end2end_results/deep-100M.json`). This file contains all metrics for all tested configurations/modes.
161-
162-
### Understanding the Results
163-
164-
#### What metrics are in the result JSON?
165-
166-
The top-level JSON typically includes `microbench_results` and `end2end_results`. Each entry in `end2end_results` corresponds to one evaluated configuration and includes:
167-
168-
- **recall**: Overall recall achieved under this configuration
169-
- **recall_target**: The target recall (if applicable) for this run
170-
- **recall_cdf**: The CDF of per-query recall (useful to see tail behavior / worst-case recall)
171-
- **recall_groupby_count**: Recall broken down by count buckets
172-
- **mean_latency**: Mean end-to-end query latency (typically in microseconds)
173-
- **latency_cdf**: Latency CDF (for tail latency analysis)
174-
- **latency_groupby_count**: Latency broken down by count buckets
175-
- **mean_cmps**: Average number of comparisons per query
176-
- **thpt**: Throughput (queries/second) for the run
177-
- **prediction_times**: Time spent on model prediction
178-
- **prediction_overhead**: Overhead attributed to prediction (if enabled)
179-
- **run_latency / generate_latency / train_latency / total_preparation_latency**: A breakdown of preparation/training-related overheads
18028

181-
It also records run context fields such as `threads`, `num_queries`, `mode`, `train_counts`, `train_queries`, and `traversal_window_size`.
182-
183-
### Plotting Results
184-
185-
If you run the default `./scripts/eval_end2end.sh` to completion and then run:
186-
187-
```bash
29+
# Plot results
18830
python ./plot/eval_end2end.py
18931
```
19032

191-
you will obtain visualizations showing a comparison of recall and search latency with different preprocessing time budgets.
192-
193-
---
194-
195-
## Comparing with Baselines
196-
197-
OMEGA is compared against:
198-
- **Normal mode**: Fixed ef parameter (baseline)
199-
- **LAET**: Learning-Augmented Early Termination
200-
- **DARTH**: Dynamic Adaptive Recall Threshold Heuristic
201-
- **Optimal**: Oracle with perfect knowledge (upper bound)
202-
203-
---
204-
205-
## References and Credits
206-
207-
### Related Work
208-
209-
We have ported the following learned ANN search systems and compared OMEGA against them:
210-
211-
- **LAET (Learned Adaptive Early Termination)**
212-
Conglong Li, Minjia Zhang, David G. Andersen, Yuxiong He.
213-
*"Improving Approximate Nearest Neighbor Search through Learned Adaptive Early Termination"*
214-
SIGMOD Conference 2020, pages 2539-2554.
215-
DOI: [10.1145/3318464.3380600](https://doi.org/10.1145/3318464.3380600)
216-
217-
- **DARTH (Declarative Recall Through Early Termination)**
218-
Manos Chatzakis, Yannis Papakonstantinou, Themis Palpanas.
219-
*"DARTH: Declarative Recall Through Early Termination for Approximate Nearest Neighbor Search"*
220-
Proceedings of the ACM on Management of Data (PACMMOD), 3(4), Article 242, 2025.
221-
DOI: [10.1145/3749160](https://doi.org/10.1145/3749160)
222-
223-
### Original Projects
224-
225-
This project builds upon:
226-
227-
- **Big ANN Benchmarks**: See [README_Big_ANN_Benchmarks.md](./README_Big_ANN_Benchmarks.md) for the original benchmark framework
228-
- Original project: [ann-benchmarks](https://github.com/erikbern/ann-benchmarks) by Erik Bernhardsson
229-
- NeurIPS competitions: [big-ann-benchmarks.com](http://big-ann-benchmarks.com/)
230-
231-
- **hnswlib**: Fast approximate nearest neighbor search library
232-
- GitHub: [nmslib/hnswlib](https://github.com/nmslib/hnswlib)
33+
Results are saved to `./plot/eval_end2end_results/` as JSON files with recall, latency, and throughput metrics.
23334

234-
- **LightGBM**: Gradient boosting framework
235-
- GitHub: [microsoft/LightGBM](https://github.com/microsoft/LightGBM)
35+
## References
23636

237-
### License
37+
Built on [big-ann-benchmarks](./README_Big_ANN_Benchmarks.md), [hnswlib](https://github.com/nmslib/hnswlib), and [LightGBM](https://github.com/microsoft/LightGBM).
23838

239-
This project respects the licenses of all incorporated components. See individual project directories for specific license information.
39+
Baselines implemented: [LAET](https://doi.org/10.1145/3318464.3380600) (SIGMOD'20), [DARTH](https://doi.org/10.1145/3749160) (PACMMOD'25).

‎scripts/eval_end2end.sh‎

Lines changed: 10 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -2,28 +2,22 @@
22

33
algorithm=hnswlib
44

5-
datasets=(bigann-100M deep-100M gist-1M)
6-
runbook_list=(bench_runbook_100M_0M.json bench_runbook_100M_0M.json bench_runbook_1M_0M.json)
7-
counts_path_list=(traces/bj.json traces/bj.json traces/bj.json)
5+
datasets=(deep-100M)
6+
runbook_list=(bench_runbook_100M_0M.json)
7+
counts_path_list=(traces/bj.json)
88

9-
input_list=(input_data_step101.npy input_data_step101.npy input_data_step2.npy)
10-
output_list=(output_data_step101.npy output_data_step101.npy output_data_step2.npy)
11-
gt_cmps_list=(gt_cmps_step101.npy gt_cmps_step101.npy gt_cmps_step2.npy)
12-
latency_list=(latency_step101.json latency_step101.json latency_step2.json)
9+
input_list=(input_data_step101.npy)
10+
output_list=(output_data_step101.npy)
11+
gt_cmps_list=(gt_cmps_step101.npy)
12+
latency_list=(latency_step101.json)
1313

14-
gt_collected_list=(gt_collected_step101.npy gt_collected_step101.npy gt_collected_step2.npy)
15-
gt_cmps_all_list=(gt_cmps_all_step101.npy gt_cmps_all_step101.npy gt_cmps_all_step2.npy)
16-
# lower_bound_gt_collected_list=(lower_bound_gt_collected_step2.npy)
17-
# upper_bound_gt_collected_list=(upper_bound_gt_collected_step2.npy)
14+
gt_collected_list=(gt_collected_step101.npy)
15+
gt_cmps_all_list=(gt_cmps_all_step101.npy)
1816

1917
Ks_end2end_train_opt=(1)
20-
Ks_end2end_train_bigann_100M=(1 10 50 100 200)
2118
Ks_end2end_train_deep_100M=(1 10 50 100 200)
22-
Ks_end2end_train_gist_1M=(1 10 50 100 200)
2319

24-
maxK_end2end_bigann_100M=200
2520
maxK_end2end_deep_100M=200
26-
maxK_end2end_gist_1M=200
2721

2822
results_base_path=eval_end2end_results
2923
num_train_queries=4000
@@ -49,11 +43,8 @@ for i in "${!datasets[@]}"; do
4943
output=${output_list[$i]}
5044
gt_cmps=${gt_cmps_list[$i]}
5145
latency=${latency_list[$i]}
52-
5346
gt_collected=${gt_collected_list[$i]}
5447
gt_cmps_all=${gt_cmps_all_list[$i]}
55-
# lower_bound_gt_collected=${lower_bound_gt_collected_list[$i]}
56-
# upper_bound_gt_collected=${upper_bound_gt_collected_list[$i]}
5748

5849
dataset_normalized=${dataset//-/_}
5950
declare -n Ks_end2end_train="Ks_end2end_train_${dataset_normalized}"
@@ -73,52 +64,12 @@ for i in "${!datasets[@]}"; do
7364
--training \
7465
--num_queries $num_train_queries
7566

76-
# # generate query sequence
67+
# generate query sequence
7768
python benchmark/streaming/generate_query_sequence.py \
7869
--dataset $dataset \
7970
--counts_path $counts_path \
8071
--recall_target 0.95
8172

82-
# # train_rem & rem
83-
# train_mode=train_rem_normal
84-
# mode=rem_normal
85-
86-
# python run.py \
87-
# --neurips23track streaming \
88-
# --algorithm $algorithm \
89-
# --dataset $dataset \
90-
# --runbook_path neurips23/runbooks/$runbook \
91-
# --results_base_path $results_base_path \
92-
# --mode $train_mode \
93-
# --num_queries $num_train_queries \
94-
# --counts ${Ks_end2end_train[@]} \
95-
# --nodocker
96-
# for K in "${Ks_end2end_train[@]}"; do
97-
# python train_rem.py \
98-
# --neurips23track streaming \
99-
# --dataset $dataset \
100-
# --runbook_path neurips23/runbooks/$runbook \
101-
# --results_base_path $results_base_path \
102-
# --mode $train_mode \
103-
# --num_queries $num_train_queries \
104-
# --model_base_path models/$dataset/$runbook/$train_mode/$K/ \
105-
# --count $K
106-
# # remove training data
107-
# rm $results_base_path/neurips23/streaming/$runbook/$dataset/$K/$algorithm/$train_mode*.hdf5
108-
# done
109-
110-
# model_base_paths=("models/$dataset/$runbook/$train_mode/")
111-
# python run.py \
112-
# --neurips23track streaming \
113-
# --algorithm $algorithm \
114-
# --dataset $dataset \
115-
# --runbook_path neurips23/runbooks/$runbook \
116-
# --results_base_path $results_base_path \
117-
# --mode $mode \
118-
# --end2end \
119-
# --nodocker \
120-
# --model_base_paths ${model_base_paths[@]}
121-
12273
# --------------------------------------------- LAET ---------------------------------------------
12374
# train & early_stop
12475
train_mode=train_normal

0 commit comments

Comments
 (0)