Function: The SparselightningIndexerGradKlLoss operator is the backward operator of LightningIndexer and integrates the loss calculation function. The LightningIndexer operator selects the top K tokens with the highest intrinsic relationship between the query token and key token, and stores them in SparseIndices. This reduces the amount of attention calculation in long sequence scenarios and accelerates the inference and training performance of the network with long sequences.
Formula: The formula for calculating the top-k value can be expressed as follows:
In the formula, is the weights corresponding to the th token, is the matrix obtained by combining query heads corresponding to the th token, and is the matrix of the th row.
LightningIndexer is trained separately, and the corresponding loss function is as follows:
In the formula, is the target distribution, which is obtained by summing up all the heads of the main attention score and then performing L1 regularization on the sum result in the context direction. is the KL divergence, and its expression is as follows:
The gradient expression of the loss can be obtained by derivation as follows:
The chain rule can be used to calculate the gradients of the weights, query, and key matrices.
S is the softmax result of the QK matrix.
Description
The data format of [object Object], [object Object], and [object Object] can be interpreted from multiple dimensions. To be specific, B (Batch) indicates the size of an input sample batch, S (Seq-Length) indicates the length of the input sample sequence, H (Head-Size) indicates the size of the hidden layer, N (Head-Num) indicates the number of heads, and D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N). T indicates the total length of all input sample sequences.
-->
Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.
Parameters:
[object Object]Returns:
[object Object]: status code. For details, see .The first-phase API implements input parameter verification. The following errors may be thrown.
[object Object]
- Deterministic computation:
- By default, aclnnSparseLightningIndexerGradKLLoss is implemented in non-deterministic mode. Deterministic computing cannot be enabled using aclrtCtxSetSysParamOpt.
- Common Constraints
- Handling of the scenario where the input parameter is empty:
- If
[object Object]is an empty tensor, the result is returned directly. - The scenario where the input parameter is empty in the general restrictions is the same as that in the FAG.
- If
- Handling of the scenario where the input parameter is empty:
- Specification Restrictions[object Object]
- Typ.[object Object]
The following example is for reference only. For details, see .