Saved in:
Bibliographic Details
Main Authors: Zhang, Chenyang, Meng, Xuran, Cao, Yuan
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2504.08638
Tags: Add Tag
No Tags, Be the first to tag this record!
Table of Contents:
  • Transformers have demonstrated remarkable success across various applications. However, the success of transformers have not been understood in theory. In this work, we give a case study of how transformers can be trained to learn a classic statistical model with "group sparsity", where the input variables form multiple groups, and the label only depends on the variables from one of the groups. We theoretically demonstrate that, a one-layer transformer trained by gradient descent can correctly leverage the attention mechanism to select variables, disregarding irrelevant ones and focusing on those beneficial for classification. We also demonstrate that a well-pretrained one-layer transformer can be adapted to new downstream tasks to achieve good prediction accuracy with a limited number of samples. Our study sheds light on how transformers effectively learn structured data.