Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transformers are Provably Optimal In-context Estimators for Wireless Communications

Vishnu Teja Kunde, Vicram Rajagopalan, Chandra Shekhara Kaushik Valmeekam, Krishna Narayanan, Srinivas Shakkottai, Dileep Kalathil, Jean-Francois Chamberland

TL;DR

This work defines in-context estimation (ICE) for symbol recovery in wireless channels, showing that a single-layer softmax-attention transformer is asymptotically optimal for a subclass of ICE problems. It proves that the transformer’s optimal configuration corresponds to the MMSE trainer and demonstrates that training loss minimization yields this optimality. Beyond theory, the authors empirically validate multi-layer transformers on time-invariant and time-varying channel models, illustrating near-optimal context-aware performance with modest in-context prompts. The results suggest transformers can serve as robust, plug-in receivers that adapt to latent channel contexts without explicit parameter updates. The work provides open-source code and broad implications for adaptive, data-driven wireless receivers.

Abstract

Pre-trained transformers exhibit the capability of adapting to new tasks through in-context learning (ICL), where they efficiently utilize a limited set of prompts without explicit model optimization. The canonical communication problem of estimating transmitted symbols from received observations can be modeled as an in-context learning problem: received observations are a noisy function of transmitted symbols, and this function can be represented by an unknown parameter whose statistics depend on an unknown latent context. This problem, which we term in-context estimation (ICE), has significantly greater complexity than the extensively studied linear regression problem. The optimal solution to the ICE problem is a non-linear function of the underlying context. In this paper, we prove that, for a subclass of such problems, a single-layer softmax attention transformer (SAT) computes the optimal solution of the above estimation problem in the limit of large prompt length. We also prove that the optimal configuration of such a transformer is indeed the minimizer of the corresponding training loss. Further, we empirically demonstrate the proficiency of multi-layer transformers in efficiently solving broader in-context estimation problems. Through extensive simulations, we show that solving ICE problems using transformers significantly outperforms standard approaches. Moreover, just with a few context examples, it achieves the same performance as an estimator with perfect knowledge of the latent context. The code is available \href{https://github.com/vishnutez/in-context-estimation}{here}.

Transformers are Provably Optimal In-context Estimators for Wireless Communications

TL;DR

This work defines in-context estimation (ICE) for symbol recovery in wireless channels, showing that a single-layer softmax-attention transformer is asymptotically optimal for a subclass of ICE problems. It proves that the transformer’s optimal configuration corresponds to the MMSE trainer and demonstrates that training loss minimization yields this optimality. Beyond theory, the authors empirically validate multi-layer transformers on time-invariant and time-varying channel models, illustrating near-optimal context-aware performance with modest in-context prompts. The results suggest transformers can serve as robust, plug-in receivers that adapt to latent channel contexts without explicit parameter updates. The work provides open-source code and broad implications for adaptive, data-driven wireless receivers.

Abstract

Pre-trained transformers exhibit the capability of adapting to new tasks through in-context learning (ICL), where they efficiently utilize a limited set of prompts without explicit model optimization. The canonical communication problem of estimating transmitted symbols from received observations can be modeled as an in-context learning problem: received observations are a noisy function of transmitted symbols, and this function can be represented by an unknown parameter whose statistics depend on an unknown latent context. This problem, which we term in-context estimation (ICE), has significantly greater complexity than the extensively studied linear regression problem. The optimal solution to the ICE problem is a non-linear function of the underlying context. In this paper, we prove that, for a subclass of such problems, a single-layer softmax attention transformer (SAT) computes the optimal solution of the above estimation problem in the limit of large prompt length. We also prove that the optimal configuration of such a transformer is indeed the minimizer of the corresponding training loss. Further, we empirically demonstrate the proficiency of multi-layer transformers in efficiently solving broader in-context estimation problems. Through extensive simulations, we show that solving ICE problems using transformers significantly outperforms standard approaches. Moreover, just with a few context examples, it achieves the same performance as an estimator with perfect knowledge of the latent context. The code is available \href{https://github.com/vishnutez/in-context-estimation}{here}.
Paper Structure (20 sections, 4 theorems, 40 equations, 3 figures, 2 tables)

This paper contains 20 sections, 4 theorems, 40 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. In-Context Estimation
  3. Main Results
  4. Experiments with Multi-Layer Transformers
  5. Experimental setup
  6. Scenario 1: Time invariant process
  7. Scenario 2: Time varying process
  8. Training
  9. Baselines
  10. Results and discussion
  11. Scenario 1: Estimation under Time-invariant process
  12. Scenario 2: Estimation under Time-varying process
  13. Conclusion
  14. Related Work
  15. Appendix for Theoretical Results
  16. ...and 5 more sections

Key Result

Lemma 1

(MMSE estimate) The optimal estimator $\hat{\mathbf{x}}^{\rm MMSE}: \mathbb{R}^{2d} \times \mathbb{R}^{2d\times 2} \to \mathbb{R}^2$ for $\mathbf{x} \in \mathcal{X} \subset \mathbb{S}^2$ given $\mathbf{y} \in \mathbb{R}^{2d}, \mathbf{H} \in \mathbb{R}^{2d\times 2}$ where $\mathbf{y} = \mathbf{H} \m where $\mathbf{H} \triangleq $ for some $\tilde{\mathbf{h}} \in \mathbb{C}^d$, and $\boldsymbol{\Si

Figures (3)

  • Figure 1: Transformer model performing causal attention
  • Figure 2: Estimation error in scenario 1 as a function of number of in-context examples. The performance of the transformer (red) is close to that of the optimal context-aware estimator (black) while significantly outperforming a typical baseline (green) when latent context $\theta=0$.
  • Figure 3: Estimation error in scenario 2 as a function of number of in-context examples. The transformer (red) computes the optimal estimator (orange) with its performance approaching the context-aware estimator (black) as the number of examples increases, while significantly outperforming typical baselines (green, blue) for latent contexts $\theta=15, 30$.

Theorems & Definitions (5)

  • Definition 1
  • Lemma 1
  • Lemma 2
  • Theorem 3
  • Theorem 4