Bu, D., Huang, W., Han, A., Nitanda, A., Zhang, Q., Wong, H. S., & Suzuki, T. (2025, October). Provable In-Context Vector Arithmetic via Retrieving Task Concepts. In International Conference on Machine Learning (pp. 5669-5724). PMLR.
Abstract:
In-context learning (ICL) has garnered significant attention for its ability to grasp functions/tasks from demonstrations. Recent studies suggest the presence of a latent task/function vector in LLMs during ICL. Merullo et al. (2024) showed that LLMs leverage this vector alongside the residual stream for Word2Vec-like vector arithmetic, solving factual-recall ICL tasks. Additionally, recent work empirically highlighted the key role of Question-Answer data in enhancing factual-recall capabilities. Despite these insights, a theoretical explanation remains elusive. To move one step forward, we propose a theoretical framework building on empirically grounded hierarchical concept modeling. We develop an optimization theory, showing how nonlinear residual transformers trained via gradient descent on cross-entropy loss perform factual-recall ICL tasks via vector arithmetic. We prove 0-1 loss convergence and show the strong generalization, including robustness to concept recombination and distribution shifts. These results elucidate the advantages of transformers over static embedding predecessors. Empirical simulations corroborate our theoretical insights.
License type:
Publisher Copyright
Funding Info:
This research / project is supported by the National Research Foundation, Singapore, Infocomm Media Development Authority - Trust Tech Funding Initiative
Grant Reference no. : DTC-RGC-05
This research / project is supported by the Ministry of Digital Development and Information - AI Visiting Professorship Programme
Grant Reference no. : AIVP- 2024-004