Test-time training (TTT) enhances model performance by explicitly updating designated parameters prior to each prediction to adapt to the test data. While TTT has demonstrated considerable empirical success, its theoretical underpinnings remain limited, particularly for nonlinear models. In this paper, we investigate the combination of TTT with in-context learning (ICL), where the model is given a few examples from the target distribution at inference time. We analyze this framework in the setting of single-index models $y=σ_*(\langle β, \mathbf{x} \rangle)$, where the feature vector $β$ is drawn from a hidden low-dimensional subspace. For single-layer transformers trained with gradient-based algorithms and adopting TTT, we establish an upper bound on the prediction risk. Our theory reveals that TTT enables the single-layer transformers to adapt to both the feature vector $β$ and the link function $σ_*$, which vary across tasks. This creates a sharp contrast with ICL alone, which is theoretically difficult to adapt to shifts in the link function. Moreover, we provide the convergence rate with respect to the data length, showing the predictive error can be driven arbitrarily close to the noise level as the context size and the network width grow.

翻译：测试时训练通过在每个预测前显式更新指定参数以适应测试数据，从而提升模型性能。尽管TTT在实证中取得了显著成功，但其理论基础仍较为有限，尤其对于非线性模型。本文研究了TTT与情境学习的结合框架，其中模型在推理时从目标分布中获取少量示例。我们在单索引模型$y=σ_*(\langle β, \mathbf{x} \rangle)$设定下分析该框架，其中特征向量$β$从隐藏的低维子空间中采样。针对采用基于梯度算法训练并实施TTT的单层Transformer，我们建立了预测风险的上界。理论分析表明，TTT使单层Transformer能够适应跨任务变化的特征向量$β$与链接函数$σ_*$。这与单纯ICL形成鲜明对比——理论上ICL难以适应链接函数的偏移。此外，我们给出了关于数据长度的收敛速率，证明随着上下文规模与网络宽度的增加，预测误差可被驱动至任意接近噪声水平。