mhc_pre_sinkhorn

Applicable Products

Ascend 950PR/Ascend 950DT: Not supported
Atlas A3 training products/Atlas A3 inference products: Supported
Atlas A2 training products/Atlas A2 inference products: Supported
Atlas 200I/500 A2 inference products: Not supported
Atlas inference products: Not supported
Atlas training products: Not supported

Function

Description:

Based on a series of computations, obtains the $\mathbf{H}'_{\text{res}}$ and $\mathbf{H}_{\text{post}}$ projection matrices of the hidden layer in the MHC architecture and the input matrix $\mathbf{h}_{\text{in}}$ of the Attention or MLP layer. Performs Sinkhorn iterative normalization on the $\mathbf{H}'_{\text{res}}$ matrix to obtain the doubly stochastic matrix $\mathbf{H}_{\text{res}}$ . Supports outputting intermediate computation results for backward gradient computation.
Formula:
$\begin{aligned} \vec{x^{'}_{l}} &= \frac{1}{\sqrt{\frac{1}{d} \sum_{\dim=-2,\text{keepdim}=\text{True}} x_i^2 + \epsilon_{norm}}}\\ H^{pre}_l &= \alpha^{pre}_{l} \cdot(\vec{x^{'}_{l}}\varphi^{pre}_{l}) + b^{pre}_{l}\\ H^{post}_l &= \alpha^{post}_{l} \cdot(\vec{x^{'}_{l}}\varphi^{post}_{l}) + b^{post}_{l}\\ H^{res}_l &= \alpha^{res}_{l} \cdot(\vec{x^{'}_{l}}\varphi^{res}_{l}) + b^{res}_{l}\\ H^{pre}_l &= \sigma (H^{pre}_{l})\\ H^{post}_l &= 2\sigma (H^{post}_{l})\\ h_{in} &=\vec{x_{l}}H^{pre}_l \end{aligned}$
With $\mathbf{H^{res}_l}$ as the input, the Sinkhorn transform performs $\mathbf{numIters}$ iterations. During the iterations, intermediate normalization results $\mathbf{normOut}[k]$ and summation results $\mathbf{sumOut}[k]$ are generated. The final output is $\mathbf{normOut}$ from the last iteration as the transform result.

First iteration (initialization):
$\begin{aligned} \mathbf{normOut}[0] &= \text{softmax}(\mathbf{H^{res}_l}, \dim=-1) + \epsilon_{hc}, \\ \mathbf{sumOut}[1] &= \sum_{\dim=-2,\text{keepdim}=\text{True}} \mathbf{normOut}[0] + \epsilon_{hc}, \\ \mathbf{normOut}[1] &= \frac{\mathbf{normOut}[0]}{\mathbf{sum\_out}[1]}, \\ \end{aligned}$
The $i$ -th iteration ( $i = 1, 2, \dots, \mathbf{num\_iters}-1$ ):
$\begin{aligned} \mathbf{sumOut}[2i] &= \sum_{\dim=-1,\text{keepdim}=\text{True}} \mathbf{normOut}[2i-1] + \epsilon_{hc}, \\ \mathbf{normOut}[2i] &= \frac{\mathbf{normOut}[2i-1]}{\mathbf{sum\_out}[2i]}, \\ \mathbf{sumOut}[2i+1] &= \sum_{\dim=-2,\text{keepdim}=\text{True}} \mathbf{normOut}[2i] + \epsilon_{hc}, \\ \mathbf{normOut}[2i+1] &= \frac{\mathbf{normOut}[2i]}{\mathbf{sum\_out}[2i+1]}, \\ \end{aligned}$
Final output:
$\mathbf{normOut}[2 \times \mathbf{num\_iters} - 1]$ $\mathbf{sumOut}[2 \times \mathbf{num\_iters} - 1]$

Prototype

mhc_pre_sinkhorn(x, phi, alpha, bias, hcMult, numIters, hcEps, normEps, outFlag) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)

Parameters

Parameter	Parameter Type	Optional/Required	Description	Data Type	Dimension (shape)
x	Tensor	Required	MHC layer input data. Corresponds to x in the formula.	bfloat16	(B, S, N, C)
phi	Tensor	Required	MHC parameter matrix. Corresponds to ψ^pre, ψ^post, and ψ^res in the formula. Split along the N axis when used, into (N, NC), (N, NC), and (NN, NC).	float32	(NN+2N, N*C)
alpha	Tensor	Required	Scaling parameter of MHC. Contains three elements corresponding to the scaling coefficients of H^pre, H^post, and H^res. Corresponds to α^pre, α^post, and α^res in the corresponding formula.	float32	(3)
bias	Tensor	Required	Bias parameter of MHC. Contains 2N+N² elements. Corresponds to β^pre, β^post, and β^res in the corresponding formula, which correspond to N, N, and N*N respectively.	float32	(2N+N²)
hcMult	int	Required	Number of residual streams, that is, the HC dimension size.	-	-
numIters	int	Required	Number of Sinkhorn algorithm iterations.	-	-
hcEps	float	Optional	Eps parameter after sigmoid of H^pre. Corresponds to ε_hc in the formula. Defaults to 1e-6.	-	-
normEps	float	Optional	Anti-division-by-zero parameter of RmsNorm. Corresponds to ε_norm in the corresponding formula. Default value: 1e-6.	-	-

Output

Parameter	Parameter Type	Optional/Required	Description	Data Type	Dimension (shape)
hin	Tensor	Required	Output h_in, used as the input of the Attention/MLP layer. Corresponds to h_in in the formula.	bfloat16	(B, S, C)
hPost	Tensor	Required	Output h_post transformation matrix of MHC. Corresponds to H^post in the formula.	float32	(B, S, N)
hRes	Tensor	Required	Output MHC h_res transformation matrix (the doubly stochastic matrix after Sinkhorn normalization). Corresponds to H^res in the corresponding formula.	float32	(B, S, N*N)

Constraints

This API supports training and inference use cases.
This API supports single-operator mode calls.

Specification constraints:

Specification Item	Specification	Description
numIters	20	If the number of iterations exceeds this range, an invalid parameter error is returned.
N	4	Only 4 is supported.
C	[1280, 1920, 2560]	Only these three values are supported.

Parameter constraints:
- Empty tensor not supported for x.
- Empty tensor not supported for phi.
- Empty tensor not supported for alpha.
- Empty tensor not supported for bias.

Deterministic Computation

Defaults to a deterministic computation.

Example

Single-operator mode call:

import torch
import torch_npu
from cann_ops_transformer.ops import mhc_pre_sinkhorn

B = 1
S = 1024
N = 4
C = 1280

x = torch.randn(B, S, N, C, dtype=torch.bfloat16).npu()
phi = torch.randn(N * N + 2 * N, N * C, dtype=torch.float32).npu()
alpha = torch.randn(3, dtype=torch.float32).npu()
bias = torch.randn(N * N + 2 * N, dtype=torch.float32).npu()

hcMult = 4
numIters = 20
hcEps = 1e-6
normEps = 1e-6
outFlag = True  # Set to True when backward propagation is required.

hin, hPost, hRes = mhc_pre_sinkhorn(
  x, phi, alpha, bias, hcMult, numIters, hcEps, normEps, outFlag
)