One-Layer Transformer Provably Learns One-Nearest Neighbor In Context

• URL: http://arxiv.org/abs/2411.10830v1
• Date: Sat, 16 Nov 2024 16:12:42 GMT
• Title: One-Layer Transformer Provably Learns One-Nearest Neighbor In Context
• Authors: Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason M. Klusowski, Jianqing Fan, Mengdi Wang,
• Abstract summary: We study the capability of one-layer transformers learning the one-nearest neighbor rule. A single softmax attention layer can successfully learn to behave like a one-nearest neighbor.
• Score: 48.4979348643494
• License:
• Abstract: Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

Related papers

• Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? [69.4145579827826]
  We show a fast flow on the regression loss despite the gradient non-ity algorithms for our convergence landscape. This is the first theoretical analysis for multi-layer Transformer in this setting.
  arXiv  Detail & Related papers  (2024-10-10T18:29:05Z)
• In-Context Learning with Representations: Contextual Generalization of Trained Transformers [66.78052387054593]
  In-context learning (ICL) refers to a capability of pretrained large language models, which can learn a new task given a few examples during inference. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks.
  arXiv  Detail & Related papers  (2024-08-19T16:47:46Z)
• Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysis [63.66763657191476]
  We show that efficient numerical training and inference algorithms as low-rank computation have impressive performance for learning Transformer-based adaption. We analyze how magnitude-based models affect generalization while improving adaption. We conclude that proper magnitude-based has a slight on the testing performance.
  arXiv  Detail & Related papers  (2024-06-24T23:00:58Z)
• In-Context Convergence of Transformers [63.04956160537308]
  We study the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent. For data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process.
  arXiv  Detail & Related papers  (2023-10-08T17:55:33Z)
• Trained Transformers Learn Linear Models In-Context [39.56636898650966]
  Attention-based neural networks as transformers have demonstrated a remarkable ability to exhibit inattention learning (ICL) We show that when transformer training over random instances of linear regression problems, these models' predictions mimic nonlinear of ordinary squares.
  arXiv  Detail & Related papers  (2023-06-16T15:50:03Z)
• Transformers learn to implement preconditioned gradient descent for in-context learning [41.74394657009037]
  Several recent works demonstrate that transformers can implement algorithms like gradient descent. We ask: Can transformers learn to implement such algorithms by training over random problem instances? For a transformer with $L$ attention layers, we prove certain critical points of the training objective implement $L$ iterations of preconditioned gradient descent.
  arXiv  Detail & Related papers  (2023-06-01T02:35:57Z)
• Transformers learn in-context by gradient descent [58.24152335931036]
  Training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass.
  arXiv  Detail & Related papers  (2022-12-15T09:21:21Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.