как векторизовать операцию рассеяния

Question

У меня много матриц w1, w2, w3...wn с фигурами (k*n1, k*n2, k*n3...k*nn) и x1, x2, x3...xn с фигурами (n1*m, n2*m, n3*m...nn*m).
Я хочу получить w1@x1, w2@x2, w3@x3 ... соответственно.

Полученная матрица состоит из нескольких k*m матриц и может быть объединен в большую матрицу с формой (k*n)*m.

Умножать их по одному будет медленно. Как векторизовать эту операцию?

Примечание. Входными данными могут быть матрица k*(n1+n2+n3+...+nn) и матрица (n1+n2+n3+...+nn)*m, и мы можем использовать индекс пакета, чтобы указать эти подматрицы.

Эта операция относится к операциям рассеяния, реализованным в pytorch_scatter, поэтому я называю это "scatter_matmul".

Answer 1

Вы можете векторизовать свою операцию, создав большую матрицу с диагональю блоков W формы n*k x (n1+..+nn), где матрицы w_i являются блоками на диагонали. Затем вы можете вертикально сложить все x матрицы в X матрицу формы (n1+..+nn) x m. Умножение диагонали блока W с вертикальным стеком всех x матриц, X:

Y = W @ X

приводит к Y формы (k*n) x m, которая является точно сцепленной большая матрица, которую вы ищете.

Если форма диагональной матрицы блока W слишком велика, чтобы уместиться в память, вы можете рассмотреть вопрос о том, чтобы сделать W разреженным и вычислить произведение, используя torch.sparse.mm.