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!
_version_ 1866912321270448128
author Zhang, Chenyang
Meng, Xuran
Cao, Yuan
author_facet Zhang, Chenyang
Meng, Xuran
Cao, Yuan
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.
format Preprint
id arxiv_https___arxiv_org_abs_2504_08638
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Transformer Learns Optimal Variable Selection in Group-Sparse Classification
Zhang, Chenyang
Meng, Xuran
Cao, Yuan
Machine Learning
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.
title Transformer Learns Optimal Variable Selection in Group-Sparse Classification
topic Machine Learning
url https://arxiv.org/abs/2504.08638