如何高效将Scipy稀疏CSR矩阵转换为PPMI加权矩阵？

针对大规模稀疏CSR矩阵（如65000×65000），逐元素计算PPMI会导致极低效率甚至内存溢出。我们可以利用稀疏矩阵的特性，通过批量操作完成转换，全程无需转为稠密矩阵。

PPMI公式转化 #

首先回顾PPMI公式：

PPMI(i, j) = max(log₂[P(i,j)/(P(i)×P(j))], 0)

其中：

• P(i,j) = 共现次数mat(i,j) / 总共现次数tot
• P(i) = 词i的总共现次数row_sum[i] / tot
• P(j) = 词j的总共现次数row_sum[j] / tot

代入后可简化公式：

PPMI(i,j) = max( log₂(mat(i,j)) + log₂(tot) - log₂(row_sum[i]) - log₂(row_sum[j]), 0 )

高效实现步骤 #

1. 计算总共现次数：所有非零元素的和（注意不是非零元素个数）
2. 计算每个词的总共现次数（矩阵每行的和）
3. 批量计算所有非零元素的PPMI值，过滤掉负值
4. 重新构建稀疏PPMI矩阵

完整代码实现 #

import numpy as np
import scipy.sparse as sp
from sklearn.feature_extraction.text import CountVectorizer

# 生成示例共现矩阵
sentences = ["The cat is on the table",
             "I have seen a cat in the office",
             "You shall feed the cat before it gets dark",
             "I have many pets in my house, but my favourite is my cat",
             "Dogs are nice, but cats are far nicer in my opinion"]
count_model = CountVectorizer(ngram_range=(1,1))
X = count_model.fit_transform(sentences)
Xc = (X.T * X)
Xc.setdiag(0)  # 移除对角线元素（自身共现）

# 高效计算PPMI
def csr_to_ppmi(matrix):
    tot = matrix.sum()
    row_sum = matrix.sum(axis=1).A1  # 获取每行的和，转为一维数组
    
    # 计算所需的对数项
    log_tot = np.log2(tot)
    log_row_sum = np.log2(row_sum)
    
    # 获取所有非零元素的位置和值
    i, j = matrix.nonzero()
    data = matrix.data
    
    # 计算每个非零元素的PPMI中间值
    log_vals = np.log2(data) + log_tot - log_row_sum[i] - log_row_sum[j]
    
    # 过滤负值，保留非零的PPMI值
    ppmi_data = np.maximum(log_vals, 0)
    # 只保留PPMI值大于0的元素（进一步压缩稀疏矩阵）
    mask = ppmi_data > 0
    i_filtered = i[mask]
    j_filtered = j[mask]
    data_filtered = ppmi_data[mask]
    
    # 构建新的CSR格式PPMI矩阵
    ppmi_matrix = sp.csr_matrix((data_filtered, (i_filtered, j_filtered)), shape=matrix.shape)
    return ppmi_matrix

# 转换为PPMI矩阵
ppmi_mat = csr_to_ppmi(Xc)

效率说明 #

• 全程基于稀疏矩阵的非零元素操作，无需加载整个稠密矩阵，内存占用仅与非零元素数量相关
• 所有计算均为向量/数组级别的批量运算，避免了逐元素循环的低效
• 最终输出仍为CSR稀疏矩阵，保持了后续处理的高效性

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