Two-Pass SoftMax Computation Without Division Bottlenecks
Find Innovative SolutionsGenerate Solutions
Solution Overview
Problem
Transformer-based neural networks face significant computational bottlenecks due to the high runtime consumption of the SoftMax operation, particularly in longer sequence lengths and larger models, which can account for 20%-40% of the overall processor runtime.
Innovation Solution
A modified 2-pass SoftMax operation is implemented, eliminating inverse multiplications or divisions by modifying the first pass to include scalar operations for calculating the logarithm of the denominator and an operand value, and the second pass to perform addition and exponentiation, maintaining mathematical equivalence and model accuracy without the need for approximation.
Engineering Contradictions & Design Principles
Engineering Contradiction Analysis
1Measurement precision
If the conventional SoftMax operation is used in transformer-based neural networks, then model accuracy is maintained, but computational runtime increases significantly (consuming 20%-40% of overall processor runtime)
Solution Approach 1:
The patent transforms the SoftMax computation by changing the mathematical parameters and operations involved. Instead of directly computing e^(x_i - max(x)) / sum(e^(x_j - max(x))), the invention uses logarithmic transformation to convert division operations into subtraction operations: log(softmax(x)) = x - max(x) - log(sum(e^(x_j - max(x)))). This parameter transformation eliminates the computationally expensive division operations while maintaining numerical stability and accuracy.
Solution Approach 2:
The patent extracts and eliminates the harmful computational operations (inverse multiplications/divisions) from the SoftMax function. By separating the logarithmic computation from the exponential computation and using pre-computed logarithmic values, the invention removes the bottleneck operations that cause 20%-40% runtime consumption while preserving the essential functionality of the SoftMax operation.
2Productivity
If the number of passes in SoftMax computation is reduced from 3 to 2, then computational efficiency improves, but implementation complexity increases due to modified scalar operations
Solution Approach 1:
The patent merges the computation of the logarithm of the denominator with the exponentiation operation in the two-pass algorithm. Instead of computing the denominator sum separately and then performing division, the invention computes log(sum(e^(x_j - max(x)))) during the first pass and uses this pre-computed value in the second pass, merging multiple operations into a unified computational flow that reduces passes while managing complexity.
Solution Approach 2:
The patent performs preliminary computation of logarithmic values during the first pass before the actual SoftMax computation in the second pass. By pre-computing log(sum(e^(x_j - max(x)))) and storing it as a cached value, the invention eliminates the need for expensive division operations in the second pass, improving efficiency while organizing complexity into manageable preliminary and final computation stages.
Data Source
Figure 1
Figure 2
Figure 3
AI summary
A modified 2-pass version of the SoftMax operation can be implemented to address reduce computational cost without loss of accuracy, in particular for deep learning neural networks such as transformer-based neural networks and large language models (LLMs). The first pass is modified to include two scalar operations at the end. At the end of the first pass, a first scalar operation is performed to calculate a logarithm of the denominator, and a second scalar operation is performed to calculate an operand value based on a sum of the logarithm of the denominator and the maximum value. The second pass is modified to perform addition and exponentiation. In the second pass, an element of an input tensor is subtracted by the operand value to obtain an exponent, and a base is raised to the exponent. The second pass avoids divisions.