Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks

• URL: http://arxiv.org/abs/2406.02550v2
• Date: Mon, 04 Nov 2024 16:04:27 GMT
• Title: Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks
• Authors: Tianyu He, Darshil Doshi, Aritra Das, Andrey Gromov,
• Abstract summary: We study the emergence of in-context learning and skill composition in a collection of modular arithmetic tasks. Specifically, we consider a finite collection of linear modular functions $z = a, x + b, y ;mathrmmod; p$ labeled by the vector $(a, b) in mathbbZ_p2$.
• Score: 5.358878931933351
• License:
• Abstract: Large language models can solve tasks that were not present in the training set. This capability is believed to be due to in-context learning and skill composition. In this work, we study the emergence of in-context learning and skill composition in a collection of modular arithmetic tasks. Specifically, we consider a finite collection of linear modular functions $z = a \, x + b \, y \;\mathrm{mod}\; p$ labeled by the vector $(a, b) \in \mathbb{Z}_p^2$. We use some of these tasks for pre-training and the rest for out-of-distribution testing. We empirically show that a GPT-style transformer exhibits a transition from in-distribution to out-of-distribution generalization as the number of pre-training tasks increases. We find that the smallest model capable of out-of-distribution generalization requires two transformer blocks, while for deeper models, the out-of-distribution generalization phase is \emph{transient}, necessitating early stopping. Finally, we perform an interpretability study of the pre-trained models, revealing highly structured representations in both attention heads and MLPs; and discuss the learned algorithms. Notably, we find an algorithmic shift in deeper models, as we go from few to many in-context examples.

Related papers

• 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)
• Towards Better Understanding of In-Context Learning Ability from In-Context Uncertainty Quantification [7.869708570399577]
  We consider a bi-objective prediction task of predicting both the conditional expectation $mathbbE[Y|X]$ and the conditional variance Var$(Y|X)$. Theoretically, we show that the trained Transformer reaches near Bayes-optimum, suggesting the usage of the information of the training distribution.
  arXiv  Detail & Related papers  (2024-05-24T00:08:55Z)
• Limits of Transformer Language Models on Learning to Compose Algorithms [77.2443883991608]
  We evaluate training LLaMA models and prompting GPT-4 and Gemini on four tasks demanding to learn a composition of several discrete sub-tasks. Our results indicate that compositional learning in state-of-the-art Transformer language models is highly sample inefficient.
  arXiv  Detail & Related papers  (2024-02-08T16:23:29Z)
• In-context Learning Generalizes, But Not Always Robustly: The Case of Syntax [36.98247762224868]
  In-context learning (ICL) is now a common method for teaching large language models (LLMs) new tasks. Do models infer the underlying structure of the task defined by the context, or do they rely on superficial generalizations that only generalize to identically distributed examples? In experiments with models from the GPT, PaLM, and Llama 2 families, we find large variance across LMs. The variance is explained more by the composition of the pre-training corpus and supervision methods than by model size.
  arXiv  Detail & Related papers  (2023-11-13T23:52:43Z)
• A Closer Look at In-Context Learning under Distribution Shifts [24.59271215602147]
  We aim to better understand the generality and limitations of in-context learning from the lens of the simple yet fundamental task of linear regression. We find that both transformers and set-based distributions exhibit in-context learning under-distribution evaluations, but transformers more closely emulate the performance of ordinary least squares (OLS) Transformers also display better resilience to mild distribution shifts, where set-based distributions falter.
  arXiv  Detail & Related papers  (2023-05-26T07:47:21Z)
• Language Models Implement Simple Word2Vec-style Vector Arithmetic [32.2976613483151]
  A primary criticism towards language models (LMs) is their inscrutability. This paper presents evidence that, despite their size and complexity, LMs sometimes exploit a simple vector arithmetic style mechanism to solve some relational tasks.
  arXiv  Detail & Related papers  (2023-05-25T15:04:01Z)
• Categorical semantics of compositional reinforcement learning [25.752647944862183]
  Reinforcement learning (RL) often requires decomposing a problem into subtasks and composing learned behaviors on these tasks. We develop a framework for a emphcompositional theory of RL using a categorical point of view. We show that $mathsfMDP$ admits natural compositional operations, such as certain fiber products and pushouts.
  arXiv  Detail & Related papers  (2022-08-29T15:51:36Z)
• Inducing Transformer's Compositional Generalization Ability via Auxiliary Sequence Prediction Tasks [86.10875837475783]
  Systematic compositionality is an essential mechanism in human language, allowing the recombination of known parts to create novel expressions. Existing neural models have been shown to lack this basic ability in learning symbolic structures. We propose two auxiliary sequence prediction tasks that track the progress of function and argument semantics.
  arXiv  Detail & Related papers  (2021-09-30T16:41:19Z)
• Text Modular Networks: Learning to Decompose Tasks in the Language of Existing Models [61.480085460269514]
  We propose a framework for building interpretable systems that learn to solve complex tasks by decomposing them into simpler ones solvable by existing models. We use this framework to build ModularQA, a system that can answer multi-hop reasoning questions by decomposing them into sub-questions answerable by a neural factoid single-span QA model and a symbolic calculator.
  arXiv  Detail & Related papers  (2020-09-01T23:45:42Z)
• On the Theory of Transfer Learning: The Importance of Task Diversity [114.656572506859]
  We consider $t+1$ tasks parameterized by functions of the form $f_j circ h$ in a general function class $mathcalF circ mathcalH$. We show that for diverse training tasks the sample complexity needed to learn the shared representation across the first $t$ training tasks scales as $C(mathcalH) + t C(mathcalF)$.
  arXiv  Detail & Related papers  (2020-06-20T20:33:59Z)
• Neural Bayes: A Generic Parameterization Method for Unsupervised Representation Learning [175.34232468746245]
  We introduce a parameterization method called Neural Bayes. It allows computing statistical quantities that are in general difficult to compute. We show two independent use cases for this parameterization.
  arXiv  Detail & Related papers  (2020-02-20T22:28:53Z)

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.