Attention-Only Transformers via Unrolled Subspace Denoising

TL;DR

This paper introduces a simplified, interpretable transformer architecture based on unrolled subspace denoising, which achieves competitive performance on vision and language tasks by iteratively refining token representations through self-attention.

Contribution

The authors propose a new transformer design that uses unrolled denoising operations with only self-attention and skip connections, providing mathematical justification and interpretability.

Findings

Achieves performance close to GPT-2 and CRATE on vision and language tasks.

Each layer denoises token representations at a linear rate.

Simplifies transformer architecture to only self-attention with skip connections.

Abstract

Despite the popularity of transformers in practice, their architectures are empirically designed and neither mathematically justified nor interpretable. Moreover, as indicated by many empirical studies, some components of transformer architectures may be redundant. To derive a fully interpretable transformer architecture with only necessary components, we contend that the goal of representation learning is to compress a set of noisy initial token representations towards a mixture of low-dimensional subspaces. To compress these noisy token representations, an associated denoising operation naturally takes the form of a multi-head (subspace) self-attention. By unrolling such iterative denoising operations into a deep network, we arrive at a highly compact architecture that consists of \textit{only} self-attention operators with skip connections at each layer. Moreover, we show that each layer performs highly efficient denoising: it improves the signal-to-noise ratio of token representations \textit{at a linear rate} with respect to the number of layers. Despite its simplicity, extensive experiments on vision and language tasks demonstrate that such a transformer achieves performance close to that of standard transformer architectures such as GPT-2 and CRATE.