Transformers: A Graph Processing Perspective

Transformers, a variant of AI/ML models that utilize the attention mechanism to capture interactions in sequential data, have significantly advanced a variety of scientific and engineering applications. However, attention implementations suffer from high computational and memory complexity requirements and despite efforts to reduce their complexity, the context length over which they can capture interactions is still several orders of magnitude lower than desired. The fundamental problem is that even though these models capture pairwise interactions between data elements, they are viewed as sequence of tensor operations on tensor data, thereby, limiting avenues of optimization. In this work, we take the first steps towards reimagining transformer based models as graph computing pipelines by implementing the attention mechanism. Our expectation is that this view will not only dramatically increase scalability by unlocking parallelism along the context length dimension, but will also enable researchers to apply the vast library of graph analytics algorithms to obtain better insights into the inner working of these models. We implement graph algorithms for attention mechanism and conduct extensive experimentation by varying sequence length, token dimensions, and sparsity factor and observer near linear reduction in computation time with sparsity factor.