Polynomial Regression as a Task for Understanding In-context Learning Through Finetuning and Alignment

TL;DR

The paper introduces univariate polynomial regression based on Chebyshev polynomials as a lightweight, structured proxy to study prompting and alignment in in-context learning. It implements a GPT2-style model and compares parameter-efficient fine-tuning methods (LoRA) with soft prompting on polynomial tasks, including a clamping-based alignment toy and a jailbreaking scenario. The results show LoRA generally outperforms soft prompting in low-data regimes, that in-context learning can capture polynomial regression dynamics, and that alignment-like behavior can be learned in-context and manipulated via jailbreaking, echoing observations from larger LLMs. The work argues that this toy setup captures essential ICL and alignment dynamics while enabling clearer visualization and reduced computational costs, offering a useful testbed for future analyses.

Abstract

Simple function classes have emerged as toy problems to better understand in-context-learning in transformer-based architectures used for large language models. But previously proposed simple function classes like linear regression or multi-layer-perceptrons lack the structure required to explore things like prompting and alignment within models capable of in-context-learning. We propose univariate polynomial regression as a function class that is just rich enough to study prompting and alignment, while allowing us to visualize and understand what is going on clearly.

Figures (18)

• Figure 1: Sampled density heatmap of joint (x, y) distribution for uniformly random x-values and linear combinations of Chebyshev polynomials. See appendix for a more detailed explanation (section \ref{['sec:app_dist_explenation']}). (a) Polynomials of random degree up to 11. The distribution of $y$ is approximately independent of $x$. (b) Polynomials of degree 5 (specifically), but the first 3 linear coefficients are fixed. The distribution of $y$ is not independent of $x$.
• Figure 2: Performance of pretrained model on normal random linear combinations of Chebyshev polynomials of a random degree between 0 and 11. Polynomial regression and ridge-regularized polynomial regression are used as baselines (See section \ref{['sec:app_baselines']}). Using polynomial regression, 12 examples are theoretically needed to achieve $0 \approx 10^{-\infty}$ error for a degree 11 polynomial. Shaded areas represent the 95% bootstrap confidence interval (\ref{['sec:app_bootstrap']}). See Appendix \ref{['sec:normal_scale_lora_vs_soft_prompt']} for plot using linear y-scale.
• Figure 3: Performance of finetuning methods for normal random linear combinations of Chebyshev polynomials of degree at most 5, and with the first $n$ first linear combination coefficients fixed. As the number of fixed coefficients increases, $y$ becomes more dependent on $x$. Polynomial regression and ridge-regularized polynomial regression are used as baselines (See section \ref{['sec:app_baselines']}). Shaded areas represent the 95% bootstrap confidence interval (\ref{['sec:app_bootstrap']}). See Appendix \ref{['sec:normal_scale_lora_vs_soft_prompt']} for plot using linear y-scale.
• Figure 4: (\ref{['fig:aligned-in-context']}) Increased context length increases performance for clamped polynomials, although the model were not trained to do this. Note that on points above the threshold loss increases, whereas the same is not true for points below the threshold. (\ref{['fig:polynomial-plot']}) Adding in a clamped context leads to good predictions in which our predicted polynomial looks like a smoothed version of the actual clamped polynomial. This example is a degree 4 polynomial with 25 in-context examples of which 9 are clamped.
• Figure 5: (\ref{['fig:finetuned-alignment']}) The transformer model is able to perform the task of clamping values in-context after being additionally finetuned. (\ref{['fig:jailbreak']}) As the context length increases, so does the number of jailbroken examples. Similar to language models, as the number of jailbroken examples increases, it becomes more likely for alignment to be broken.