基于指示矩阵的高效受限矩阵乘法实现及存储优化技术问询

Great question! Let's tackle your two problems one by one, with practical code and explanations to get you the efficiency and memory savings you need.

1. Efficiently Implement the Sparse Dot Product #

Your current loop-based approach is slow because it's executing Python-level iterations instead of leveraging PyTorch's optimized vectorized operations and hardware acceleration (GPU or optimized CPU kernels). Here's a much faster way to compute the sparse dot product:

Optimized Code #

import torch
import time
import random
from sklearn.neighbors import kneighbors_graph

n = 1024
d = 512
knn = 64

q = torch.randn(n, d)
k = torch.randn(d, n)

coords = torch.tensor([(random.random(), random.random()) for _ in range(n)])
# Get sparse connectivity matrix directly (avoid dense conversion)
sparse_A = kneighbors_graph(coords, knn, mode='connectivity', include_self=False)

# Convert sparse matrix to a (n, knn) index tensor of relevant columns
indices = torch.tensor(sparse_A.indices).t()  # Shape: (n*knn, 2) - (row, col) pairs
row_ids, col_ids = indices[:, 0], indices[:, 1]
# Group columns by their row to get per-row knn indices
sorted_col_indices = col_ids[row_ids.argsort()].reshape(n, knn)

# Fast sparse dot product using batch operations
t1_start = time.time()
# Extract relevant columns of K and reshape for batch matmul
k_selected = k[:, sorted_col_indices].transpose(0, 1)  # Shape: (n, d, knn)
# Batch matrix multiplication: (n, 1, d) @ (n, d, knn) → (n, 1, knn)
outs = torch.bmm(q.unsqueeze(1), k_selected).squeeze(1)
t1_stop = time.time()
print(f"Sparse optimized time: {t1_stop - t1_start:.6f}")

# Compare with full dot product
t1_start = time.time()
full_outs = q.matmul(k)
t1_stop = time.time()
print(f"Full dot product time: {t1_stop - t1_start:.6f}")

Why This Works #

• We eliminate Python-level loops entirely, letting PyTorch handle the computation in optimized C++/CUDA code.
• For large n (e.g., n=10,000), this approach will outperform the full dot product dramatically—since it only computes n*knn operations instead of n². On GPU, the speedup will be even more significant.

2. Memory-Optimal Storage for Matrix A #

Storing A as a dense boolean matrix is extremely inefficient for large n—it uses n² bits/bytes, which becomes prohibitive when n scales to tens of thousands. Instead, use sparse formats or direct index storage:

Option 1: PyTorch Sparse Tensors #

PyTorch's native sparse tensors only store coordinates of non-zero elements, cutting memory usage to O(n*knn):

# Create a sparse COO tensor from scikit-learn's sparse matrix
values = torch.ones(sparse_A.nnz(), dtype=torch.bool)
sparse_A_torch = torch.sparse_coo_tensor(
    torch.tensor(sparse_A.indices),
    values,
    size=(n, n),
    dtype=torch.bool
)

# To retrieve per-row indices later:
row_indices, col_indices = sparse_A_torch.indices()

Option 2: Store Per-Row Indices Directly #

If you only need A to filter columns for the dot product, skip storing the matrix entirely and keep just the (n, knn) index tensor we created earlier:

# This tensor holds exactly the knn column indices for each row of Q
sorted_col_indices = col_ids[row_ids.argsort()].reshape(n, knn)

# For extra savings, use smaller dtypes if possible:
sorted_col_indices = sorted_col_indices.to(torch.int16)  # Works if n < 65536

This is the most memory-efficient option—you only store the data you actually need, no extra overhead from sparse matrix metadata.

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