# In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick

arXiv: https://arxiv.org/abs/2307.02419 · HTML: https://mlsanity.com/paper/2307.02419/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 49 references that are in the map.
- First posted on arXiv 5 Jul 2023.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [MetaICL: Learning to Learn In Context](https://mlsanity.com/paper/2110.15943/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning](https://mlsanity.com/paper/2301.11916/) (2023)
- [General-Purpose In-Context Learning by Meta-Learning Transformers](https://mlsanity.com/paper/2212.04458/) (2022)
- [Fast Attention Requires Bounded Entries](https://mlsanity.com/paper/2302.13214/) (2023)
- [Language Modeling with Deep Transformers](https://mlsanity.com/paper/1905.04226/) (2019)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Structured Prompting: Scaling In-Context Learning to 1,000 Examples](https://mlsanity.com/paper/2212.06713/) (2022)
- [Training Multi-Layer Over-Parametrized Neural Network in Subquadratic Time](https://mlsanity.com/paper/2112.07628/) (2021)
- [KDEformer: Accelerating Transformers via Kernel Density Estimation](https://mlsanity.com/paper/2302.02451/) (2023)
- [Meta Learning Backpropagation And Improving It](https://mlsanity.com/paper/2012.14905/) (2020)
- [Attention Scheme Inspired Softmax Regression](https://mlsanity.com/paper/2304.10411/) (2023)
- [Relative Error Tensor Low Rank Approximation](https://mlsanity.com/paper/1704.08246/) (2017)
- [An Over-parameterized Exponential Regression](https://mlsanity.com/paper/2303.16504/) (2023)
- [Solving Regularized Exp, Cosh and Sinh Regression Problems](https://mlsanity.com/paper/2303.15725/) (2023)
- [The Closeness of In-Context Learning and Weight Shifting for Softmax Regression](https://mlsanity.com/paper/2304.13276/) (2023)
- [Approximating How Single Head Attention Learns](https://mlsanity.com/paper/2103.07601/) (2021)

## Built on by (most cited first, 30 of 32)

- [A Theoretical Understanding of Self-Correction through In-context Alignment](https://mlsanity.com/paper/2405.18634/) (2024)
- [On Statistical Rates and Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)](https://mlsanity.com/paper/2407.01079/) (2024)
- [Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers](https://mlsanity.com/paper/2405.05219/) (2024)
- [Computational Limits of Low-Rank Adaptation (LoRA) Fine-Tuning for Transformer Models](https://mlsanity.com/paper/2406.03136/) (2024)
- [A Fast Optimization View: Reformulating Single Layer Attention in LLM Based on Tensor and SVM Trick, and Solving It in Matrix Multiplication Time](https://mlsanity.com/paper/2309.07418/) (2023)
- [The Fine-Grained Complexity of Gradient Computation for Training Large Language Models](https://mlsanity.com/paper/2402.04497/) (2024)
- [GradientCoin: A Peer-to-Peer Decentralized Large Language Models](https://mlsanity.com/paper/2308.10502/) (2023)
- [The Computational Limits of State-Space Models and Mamba via the Lens of Circuit Complexity](https://mlsanity.com/paper/2412.06148/) (2024)
- [Looped ReLU MLPs May Be All You Need as Practical Programmable Computers](https://mlsanity.com/paper/2410.09375/) (2024)
- [Beyond Classical Attention: Quantum Attention for Scalable Computation](https://mlsanity.com/paper/2307.08045/) (2023)
- [Zero-th Order Algorithm for Softmax Attention Optimization](https://mlsanity.com/paper/2307.08352/) (2023)
- [Why Softmax Attention Outperforms Linear Attention](https://mlsanity.com/paper/2310.11685/) (2023)
- [Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond](https://mlsanity.com/paper/2405.03251/) (2024)
- [A Unified Scheme of ResNet and Softmax](https://mlsanity.com/paper/2309.13482/) (2023)
- [Fine-tune Language Models to Approximate Unbiased In-context Learning](https://mlsanity.com/paper/2310.03331/) (2023)
- [Unmasking Transformers: A Theoretical Approach to Data Recovery via Attention Weights](https://mlsanity.com/paper/2310.12462/) (2023)
- [Efficient SGD Neural Network Training via Sublinear Activated Neuron Identification](https://mlsanity.com/paper/2307.06565/) (2023)
- [Theoretical Guarantees for High Order Trajectory Refinement in Generative Flows](https://mlsanity.com/paper/2503.09069/) (2025)
- [An Automatic Learning Rate Schedule Algorithm for Achieving Faster Convergence and Steeper Descent](https://mlsanity.com/paper/2310.11291/) (2023)
- [Local Convergence of Approximate Newton Method for Two Layer Nonlinear Regression](https://mlsanity.com/paper/2311.15390/) (2023)
- [The Expressibility of Polynomial based Attention Scheme](https://mlsanity.com/paper/2310.20051/) (2023)
- [Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining](https://mlsanity.com/paper/2502.05790/) (2025)
- [Convergence of Two-Layer Regression with Nonlinear Units](https://mlsanity.com/paper/2308.08358/) (2023)
- [A Theoretical Insight into Attack and Defense of Gradient Leakage in Transformer](https://mlsanity.com/paper/2311.13624/) (2023)
- [One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space](https://mlsanity.com/paper/2311.14652/) (2023)
- [Advancing the Understanding of Fixed Point Iterations in Deep Neural Networks: A Detailed Analytical Study](https://mlsanity.com/paper/2410.11279/) (2024)
- [Neural Algorithmic Reasoning for Hypergraphs with Looped Transformers](https://mlsanity.com/paper/2501.10688/) (2025)
- [Scaling Law Phenomena Across Regression Paradigms: Multiple and Kernel Approaches](https://mlsanity.com/paper/2503.01314/) (2025)
- [Time and Memory Trade-off of KV-Cache Compression in Tensor Transformer Decoding](https://mlsanity.com/paper/2503.11108/) (2025)
- [Theoretical Foundation of Flow-Based Time Series Generation: Provable Approximation, Generalization, and Efficiency](https://mlsanity.com/paper/2503.14076/) (2025)

## Abstract (arXiv)

Large language models (LLMs) have brought significant and transformative changes in human society. These models have demonstrated remarkable capabilities in natural language understanding and generation, leading to various advancements and impacts across several domains. We consider the in-context learning under two formulation for attention related regression in this work. Given matrices A1 ∈ ℝn × d, and A2 ∈ ℝn × d and B ∈ ℝn × n, the purpose is to solve some certain optimization problems: Normalized version minX | D(X)-1 exp(A1 X A2^top) - B |F2 and Rescaled version | exp(A1 X A2^top) - D(X) · B |F2. Here D(X) := diag( exp(A1 X A2^top) bf 1n ). Our regression problem shares similarities with previous studies on softmax-related regression. Prior research has extensively investigated regression techniques related to softmax regression: Normalized version | langle exp(Ax) , bf 1n rangle-1 exp(Ax) - b |22 and Resscaled version | exp(Ax) - langle exp(Ax), bf 1n rangle b |22  In contrast to previous approaches, we adopt a vectorization technique to address the regression problem in matrix formulation. This approach expands the dimension from d to d2, resembling the formulation of the regression problem mentioned earlier. Upon completing the lipschitz analysis of our regression function, we have derived our main result concerning in-context learning.