Как быстро отранжировать вектора по близости к запросу?

Как быстро отранжировать вектора по близости к запросу?

Ранжирование векторов по близости к запросу — фундаментальная задача в информационных поиске, рекомендациях и semanticсearch. Существует несколько подходов с разными trade-offs между скоростью и точностью.

1. Быстрые приблизительные методы (ANN)

Approximate Nearest Neighbors (ANN) — основной инструмент для быстрого поиска в миллионах векторов:

FAISS (Facebook AI Similarity Search)

import numpy as np
import faiss
from time import time

# Генерируем датасет из 1 миллиона 768-мерных векторов (как в BERT)
n_samples = 1_000_000
dim = 768

# Случайные векторы
db_vectors = np.random.random((n_samples, dim)).astype(float32)
query = np.random.random((1, dim)).astype(float32)

# Способ 1: Точный поиск (brute force) — медленно
print("===== Exact Search =====")
index = faiss.IndexFlatL2(dim)  # L2 расстояние
index.add(db_vectors)

start = time()
distances, indices = index.search(query, k=10)  # ищем топ-10
print(f"Time: {time() - start:.4f}s")
print(f"Top 10 indices: {indices[0]}")

# Способ 2: Быстрый приблизительный поиск с HNSW
print("\n===== HNSW (Hierarchical Navigable Small World) =====")
index = faiss.IndexHNSWFlat(dim, 32)  # 32 соседей для навигации
index.add(db_vectors)
index.hnsw.efConstruction = 200  # параметр построения

start = time()
distances, indices = index.search(query, k=10)
print(f"Time: {time() - start:.4f}s")
print(f"Top 10 indices: {indices[0]}")

# Способ 3: Быстрый поиск с квантизацией (экономит память в 64x раз)
print("\n===== Product Quantization =====")
index = faiss.IndexPQ(dim, 64, 8)  # 64 подпространства, 8 бит
index.train(db_vectors)  # требует обучения
index.add(db_vectors)

start = time()
distances, indices = index.search(query, k=10)
print(f"Time: {time() - start:.4f}s (с квантизацией)")
print(f"Top 10 indices: {indices[0]}")

# Способ 4: Комбинация IVF + PQ (лучшее по скорости/качеству)
print("\n===== IVF + PQ =====")
quantizer = faiss.IndexFlatL2(dim)  # базовый индекс для квантизатора
index = faiss.IndexIVFPQ(quantizer, dim, 100, 64, 8)  # 100 кластеров, PQ с 64x8
index.train(db_vectors)
index.add(db_vectors)
index.nprobe = 10  # количество кластеров для поиска

start = time()
distances, indices = index.search(query, k=10)
print(f"Time: {time() - start:.4f}s")
print(f"Top 10 indices: {indices[0]}")

Методы FAISS — сравнение

2. Использование специализированных БД

Milvus (Open Source Vector Database)

from pymilvus import Collection, connections

# Подключаемся к Milvus
connections.connect("default", host="localhost", port=19530)

# Создаём коллекцию с индексом HNSW
from pymilvus import FieldSchema, CollectionSchema

fields = [
    FieldSchema(name="id", dtype="int64", is_primary=True),
    FieldSchema(name="vector", dtype="float_vector", dim=768)
]
schema = CollectionSchema(fields, "Document collection")
collection = Collection("documents", schema)

# Добавляем векторы
import numpy as np
data = {
    "id": list(range(1_000_000)),
    "vector": np.random.random((1_000_000, 768)).astype(float32).tolist()
}
collection.insert(data)

# Создаём индекс
index_params = {
    "metric_type": "L2",
    "index_type": "HNSW",
    "params": {"M": 30, "efConstruction": 200}
}
collection.create_index("vector", index_params)

# Быстрый поиск
query_vector = np.random.random((1, 768)).astype(float32)
results = collection.search(
    query_vector,
    "vector",
    {"metric_type": "L2", "params": {"ef": 64}},
    limit=10
)
for result in results:
    print(f"ID: {result.id}, Distance: {result.distance}")

Pinecone (Managed Vector DB)

import pinecone

# Инициализируем Pinecone
pinecone.init(api_key="YOUR_API_KEY", environment="us-west1-gcp")

# Создаём индекс
pinecone.create_index(
    "documents",
    dimension=768,
    metric="cosine",
    pod_type="p1.x1"  # облачный сервер
)

# Загружаем векторы (батчами для скорости)
index = pinecone.Index("documents")
for i in range(0, 1_000_000, 100):
    vectors = [
        (str(j), np.random.random(768).tolist(), {"text": f"doc_{j}"})
        for j in range(i, min(i+100, 1_000_000))
    ]
    index.upsert(vectors=vectors)

# Быстрый поиск
query_vector = np.random.random(768).tolist()
results = index.query(query_vector, top_k=10, include_metadata=True)
for match in results.matches:
    print(f"ID: {match.id}, Score: {match.score}, Text: {match.metadata[text]}")

3. Оптимизированные метрики расстояния

import numpy as np
from scipy.spatial.distance import cdist
import time

n_vectors = 100_000
dim = 768
vectors = np.random.random((n_vectors, dim)).astype(float32)
query = np.random.random((1, dim)).astype(float32)

# L2 расстояние (Euclidean)
start = time.time()
distances_l2 = np.linalg.norm(vectors - query, axis=1)
top_10_l2 = np.argsort(distances_l2)[:10]
print(f"L2 distance: {time.time() - start:.4f}s")

# Cosine сходство (нормализованные векторы — быстрее)
start = time.time()
vectors_norm = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
query_norm = query / np.linalg.norm(query)
similarities = np.dot(vectors_norm, query_norm.T).flatten()
top_10_cos = np.argsort(-similarities)[:10]  # минус для макса
print(f"Cosine distance: {time.time() - start:.4f}s")

# Inner product (если векторы уже нормализованы — самый быстрый)
start = time.time()
scores = np.dot(vectors_norm, query_norm.T).flatten()
top_10_ip = np.argsort(-scores)[:10]
print(f"Inner product: {time.time() - start:.4f}s")

4. Практический пример с embedding моделью

from sentence_transformers import SentenceTransformer
import faiss

# Загружаем модель
model = SentenceTransformer(all-MiniLM-L6-v2)

# Документы
documents = [
    "XGBoost это градиентный бустинг для классификации",
    "Random Forest использует множество деревьев",
    "Нейросети используют слои и активации",
    "Кластеризация находит группы схожих данных",
    # ... миллионы документов
]

# Кодируем документы в векторы
print("Encoding documents...")
embeddings = model.encode(documents, batch_size=128, show_progress_bar=True)
embeddings = embeddings.astype(float32)

# Создаём FAISS индекс
dim = embeddings.shape[1]
index = faiss.IndexIVFPQ(faiss.IndexFlatL2(dim), dim, 100, 64, 8)
index.train(embeddings)
index.add(embeddings)
index.nprobe = 10

# Запрос
query = "Как работает XGBoost?"
query_embedding = model.encode([query], convert_to_tensor=False).astype(float32)

# Поиск
print("\nSearching...")
import time
start = time.time()
distances, indices = index.search(query_embedding, k=10)
print(f"Time: {time.time() - start:.4f}s")

for i, idx in enumerate(indices[0]):
    print(f"{i+1}. {documents[idx]} (distance: {distances[0][i]:.4f})")

5. Рекомендации по выбору метода

< 100K векторов, требуется точность 100%:

faiss.IndexFlatL2(dim)  # Brute force

100K - 10M векторов, требуется баланс скорости/точности:

faiss.IndexHNSWFlat(dim, 32)  # HNSW

10M+ векторов, критична скорость:

index = faiss.IndexIVFPQ(faiss.IndexFlatL2(dim), dim, 1000, 64, 8)

Production, высокая нагрузка:

# Используйте управляемый сервис: Pinecone, Weaviate, Qdrant

Ключевые выводы

1. Точный поиск (brute force) работает только для маленьких датасетов
2. HNSW — лучший выбор для баланса скорости и точности
3. IVF + PQ — оптимален для больших датасетов (миллионы векторов)
4. Специализированные БД (Milvus, Pinecone) — для production систем
5. Cosine сходство быстрее L2 для нормализованных векторов
6. Батчевая обработка критична для скорости
7. Параметр k в поиске влияет на скорость больше, чем размер датасета