Softmax $\geq$ Linear: Transformers may learn to classify in-context by kernel gradient descent
This work provides theoretical insights into in-context learning algorithms for transformers in more realistic classification scenarios, which is incremental but enhances understanding for researchers in machine learning and AI.
The paper tackles the problem of understanding how transformers perform in-context learning for classification tasks, bridging from existing linear regression theories to non-linear settings with softmax activation, and finds that transformers learn to do gradient descent in kernel feature space with context-adaptive learning rates.
The remarkable ability of transformers to learn new concepts solely by reading examples within the input prompt, termed in-context learning (ICL), is a crucial aspect of intelligent behavior. Here, we focus on understanding the learning algorithm transformers use to learn from context. Existing theoretical work, often based on simplifying assumptions, has primarily focused on linear self-attention and continuous regression tasks, finding transformers can learn in-context by gradient descent. Given that transformers are typically trained on discrete and complex tasks, we bridge the gap from this existing work to the setting of classification, with non-linear (importantly, softmax) activation. We find that transformers still learn to do gradient descent in-context, though on functionals in the kernel feature space and with a context-adaptive learning rate in the case of softmax transformer. These theoretical findings suggest a greater adaptability to context for softmax attention, which we empirically verify and study through ablations. Overall, we hope this enhances theoretical understanding of in-context learning algorithms in more realistic settings, pushes forward our intuitions and enables further theory bridging to larger models.