Method and device for accelerating attention computation in transformer-based neural network

Abstract

This invention proposes a method and apparatus for accelerating attention computation in a neural network based on a transformation model. The method includes receiving an input attention map, generating a structured sparse attention map based on an average attention map, and encoding the structured sparse attention map into compressed data. The average attention map is pre-computed based on multiple training attention maps. The structured sparse attention map includes at least one of a column-oriented distribution pattern, a row-oriented distribution pattern, and a diagonal distribution pattern. The compressed data includes the positions and corresponding values ​​of retained items within the input attention map.

Claims

1. A method for accelerating attention computation in a neural network based on a transformation model, comprising: Received input attention map; Based on the average attention map, a structured sparse attention map is generated, wherein the average attention map is pre-computed based on multiple training attention maps, and the structured sparse attention map includes at least one of column distribution pattern, row distribution pattern and diagonal distribution pattern; and the structured sparse attention map is encoded into compressed data, wherein the compressed data includes the position and corresponding value of the retained items in the input attention map.

2. The method as described in claim 1, wherein prior to the step of receiving the input attention map, the method further comprises: During the training phase, for different training samples, the average attention map is generated by aggregating multiple training attention maps obtained from different attention heads of different layers of the neural network based on the transformation model.

3. The method as described in claim 1, wherein the step of generating the structured sparse attention map based on the average attention map includes: Based on the values ​​of the average attention map, filtering is performed on the average attention map to generate the structured sparse attention map.

4. The method of claim 3, wherein the step of performing filtering on the average attention map based on the value of the average attention map to generate the structured sparse attention map includes: For a specified column of the average attention map: determine the number of elements whose values ​​are greater than a first threshold; retain the specified column when the number of elements is greater than a second threshold; and filter out the specified column when the number of elements is not greater than the second threshold.

5. The method of claim 3, wherein the step of performing filtering on the average attention map based on the value of the average attention map to generate the structured sparse attention map includes: For a specified row of the average attention map: determine the number of elements whose values ​​are greater than a first threshold; When the number of elements is greater than the second threshold, the specified row is retained; and when the number of elements is not greater than the second threshold, the specified row is filtered out.

6. The method of claim 3, wherein the step of performing filtering on the average attention map based on the value of the average attention map to generate the structured sparse attention map includes: For a specified diagonal of the average attention map: determine the number of elements whose value is greater than a first threshold; when the number of elements is greater than a second threshold, retain the specified diagonal; and when the number of elements is not greater than the second threshold, filter out the specified diagonal.

7. The method as described in claim 1, wherein the structured sparse attention graph includes items belonging to at least one of a column-oriented distribution pattern, a row-oriented distribution pattern, and a diagonal distribution pattern, wherein items not belonging to the column-oriented distribution pattern, the row-oriented distribution pattern, and the diagonal distribution pattern are set to zero.

8. The method of claim 1, wherein the column distribution pattern represents a reserved position along at least one column of the input attention map, wherein the row distribution pattern represents a reserved position along at least one row of the input attention map, and wherein the diagonal distribution pattern represents a reserved position along at least one diagonal of the input attention map.

9. The method of claim 1, wherein after the step of encoding the structured sparse attention map into the compressed data, the method further comprises: The compressed data is decoded to obtain a decoded structured sparse attention map.

10. The method of claim 9, wherein the column-oriented distribution pattern is encoded as a column-oriented mask, and the step of obtaining the decoded structured sparse attention map includes: The decoded structured sparse attention graph is reconstructed using the corresponding values ​​of the retained items and the column-directed mask.

11. The method of claim 9, wherein the row distribution pattern is encoded as a row mask, and the step of obtaining the decoded structured sparse attention map includes: The decoded structured sparse attention graph is reconstructed using the corresponding values ​​of the retained items and the row direction mask.

12. The method of claim 9, wherein the diagonal distribution pattern is encoded as a diagonal mask, and the step of obtaining the decoded structured sparse attention map comprises: The decoded structured sparse attention map is reconstructed by using the corresponding values ​​of the retained items and the diagonal mask.

13. The method as described in claim 9, further comprising: The non-zero items in the structured sparse attention map within the compressed data are reordered to produce a reordered sequence; And based on the reordered sequence, the compressed data is decoded to generate the decoded structured sparse attention map.

14. The method as described in claim 13, wherein the step of reordering the non-zero items includes: Based on the different memory partitions of the memory circuit, the non-zero items are grouped.

15. The method of claim 13, wherein after the reordering step for the non-zero items, the method further comprises: The reordered sparse attention map is multiplied by the value matrix to produce a dense output matrix.

16. An apparatus for accelerating attention computation in a neural network based on a transformation model, comprising: Memory circuitry; And at least one processing circuit for: receiving the input attention map; Based on the average attention map, a structured sparse attention map is generated, wherein the average attention map is pre-computed based on multiple training attention maps, and the structured sparse attention map includes at least one of column distribution pattern, row distribution pattern and diagonal distribution pattern; and the structured sparse attention map is encoded into compressed data, wherein the compressed data includes the position and corresponding value of the retained items in the input attention map.

Images