Changhyeon Lee

Softmax Output Approximation for Activation Memory-Efficient Training of Attention-based Networks

In this paper, we propose to approximate the softmax output, which is the key product of the attention mechanism, to reduce its activation memory usage when training attention-based networks (aka Transformers). During the forward pass of the network, the proposed softmax output approximation method…