优化词嵌入邻接矩阵构建的内存与时间效率问询

这个问题我太有共鸣了——处理10万级高维向量的邻接矩阵时，pdist的O(n²)内存开销简直是噩梦，手动迭代又慢到让人崩溃。咱们直接上最实用的解决方案，分优先级推荐：

1. 首选：近似最近邻(ANN)库（FAISS）——速度内存双优 #

核心思路是只计算我们需要的邻居关系，而非所有两两距离，这直接规避了O(n²)的内存爆炸。FAISS是Facebook开源的高性能向量检索库，用C++实现，速度拉满，内存占用极低。

欧氏距离实现代码 #

import faiss
import numpy as np
from scipy.sparse import csr_matrix

# 第一步：把float64转成float32，直接省一半内存（GloVe精度完全够用）
common_model = common_model.astype(np.float32)
n, d = common_model.shape
distance_threshold = 0.4

# 构建FAISS的L2距离索引（欧氏距离的平方，所以后面阈值要平方）
index = faiss.IndexFlatL2(d)
index.add(common_model)

# 范围搜索：找出每个向量距离<=阈值的所有邻居
# 返回的indptr是每个向量的邻居起始位置，indices是邻居的ID
_, indptr, indices = index.range_search(common_model, distance_threshold ** 2)

# 构建CSR邻接矩阵
data = np.ones_like(indices, dtype=np.uint8)
adjacency = csr_matrix((data, indices, indptr), shape=(n, n), dtype=np.uint8)

余弦距离适配 #

如果要使用余弦距离，先把向量归一化，然后用内积索引（归一化后内积等价于余弦相似度，距离=2*(1-内积)）：

# 归一化向量，L2范数为1
common_model = common_model / np.linalg.norm(common_model, axis=1, keepdims=True)
cos_threshold = 0.4
# 余弦距离<=0.4 等价于 内积>=1 - 0.4/2 = 0.8
inner_product_threshold = 1 - cos_threshold / 2

index = faiss.IndexFlatIP(d)
index.add(common_model)
_, indptr, indices = index.range_search(common_model, inner_product_threshold)

# 同样构建CSR矩阵
data = np.ones_like(indices, dtype=np.uint8)
adjacency = csr_matrix((data, indices, indptr), shape=(n, n), dtype=np.uint8)

为什么选FAISS？

• 内存：只需要存储最终的稀疏邻接矩阵，内存占用是O(E)（E是边的数量），远小于O(n²)
• 速度：C++底层优化，比numpy向量化还快，10万级向量几分钟就能搞定

2. 次选：分块向量化处理（无第三方库依赖） #

如果不能安装FAISS，用numpy+scipy的分块方法也能大幅降低内存占用，同时比手动迭代快N倍。核心是把大矩阵切成小块，每次计算一块与整个矩阵的距离，避免一次性生成O(n²)的距离矩阵。

import numpy as np
from scipy.sparse import csr_matrix, vstack

common_model = common_model.astype(np.float32)
n, d = common_model.shape
distance_threshold = 0.4
block_size = 1000  # 可根据你的内存调整，比如2000或5000

sparse_blocks = []
for start_idx in range(0, n, block_size):
    # 取出当前块的向量
    block_vecs = common_model[start_idx:start_idx+block_size]
    block_len = len(block_vecs)
    
    # 计算当前块与所有向量的欧氏距离平方（不开根号，省计算量）
    dist_sq = np.sum((block_vecs[:, np.newaxis] - common_model)**2, axis=-1)
    # 筛选距离<=阈值的位置
    mask = dist_sq <= distance_threshold ** 2
    
    # 构建当前块的稀疏矩阵
    row_ids = np.repeat(np.arange(block_len), mask.sum(axis=1))
    col_ids = np.tile(np.arange(n), block_len)[mask.ravel()]
    block_data = np.ones(len(row_ids), dtype=np.uint8)
    block_csr = csr_matrix((block_data, (row_ids, col_ids)), shape=(block_len, n))
    sparse_blocks.append(block_csr)

# 拼接所有块得到最终邻接矩阵
adjacency = vstack(sparse_blocks)

优势：无需额外安装库，内存占用可控（每次只需要O(block_size * n)的内存），numpy向量化操作比手动循环快数十倍。

关键优化小技巧 #

• 永远把float64转成float32：GloVe预训练嵌入的精度用float32完全足够，直接省一半内存
• 用距离平方代替欧氏距离：避免开根号的计算开销，比较时用阈值的平方即可
• 优先用稀疏矩阵存储邻接关系：0-1矩阵绝大多数是0，稀疏矩阵能把内存占用降到最低

内容的提问来源于stack exchange，提问作者neurotroph