Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Julistiono, Addison Kristanto, Tarzanagh, Davoud Ataee, Azizan, Navid
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918316632702976
author Julistiono, Addison Kristanto
Tarzanagh, Davoud Ataee
Azizan, Navid
author_facet Julistiono, Addison Kristanto
Tarzanagh, Davoud Ataee
Azizan, Navid
contents Attention mechanisms have revolutionized several domains of artificial intelligence, such as natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. While recent work has characterized the optimization dynamics of gradient descent (GD) in attention-based models and the structural properties of its preferred solutions, less is known about more general optimization algorithms such as mirror descent (MD). In this paper, we investigate the convergence properties and implicit biases of a family of MD algorithms tailored for softmax attention mechanisms, with the potential function chosen as the $p$-th power of the $\ell_p$-norm. Specifically, we show that these algorithms converge in direction to a generalized hard-margin SVM with an $\ell_p$-norm objective when applied to a classification problem using a softmax attention model. Notably, our theoretical results reveal that the convergence rate is comparable to that of traditional GD in simpler models, despite the highly nonlinear and nonconvex nature of the present problem. Additionally, we delve into the joint optimization dynamics of the key-query matrix and the decoder, establishing conditions under which this complex joint optimization converges to their respective hard-margin SVM solutions. Lastly, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.
format Preprint
id arxiv_https___arxiv_org_abs_2410_14581
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection
Julistiono, Addison Kristanto
Tarzanagh, Davoud Ataee
Azizan, Navid
Machine Learning
Artificial Intelligence
Computation and Language
Attention mechanisms have revolutionized several domains of artificial intelligence, such as natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. While recent work has characterized the optimization dynamics of gradient descent (GD) in attention-based models and the structural properties of its preferred solutions, less is known about more general optimization algorithms such as mirror descent (MD). In this paper, we investigate the convergence properties and implicit biases of a family of MD algorithms tailored for softmax attention mechanisms, with the potential function chosen as the $p$-th power of the $\ell_p$-norm. Specifically, we show that these algorithms converge in direction to a generalized hard-margin SVM with an $\ell_p$-norm objective when applied to a classification problem using a softmax attention model. Notably, our theoretical results reveal that the convergence rate is comparable to that of traditional GD in simpler models, despite the highly nonlinear and nonconvex nature of the present problem. Additionally, we delve into the joint optimization dynamics of the key-query matrix and the decoder, establishing conditions under which this complex joint optimization converges to their respective hard-margin SVM solutions. Lastly, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.
title Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2410.14581