Attention layers provably solve single-location regression

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Marion, Pierre, Berthier, Raphaël, Biau, Gérard, Boyer, Claire
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916630097821696
author Marion, Pierre
Berthier, Raphaël
Biau, Gérard
Boyer, Claire
author_facet Marion, Pierre
Berthier, Raphaël
Biau, Gérard
Boyer, Claire
contents Attention-based models, such as Transformer, excel across various tasks but lack a comprehensive theoretical understanding, especially regarding token-wise sparsity and internal linear representations. To address this gap, we introduce the single-location regression task, where only one token in a sequence determines the output, and its position is a latent random variable, retrievable via a linear projection of the input. To solve this task, we propose a dedicated predictor, which turns out to be a simplified version of a non-linear self-attention layer. We study its theoretical properties, by showing its asymptotic Bayes optimality and analyzing its training dynamics. In particular, despite the non-convex nature of the problem, the predictor effectively learns the underlying structure. This work highlights the capacity of attention mechanisms to handle sparse token information and internal linear structures.
format Preprint
id arxiv_https___arxiv_org_abs_2410_01537
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Attention layers provably solve single-location regression
Marion, Pierre
Berthier, Raphaël
Biau, Gérard
Boyer, Claire
Machine Learning
Attention-based models, such as Transformer, excel across various tasks but lack a comprehensive theoretical understanding, especially regarding token-wise sparsity and internal linear representations. To address this gap, we introduce the single-location regression task, where only one token in a sequence determines the output, and its position is a latent random variable, retrievable via a linear projection of the input. To solve this task, we propose a dedicated predictor, which turns out to be a simplified version of a non-linear self-attention layer. We study its theoretical properties, by showing its asymptotic Bayes optimality and analyzing its training dynamics. In particular, despite the non-convex nature of the problem, the predictor effectively learns the underlying structure. This work highlights the capacity of attention mechanisms to handle sparse token information and internal linear structures.
title Attention layers provably solve single-location regression
topic Machine Learning
url https://arxiv.org/abs/2410.01537