Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

2024-06-11 · Zixuan Wang, Stanley Wei, Daniel Hsu, Jason D. Lee

The transformer architecture has prevailed in various deep learning settings due to its exceptional capabilities to select and compose structural information. Motivated by these capabilities, Sanford et al. proposed the sparse token selection task, in which transformers excel while fully-connected networks (FCNs) fail in the worst case. Building upon that, we strengthen the FCN lower bound to an average-case setting and establish an algorithmic separation of transformers over FCNs. Specifically, a one-layer transformer trained with gradient descent provably learns the sparse token selection task and, surprisingly, exhibits strong out-of-distribution length generalization. We provide empirical simulations to justify our theoretical findings.

📄 PDF Abstract BibTeX arXiv:2406.06893

Code (0)

등록된 구현이 없습니다.

Methods 이 논문이 사용한 방법론

Max Pooling Max Pooling is a pooling operation that calculates the maximum value for patches of a feature map, and uses it to create a downsampled (pooled) feature map. It is usually…

Convolution A convolution is a type of matrix operation, consisting of a kernel, a small matrix of weights, that slides over input data performing element-wise multiplication with the…

FCN Fully Convolutional Networks, or FCNs, are an architecture used mainly for semantic segmentation. They employ solely locally connected layers, such as…

Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

2026-03-24 · Chenyang Zhang, Qingyue Zhao, Quanquan Gu, Yuan Cao arxiv

Transformers have achieved great success across a wide range of applications, yet the theoretical foundations underlying their success remain largely unexplored. To demystify the strong capacities of transformers applied…

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

2025-12-18 · Yifan Zhou, Zeqi Xiao, Tianyi Wei, Shuai Yang 외 arxiv

Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost fundamentally limits scaling to long token sequences. Recent Top-K sparse attention approaches reduce t…

Image Generation

SaiT: Sparse Vision Transformers through Adaptive Token Pruning

2022-10-11 · Ling Li, David Thorsley, Joseph Hassoun

While vision transformers have achieved impressive results, effectively and efficiently accelerating these models can further boost performances. In this work, we propose a dense/sparse training framework to obtain a uni…

Knowledge Distillation

One-Layer Transformers are Provably Optimal for In-context Reasoning and Distributional Association Learning in Next-Token Prediction Tasks

2025-05-21 · Quan Nguyen, Thanh Nguyen-Tang

We study the approximation capabilities and on-convergence behaviors of one-layer transformers on the noiseless and noisy in-context reasoning of next-token prediction. Existing theoretical results focus on understanding…

Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently

2025-11-22 · Bochen Lyu, Yiyang Jia, Xiaohao Cai, Zhanxing Zhu arxiv

Transformers can acquire Chain-of-Thought (CoT) capabilities to solve reasoning tasks via fine-tuning. Reinforcement learning (RL) and supervised fine-tuning (SFT) are two primary approaches to this end. In this work, we…

Reinforcement Learning