Attention-Based AFDM Channel Estimation Network Using Diagonal Reconstruction

Abstract

Affine Frequency Division Multiplexing (AFDM) has been proposed for future high-mobility communication scenarios. However, existing AFDM channel estimation methods suffer significant performance degradation under fractional Doppler conditions due to path energy dispersion. To address this issue, we propose a deep learning network that adaptively learns path energy dispersion through a 1D processing module and a Transformer block, based on the diagonal reconstruction of the AFDM effective channel matrix. 1D processing module employs convolutions with different kernel sizes to extract pilot features, and Transformer block models vary energy dispersion patterns. The proposed method does not require prior knowledge of the number of paths and the assumption of distinct path delays. Simulation results demonstrate that at a Signal-to-Noise Ratio (SNR) of 25 dB, the proposed method achieves up to a 4 dB gain in Normalized Mean Square Error (NMSE) and an 6 dB improvement in Bit Error Rate (BER) over existing traditional methods under fractional Doppler conditions.

1. Introduction

Next-generation wireless communication systems (6G) are envisioned to encompass high-mobility communication scenarios, such as high-speed railway networks, vehicle-to-everything communications, unmanned aerial vehicles, and low-earth-orbit satellite constellations [1]. In the current 4G and 5G standards, Orthogonal Frequency Division Multiplexing (OFDM) serves as the dominant waveform. However, its inherent advantages are severely compromised in high-mobility environments. The rapid relative motion between the transmitter and receiver induces significant Doppler shifts, destroying the orthogonality among subcarriers and leading to severe Inter-Carrier Interference (ICI) [2,3]. Consequently, the research community has shifted focus towards designing new waveforms capable of achieving full diversity in doubly dispersive channels.

To address the challenges posed by high-mobility environments, Orthogonal Time Frequency Space (OTFS) modulation was introduced [4], which transforms the time-varying fading channel into a quasi-static 2D channel to achieve full diversity [5,6]. However, the 2D nature of OTFS entails high pilot overhead. As a promising alternative, Affine Frequency Division Multiplexing (AFDM) [7] has recently emerged. AFDM achieves full diversity with lower complexity by adjusting the chirp parameters based on the channel’s delay-Doppler profile. Moreover, as a one-dimensional modulation scheme, AFDM achieves channel estimation with lower guard-symbol overhead than OTFS [8].

Given its potential, AFDM has garnered widespread attention, with recent studies expanding its applicability across various domains, including Index Modulation (IM) [9], multi-user uplink schemes [10], Reconfigurable Intelligent Surfaces (RISs) [11], and Integrated Sensing and Communication (ISAC) [12,13]. However, the foundation of all the aforementioned domains is the availability of accurate Channel State Information (CSI). While several conventional channel estimation schemes have been proposed, they exhibit significant susceptibility to fractional Doppler shifts and inter-Doppler interference (IDoI), leading to severe degradation in estimation accuracy (detailed in Section 2).

To perform AFDM channel estimation under fractional-Doppler conditions, we propose a novel deep learning network based on the Transformer architecture. Our main contributions are outlined as follows:

1.

A 1D feature processing module is applied to capture pilot features with varying kernel sizes: small kernels for distinct paths and large kernels for smeared energy clusters.

2.

Self-attention module, based on Vision Transformer(ViT) blocks, is integrated to act as a global energy aggregator. This allows the network to effectively suppress path ambiguity by correlating and gathering the dispersed energy to implicitly model path parameters.

3.

Compared to conventional methods, the proposed method requires no prior knowledge of the number of paths and the assumption of distinct path delays, and demonstrates superior performance in minimizing NMSE and BER with an efficient architecture.

The paper is organized as follows. In Section 2, the related works are reviewed. In Section 3, the AFDM system model is described, and the embedded pilot channel estimation method is discussed. In Section 4, the details of the proposed module, data collection, and training are described. In Section 5, comparative experiments and the complexity of different methods are considered. Finally, the conclusion is drawn in Section 6.

2. Related Works

2.1. Traditional Channel Estimation in AFDM

Several channel estimation (CE) schemes have been proposed. An embedded pilot-aided scheme was introduced in [14], utilizing guard bands to isolate the pilot from data symbols, thereby mitigating pilot-data interference. Exploiting the unique structure of AFDM, the Embedded Pilot-Aided Diagonal Reconstruction (EPA-DR) method [15] was proposed, which points out that the full AFDM effective channel matrix can be reconstructed from a single column using a deterministic transformation factor. Other approaches, such as the Embedded Pilot-Aided Approximated Maximum Likelihood (EPA-AML) [16], attempt to estimate channel parameters by brute searching for Doppler shifts. Additionally, Compressed Sensing (CS) based methods, including Orthogonal Matching Pursuit (OMP) and off-grid estimation [17,18].

Despite their commendable performance in ideal settings, these methods exhibit significant susceptibility to fractional Doppler shifts and inter-Doppler interference (IDoI), both of which lead to a severe degradation in estimation accuracy. EPA-DR relies on hard thresholds to detect paths. When fractional Doppler occurs, the energy of a single path spreads beyond its principal location. Because fixed thresholds lack the flexibility to adapt to varying Doppler shifts, spreading energy is often filtered out, leading to a significant loss of channel information. EPA-AML operates on the a priori assumption of distinct path delays and a known number of propagation paths. However, when paths share the same delay, IDoI significantly impairs the estimation accuracy of delay and Doppler parameters. This limitation inherently compromises the practical viability of the method in complex scattering environments.

2.2. Deep Learning-Based Channel Estimation in OTFS

The severe path energy dispersion and interference induced by fractional Doppler shifts are not unique to AFDM; they also pose significant challenges in OTFS systems. To overcome the limitations of traditional algorithms in these environments, Deep Learning (DL) has been applied for OTFS channel estimation. For instance, some approaches formulate the task as an image recovery problem, utilizing video denoising networks to exploit correlations among multiple coarse channel estimates for enhanced accuracy [19]. Deep Residual Attention Network (DRAN) employs cascaded attention mechanisms to serve as adaptive filters that explicitly suppress interference from data and noise around embedded pilots [20]. Inspired by the structural similarity between the effective channel matrices of OTFS and AFDM [21], these successes demonstrate that DL models offer a powerful alternative for AFDM.

2.3. Motivation

In AFDM channel estimation, each position in the truncated received signal inherently corresponds to a specific delay-Doppler pair. Under integer-Doppler conditions, channel parameters can be efficiently estimated by simply detecting the presence of a signal at these positions. However, when fractional Doppler shifts occur, the energy of a single path disperses infinitely across the domain, leading to severe path ambiguity that renders direct parameter extraction infeasible.

Motivated by this physical mechanism, our network first expands the dimension of the received signal to mitigate inter-path interference and alleviate initial path ambiguity. Subsequently, the received signal sequence is divided into patches, effectively grouping specific positions of the received signal that inherently correspond to delay-Doppler pairs. By processing these position-specific segments, the patch embedding layer implicitly models the underlying path parameters from the dispersed energy within each patch. The self-attention mechanism is then employed to effectively combat path ambiguity by capturing global correlations among the dispersed energy components. Finally, the network leverages the inherent diagonal reconstructability of the AFDM effective channel by utilizing transposed convolutions to simulate this reconstruction behavior. The network seamlessly maps the implicitly modeled path parameters back into the full AFDM effective channel matrix.

2.4. Notation and Normalization

To maintain consistency and clarity throughout the system model, the key physical variables and their normalized discrete-time counterparts are summarized in

Table 1. Summary of key notations and units.

Variable	Description	Unit/Scale
B	System bandwidth	Hz
N	Number of subcarriers	Dimensionless
$T_s$	Sampling interval	Seconds (s)
$T_{sym}$	AFDM symbol duration ($N{T}_s$)	Seconds (s)
$\Delta f$	Subcarrier spacing/Doppler resolution	Hz
${\tau }_i$	Physical path delay	Seconds (s)
$f_i$	Physical Doppler shift	Hz
$l_i$	Normalized path delay (${\tau }_i/T_s$)	Samples (Integer)
$v_i$	Normalized Doppler shift ($f_i/\Delta f$)	Dimensionless
$c_1,c_2$	AFDM chirp parameters	Dimensionless

. We define the bandwidth as B (Hz) and the total number of subcarriers as N. The sampling interval (the time between consecutive samples) is defined as $T_s=1/B$ (seconds). Consequently, the total duration of one AFDM symbol (excluding the prefix) is $T_{sym}=N{T}_s$ (seconds). The subcarrier spacing, which also represents the Doppler resolution, is given by $\Delta f=1/T_{sym}=B/N$ (Hz).

Physical parameters, such as path delay ${\tau }_i$ (seconds) and Doppler shift $f_i$ (Hz), are normalized to these fundamental units. Specifically, the normalized delay $l_i={\tau }_i/T_s$ is assumed to be an integer based on the high-bandwidth sampling assumption, and the normalized Doppler shift $v_i=f_i/\Delta f=f_{i}N{T}_s$ includes both integer ${\alpha }_i$ and fractional $a_i$ components.

3. Review of the AFDM System Model and Channel Estimation

In this section, the AFDM system model and the embedded pilot-aided channel estimation method are briefly introduced. The block diagram of AFDM systems is demonstrated in Figure 1.

3.1. System Model of AFDM

In the following, we consider a Single-Input, Single-Output (SISO) AFDM system. The modulation and demodulation schemes, along with the channel model of AFDM, are presented.

3.1.1. Modulation

All transmit symbols are in the Discrete Affine Fourier Transform (DAFT) domain, let $\mathbf{x}={[x\left[0\right],\dots ,x[N-1]]}^T\in {\mathbb{A}}^{N\times 1}$ denote the vector of N phase shift keying (PSK)/quadrature amplitude modulation (QAM) symbols in the DAFT domain, which includes pilots, data, and guard symbols, and $\mathbb{A}$ represent the modulation alphabet. Following the serial-to-parallel (S/P) operation, the discrete-time domain representation of the signal $\mathbf{s}\in {\mathbb{C}}^{N\times 1}$ is obtained via the Inverse Discrete Affine Fourier Transform (IDAFT), which can be specifically formulated as

$$s\left[n\right]=\sum _{m=0}^{N-1}x\left[m\right]·{\displaystyle \frac{1}{\sqrt{N}}}{e}^{j2\pi \left(c_1{n}^2+c_2{m}^2+{\displaystyle \frac{nm}{N}}\right)}$$

(1)

where n and m denote the time and DAFT domains indices, respectively, and N represents the number of subcarriers. The AFDM signal has a bandwidth $B={\displaystyle \frac{1}{T_s}}$; subcarrier spacing (Doppler resolution) $\Delta f={\displaystyle \frac{B}{N}}={\displaystyle \frac{1}{N{T}_s}}$. $c_1$, $c_2$ are adjustable parameters of AFDM; and $c_1$ determines the slope of chirp. To achieve the optimal diversity order of the given Delay-Doppler profile of the environment [16], $c_2$ can be an arbitrary irrational number or a rational number sufficiently smaller than ${\displaystyle \frac{1}{2N}}$, and $c_1$ should be tuned as follows

$$c_1={\displaystyle \frac{2({\alpha }_{\max }+k_v)+1}{2N}}$$

(2)

where $k_v$ means additional guard symbols between pilot and data, and ${\alpha }_{\max }$ represents the max normalized integer Doppler, which will be further explored later in this paper, then the modulation of AFDM in Equation (1) can be rewritten in matrix form as

$$\mathbf{s}={\mathbf{A}}^H\mathbf{x}={\mathbf{\Lambda }}_{c_1}^H{\mathbf{F}}^H{\mathbf{\Lambda }}_{c_2}^H\mathbf{x}$$

(3)

where $\mathbf{A}={\mathbf{\Lambda }}_{c_2}\mathbf{F}{\mathbf{\Lambda }}_{c_1}\in {\mathbb{C}}^{N\times N}$ represents the DAFT transform matrix, and $\mathbf{F}$ represents the N-point discrete Fourier transform (DFT) matrix with entries $e^{-j2\pi mn/N}/\sqrt{N}$, $m,n\in \{1,\dots ,N\}$, and ${\mathbf{\Lambda }}_c$ is a frequency adjustment diagonal matrix given by

$${\mathbf{\Lambda }}_c=\mathrm{diag}\left({\left[1,e^{-j2\pi c},\dots ,e^{-j2\pi c{(N-1)}^2}\right]}^T\right).$$

(4)

A chirp-periodic prefix (CPP) should be added to $\mathbf{s}$ before transmitting; this prefix, analogous to the cyclic prefix (CP) in OFDM systems, aims to mitigate inter-symbol interference caused by multipath propagation. After parallel-to-serial conversion, the signal is transmitted in the time domain.

3.1.2. Channel Model

The impulse response of a doubly selective fading channel at time n and delay l can be represented as

$$g_n\left(l\right)=\sum _{i=1}^P{h}_i{e}^{-j2\pi f_{i}n}\delta (l-{\tau }_i/T_s),$$

(5)

where P is the number of propagation paths; $\delta (·)$ is the Dirac delta function; $h_i$, $f_i$, and ${\tau }_i$ are the complex path coefficients, Doppler shift, and the delay of the i-th propagation path, respectively.

The discrete-time received signal $\mathbf{r}={[r\left[0\right],\dots ,r[N-1]]}^T\in {\mathbb{C}}^{N\times 1}$, after propagating through the time-varying multipath channel, is characterized as follows

$$r\left[n\right]=\sum _{i=1}^P{h}_{i}s[n-l_i]e^{-j{\displaystyle \frac{2\pi }{N}}{v}_{i}n}+w\left[n\right],$$

(6)

where $w\left[n\right]$ represents the additive white Gaussian noise (AWGN) modeled by a complex normal distribution $\mathcal{CN}(0,N_0)$. $N_0$ is the variance of the AWGN, which corresponds to the noise power. Furthermore, $v_i$ signifies the Doppler shift associated with the i-th propagation path, normalized relative to the subcarrier spacing. Correspondingly, $l_i$ denotes the i-th path delay, normalized by the sampling interval $T_s$. More specifically, the normalized parameters for the i-th path are defined as

$$l_i={\displaystyle \frac{{\tau }_i}{T_s}},v_i=f_{i}N{T}_s,$$

(7)

where $l_i\in [0,l_{max}]$ is the normalized delay, with $l_{max}$ denoting the maximum delay spread of the channel. The normalized Doppler shift $v_i\in [-v_{max},v_{max}]$ can be decomposed into an integer component ${\alpha }_i\in [-{\alpha }_{max},{\alpha }_{max}]$ and a fractional component $a_i\in (-1/2,1/2]$, $v_i={\alpha }_i+a_i$.

After discarding the CPP, the received signal $\mathbf{r}$ in Equation (6) can be rewritten in matrix form as

$$\mathbf{r}=\mathbf{H}\mathbf{s}+\mathbf{w}$$

(8)

where $\mathbf{w}\sim \mathcal{CN}(\mathbf{0},N_0\mathbf{I})$, while the effective channel matrix $\mathbf{H}$ is given by $\mathbf{H}={\sum }_{i=1}^P{h}_i{\mathbf{\Gamma }}_{{\mathrm{CPP}}_i}{\mathbf{\Delta }}_{f_i}{\mathbf{\Pi }}^{l_i}$. Here, $\mathbf{\Pi }$ represents the forward cyclic-shift matrix:

$$\mathbf{\Pi }={\left[\begin{array}{cccc}0& \dots & 0& 1\\ 1& \dots & 0& 0\\ ⋮& \ddots & \ddots & ⋮\\ 0& \dots & 1& 0\end{array}\right]}_{N\times N}$$

(9)

where ${\mathbf{\Delta }}_{f_i}\triangleq \mathrm{diag}(e^{-j2\pi f_{i}n},n=0,1,\dots ,N-1)$ and ${\mathbf{\Gamma }}_{{\mathrm{CPP}}_i}$ is an $N\times N$ diagonal matrix:

$${\mathbf{\Gamma }}_{{\mathrm{CPP}}_i}=\mathrm{diag}\left\{\begin{array}{cc}{e}^{-j2\pi c_1(N^2-2N(l_i-n))}\hfill & n<l_i\hfill \\ 1\hfill & n\ge l_i\hfill \end{array}\right.$$

(10)

3.1.3. Demodulation

At the receiver side, the DAFT domain demodulated symbols are obtained by the following matrix

$$\mathbf{y}={\mathbf{\Lambda }}_{c_2}\mathbf{F}{\mathbf{\Lambda }}_{c_1}\mathbf{r}={\mathbf{H}}_{\mathrm{eff}}\mathbf{x}+\tilde{\mathbf{w}}$$

(11)

where the effective channel matrix is defined as ${\mathbf{H}}_{\mathrm{eff}}\triangleq \mathbf{A}\mathbf{H}{\mathbf{A}}^H$, and the transformed noise vector is given by $\tilde{\mathbf{w}}=\mathbf{A}\mathbf{w}$. Given that $\mathbf{A}$ is a unitary operator, the statistical characteristics of $\tilde{\mathbf{w}}$ remain identical to those of the original noise vector $\mathbf{w}$.

3.1.4. Input–Output Relation

To explicitly characterize the individual contribution of each propagation path, (11) can be reformulated as follows:

$$\mathbf{y}=\sum _{i=1}^P{h}_i{\tilde{\mathbf{H}}}_i\mathbf{x}+\tilde{\mathbf{w}}$$

(12)

where ${\tilde{\mathbf{H}}}_i\triangleq \mathbf{A}{\mathbf{\Gamma }}_{{\mathrm{CPP}}_i}{\mathbf{\Delta }}_{f_i}{\mathbf{\Pi }}^{l_i}{\mathbf{A}}^H$. It can be shown that

$$\begin{array}{cc}\hfill {\tilde{H}}_i[p,q]=& {\displaystyle \frac{1}{N}}{e}^{j{\displaystyle \frac{2\pi }{N}}\left(N{c}_1{l}_i^2-q{l}_i+N{c}_2\left(q^2-p^2\right)\right)}\hfill \\ \hfill & \times \sum _{n=0}^{N-1}{e}^{-j{\displaystyle \frac{2\pi }{N}}\left(\left(p-q+v_i+2N{c}_1{l}_i\right)n\right)}\hfill \end{array}$$

(13)

The summation term in (13) varies significantly depending on whether the normalized Doppler shift takes an integer value. In light of this, the following two conditions are investigated.

(1) Integer Doppler Shifts: Specifically, when the Doppler shift is restricted to an integer value (i.e., $a_i=0$), the normalized shift simplifies to $v_i={\alpha }_i$. Under this condition, the i-th sub-channel matrix exhibits a sparse structure, where its elements $H_i[p,q]$ are non-zero only at specific indices, as expressed by

$${\tilde{H}}_i[p,q]=\left\{\begin{array}{cc}{e}^{j{\displaystyle \frac{2\pi }{N}}\left(N{c}_1{l}_i^2-q{l}_i-N{c}_2(q^2-p^2)\right)}\hfill & q={(p+lo{c}_i)}_N\hfill \\ 0\hfill & \mathrm{otherwise}\hfill \end{array}\right.$$

(14)

(2) Fractional Doppler Shifts: It is evident from Figure 2 that under fractional normalized Doppler shifts, the structure of the equivalent path matrix deviates from its original form of P diagonals. An infinite energy leakage is observed around these diagonals, yet the dominant portion of the energy is still confined to the neighborhood of the diagonals.

3.2. Embedded Pilot-Aided Channel Estimation

From Equation (11), ${\mathbf{H}}_{\mathrm{eff}}$ only have $Q+1$ non-zero elements per column under Integer Doppler condition, $Q=\left(l_{max}+1\right)\left(2\left({\alpha }_{max}+k_v\right)+1\right)-1$. As a result, Q guard symbols are set between the pilot symbol and data symbols so that the channel estimation can be done with less interference from the data symbols. In this paper, we place a single pilot at the Tx frame; the channel estimation can be performed by means of the subsequent pilot placement:

$$x\left[n\right]=\left\{\begin{array}{cc}{x}_{\mathrm{pilot}},\hfill & n=Q\hfill \\ 0,\hfill & 0\le n<Q,Q<n\le 2Q\hfill \\ x_{\mathrm{data}},\hfill & else\hfill \end{array}\right.$$

(15)

Due to the structure of ${\mathbf{H}}_{\mathrm{eff}}$, $Q+1$ elements of Rx frame that are most relevant to the transmitted pilots are considered for the channel estimation. These symbols can be derived as

$${\mathbf{y}}_{\mathrm{E}}={\mathbf{H}}_{\mathrm{eff},\mathrm{E}}{\mathbf{x}}_{\mathrm{E}}+{\tilde{\mathbf{w}}}_{\mathrm{E}}$$

(16)

where E stands for channel estimation. Define ${\mathbf{T}}_{\mathrm{x}}={\left[{\mathbf{I}}_N\right]}_{in{d}_{\mathrm{x}}}$, and ${\mathbf{T}}_{\mathrm{y}}={\left[{\mathbf{I}}_N\right]}_{in{d}_{\mathrm{y}}}$ to extract specific rows or columns from matrix, $in{d}_x=[0:2Q]$ and $in{d}_y=[({\alpha }_{max}+k_v):Q+({\alpha }_{max}+k_v)]$, ${\mathbf{x}}_{\mathrm{E}}={\mathbf{T}}_{\mathrm{x}}\mathbf{x}$, ${\mathbf{y}}_{\mathrm{E}}={\mathbf{T}}_{\mathrm{y}}\mathbf{y}$ and ${\mathbf{H}}_{\mathrm{eff},\mathrm{E}}={\mathbf{T}}_{\mathrm{y}}{\mathbf{H}}_{\mathrm{eff}}{\mathbf{T}}_{\mathrm{x}}^H$.

In [16], $l_i$, $f_i$, and $h_i$ can be estimated through the following equation, which has good performance if the paths have different delays.

$$l({\mathbf{y}}_E\mid \mathit{\theta },{\mathbf{x}}_E)={∥{\mathbf{y}}_E-{\mathbf{H}}_{\mathrm{eff},E}{\mathbf{x}}_E∥}^2.$$

(17)

According to the diagonal reconstructability [15] of the AFDM channel matrix, the whole ${\mathbf{H}}_{\mathrm{eff}}$ can be reconstructed “diagonally” as long as we acquire an arbitrary column of ${\mathbf{H}}_{\mathrm{eff}}$, which can be expressed as

$$\mathbf{H}\left[{(m+1)}_N,{\left(m^{\prime }+1\right)}_N\right]=\mathcal{T}\left(l,m,m^{\prime }\right)\mathbf{H}\left[m,m^{\prime }\right]$$

(18)

where $\mathcal{T}\left(l,m,m^{\prime }\right)$ is the transform factor between $\mathbf{H}\left[m,m^{\prime }\right]$ and $\mathbf{H}\left[{(m+1)}_N,{\left(m^{\prime }+1\right)}_N\right]$, which is just relevant to the location and the delay. In our proposed scheme, the transformation factor can be adaptively learned through network training rather than being a fixed constant value as in [15].

4. Proposed Scheme

In this section, a deep learning-based channel estimation with an attention network is proposed to extract the pilot feature from the received signal and then reconstruct ${\widehat{\mathbf{H}}}_{\mathrm{eff}}$.

4.1. 1D Process Module

The 1D process module is designed to suppress interference from data and noise and to extract pilot features across different scales. As shown in Figure 3, this module consists of one linear layer, three 1D convolutions with different kernel sizes as 1, 3, and 5, and one convolutional layer.

The input of this module ${\mathbf{Y}}_{\mathrm{E}}$ can be denoted as

$${\mathbf{Y}}_E=\left[{\mathbf{y}}_{E,R},{\mathbf{y}}_{E,I}\right]\in {\mathbb{R}}^{2\times (Q+1)\times 1}$$

(19)

where ${\mathbf{y}}_{E,R}$ and ${\mathbf{y}}_{E,I}$ represent, respectively, the real and imaginary parts of the truncated received signal ${\mathbf{y}}_E$, and Q means guard symbols between pilot and signal, which is shown in Section 3.2.

The input ${\mathbf{Y}}_{\mathrm{E}}$ is first projected into a higher-dimensional feature space through a linear layer to mitigate the inter-path interference in the original signal, thereby enhancing the discriminability of different paths for subsequent processing stages. To adaptively capture diverse channel impulse patterns in the DAFT domain, the signal is subsequently passed through three convolutional layers with varying kernel sizes to extract multi-scale pilot features.

Specifically, the small-scale kernels are designed to maintain high-resolution localization for paths with integer Doppler shifts or low-velocity scenarios, where the path energy remains highly concentrated. Conversely, the large-scale kernels are employed to capture the extended energy dispersion patterns characteristic of high-mobility scenarios with path energy leakage. This multi-scale architecture allows the 1D processing module to effectively gather dispersed energy across different Doppler spread levels.

All convolutional layers are designed to preserve the spatial dimensions of the feature maps while modifying only the channel dimensions in this paper. The dimensional configurations of the linear layers and the channel settings of the convolutional layers are illustrated in Figure 3. And the output of the 1D process module can be expressed as follows:

$${\mathbf{S}}_{1D}=f_{1D}\left({\mathbf{Y}}_E\right)\in {\mathbb{R}}^{8\times \left(32\right)\times 1}$$

(20)

4.2. Transform-Based Module

In this module, the self-attention mechanism [22] is employed to capture path energy dispersion under varying fractional Doppler conditions. By exploiting the inherent diagonal reconstruction properties of the AFDM effective channel [15], the proposed module effectively completes high-precision channel estimation. As shown in Figure 3, ${\mathbf{S}}_{1D}$ first undergoes segmentation with different patch sizes. Then, these patches are linearly projected into a sequence of token embeddings and augmented with standard positional encodings to retain sequential information. Subsequently, the multi-head self-attention (MHA) mechanism within the Transformer blocks computes global correlations across all tokens. This global receptive field enables the model to capture long-range dependencies and aggregate dispersed path energy across the entire sequence.

As illustrated in Figure 4, the Transformer block begins with a LayerNorm layer, followed by a multi-head self-attention layer that incorporates a residual connection. In the following, a second LayerNorm layer and an MLP are applied, where an additional residual connection is utilized to facilitate gradient flow.

The architecture of the MHA is shown in Figure 5. Notably, the multi-head mechanism linearly projects the input into multiple independent representation subspaces, allowing different heads to concurrently capture diverse energy dispersion patterns. The calculation formula for the attention score is as follows.

$$Attention(\mathbf{Q},\mathbf{K},\mathbf{V})=Softmax\left({\displaystyle \frac{\mathbf{Q}{\mathbf{K}}^{\mathsf{T}}}{\sqrt{d}}}\right)\mathbf{V}$$

(21)

At the end of the module, a linear projection and a reshape operation are applied after the Transformer blocks to map the extracted features into a higher dimension and reshape them into 2D structural representations. These representations then undergo progressive upsampling via Transposed Conv2D layers. Meanwhile, subsequent convolutional layers are utilized to adaptively reweight the upsampled features, thereby fine-tuning the transformation factors across different columns to ensure accurate reconstruction of the effective channel matrix. The detailed step-by-step operations and exact dimensional transitions of this entire process are summarized in Algorithm 1.

Algorithm 1 Forward Pass of the Proposed AFDM Channel Estimation Network

Require: Truncated received signal $\mathbf{X}\in {\mathbb{R}}^{B\times 2\times 15}$ Ensure: Estimated effective channel matrix ${\widehat{\mathbf{H}}}_{\mathrm{eff}}\in {\mathbb{R}}^{B\times 2\times 128\times 128}$   1: % Phase 1: 1D Feature Processing   2: ${\mathbf{X}}_1\leftarrow {\mathrm{Linear}}_{15\to 32}\left(\mathbf{X}\right)$ {Initial dimension expansion, Shape: $B\times 2\times 32$}   3: ${\mathbf{X}}_2\leftarrow 1\mathrm{D}\text{-}\mathrm{process}\left({\mathbf{X}}_1\right)$ {Shape: $B\times 8\times 32$}   4: Phase 2: Vision Transformer (ViT) Blocks   5: {Three parallel ViT branches processing the same 1D features}   6: ${\mathbf{F}}_{\mathrm{v}1}\leftarrow \mathrm{ViT}\_\mathrm{1D}({\mathbf{X}}_2,p=4,\mathrm{upsample}=2)$ {Branch 1, Shape: $B\times 8\times 64\times 64$}   7: ${\mathbf{F}}_{\mathrm{v}2}\leftarrow \mathrm{ViT}\_\mathrm{1D}({\mathbf{X}}_2,p=2,\mathrm{upsample}=1)$ {Branch 2, Shape: $B\times 16\times 32\times 32$}   8: ${\mathbf{F}}_{\mathrm{v}3}\leftarrow \mathrm{ViT}\_\mathrm{1D}({\mathbf{X}}_2,p=1,\mathrm{upsample}=0)$ {Branch 3, Shape: $B\times 32\times 16\times 16$}   9: % Phase 3: Upsample 10: ${\mathbf{U}}_3\leftarrow {\mathrm{ConvTranspose}2\mathrm{D}}_{k=4,s=2}\left({\mathbf{F}}_{\mathrm{v}3}\right)$ {Shape: $B\times 16\times 32\times 32$} 11: ${\mathbf{C}}_2\leftarrow \mathrm{Concat}({\mathbf{U}}_3,{\mathbf{F}}_{\mathrm{v}2},\dim =1)$ {Shape: $B\times 32\times 32\times 32$} 12: ${\mathbf{U}}_2\leftarrow {\mathrm{ConvTranspose}2\mathrm{D}}_{k=4,s=2}\left({\mathbf{C}}_2\right)$ {Shape: $B\times 8\times 64\times 64$} 13: ${\mathbf{C}}_1\leftarrow \mathrm{Concat}({\mathbf{U}}_2,{\mathbf{F}}_{\mathrm{v}1},\dim =1)$ {Shape: $B\times 16\times 64\times 64$} 14: ${\mathbf{U}}_1\leftarrow {\mathrm{ConvTranspose}2\mathrm{D}}_{k=4,s=2}\left({\mathbf{C}}_1\right)$ {Shape: $B\times 8\times 128\times 128$} 15: % Phase 4: Final Feature Refinement 16: ${\mathbf{M}}_{\mathrm{out}}\leftarrow {\mathrm{Conv}2\mathrm{D}}_{1\times 1}\left({\mathbf{U}}_1\right)$ {Channel-wise refinement, Shape: $B\times 8\times 128\times 128$} 17: ${\widehat{\mathbf{H}}}_{\mathrm{eff}}\leftarrow {\mathrm{Conv}2\mathrm{D}}_{3\times 3}\left({\mathbf{M}}_{\mathrm{out}}\right)$ {Final output mapping, Shape: $B\times 2\times 128\times 128$} 18: return  ${\widehat{\mathbf{H}}}_{\mathrm{eff}}$

The key settings of each transformer block are shown in

Table 2. Parameters of the Transformer block.

Parameter	Vit-Third	Vit-Second	Vit-First
patch size	1	2	4
depth	2	2	2
dim	64	48	48
heads	8	8	8
mlp-dim	64	64	64
linear-output	256	256	256

. where patch size determines the division of input patches, and depth denotes the number of stacked layers within each Transformer block. The parameter dim represents the feature dimension of both the input and output of the Transformer block, and heads specifies the number of attention heads. The mlp-dim corresponds to the hidden dimension of the MLP illustrated in Figure 4. In addition, linear-output refers to the output dimension of a linear layer applied to expand the feature dimension from dim, serving as a preparatory step for the subsequent transformation into a 2D matrix.

4.3. Data Collection and TRAINING

All datasets are generated in MATLAB R2023a to provide a controlled environment for performance evaluation. The number of propagation paths P is uniformly distributed between 1 and 5. The delay for each path is randomly chosen from a uniform distribution within $[0,l_{max}]$. In practical wideband systems, the sampling frequency is sufficiently high, allowing the actual physical delays to be accurately approximated as integer multiples of the sampling interval, so the path delays are set as integers, consistent with the settings in [7,15,16].

The channel follows a uniform power-delay profile (PDP), where the complex path coefficients $h_i$ are generated as independent complex Gaussian random variables with zero mean and a variance of 1.

The Doppler shift of each path is generated using $v_i={\alpha }_{max}cos\left({\theta }_i\right)$, where ${\theta }_i$ is uniformly distributed over $[-\pi ,\pi ]$ to simulate realistic mobility.

The system signal-to-noise (SNR) is defined as ${\displaystyle \frac{E\left\{{\left|x\left(n\right)\right|}^2\right\}}{{\sigma }^2}}$, where ${\sigma }^2$ means the energy of noise and $E\{·\}$ is the expected operation. For the pilot, $10{log}_{10}{\displaystyle \frac{|x_p{|}^2}{E\left\{{\left|x\left(n\right)\right|}^2\right\}}}=20\mathrm{dB}$ is obeyed, and other settings can be found in

Table 3. Parameters of the AFDM system.

Parameter	Value
Carrier frequency $f_c$ (GHz)	28
Constellation	16-QAM
Max normalized Doppler	1
Max normalized delay	2
Bandwidth (MHz)	2
Sampling frequency	2 × Bandwidth
Guard symbols	$\left(l_{max}+1\right)\left(2\left({\alpha }_{max}+k_v\right)+1\right)-1$
$k_v$	1
Max UE speed (kmph)	600
Number of paths (p)	1–5

. The training, validation, and test datasets are generated at each SNR ranging from 0 dB to 25 dB with a 5 dB step size, containing 80,000, 20,000, and 20,000 samples per SNR level, respectively. Specifically, independent models are trained for each specific SNR level to ensure optimal estimation performance across different noise conditions. To train this network, we chose MSE as the loss function to improve performance.

$${\mathcal{L}}_{\mathrm{MSE}}={\displaystyle \frac{1}{N^2}}\sum _{i=1}^N\sum _{j=1}^N{\left|{\mathit{H}}_{\mathrm{eff}}[i,j]-{\widehat{\mathit{H}}}_{\mathrm{eff}}[i,j]\right|}^2$$

(22)

The proposed network is implemented using the PyTorch 2.6.0 framework on a desktop workstation equipped with an 10th Gen Intel(R) Core(TM) i9-10900K CPU (Intel Corporation, Santa Clara, CA, USA) and an NVIDIA RTX 2080 GPU (NVIDIA Corporation, Santa Clara, CA, USA). The network is trained from scratch for 300 epochs with a batch size of 128. To optimize the network, we employ the Adam optimizer with an initial learning rate of 0.001. Regarding internal configurations, we use GELU as the nonlinear activation function throughout the network, along with Layer Normalization (LayerNorm) in the ViT blocks. The network weights are initialized using PyTorch’s default methods (e.g., Kaiming initialization for convolutional layers). Furthermore, the exact input dimension shaping and the specific channel transitions across all intermediate layers are comprehensively detailed in Algorithm 1. The entire training process takes approximately a few hours.

5. Simulation and Results

In this section, we consider an SISO-AFDM system with N = 128 subcarriers. One pilot is inserted into the Q-th subcarrrier and the arrangement of pilot and data symbols is shown in Section 3.2, and $k_v=1$. For a fair comparison, we adopt the same dual-dispersion channel model as [15,16,18], the parameters are generated in the same manner as described in Data collection. The remaining AFDM parameters can be found in Table 3.

5.1. Simulation Results

In accordance with the configurations adopted in prior studies, the Doppler grid resolution for the OMP scheme is 0.2 [17], and the Doppler grid resolution for the EPA-AML is 0.1 [17]. To quantify the accuracy of various channel estimation techniques, the Normalized Mean Squared Error (NMSE) is employed as the primary performance metric, which is defined as follows:

$$\mathrm{NMSE}={\displaystyle \frac{{∥{\mathit{H}}_{\mathrm{eff}}-{\widehat{\mathit{H}}}_{\mathrm{eff}}∥}^2}{{∥{\mathit{H}}_{\mathrm{eff}}∥}^2}}$$

(23)

Figure 6 and Figure 7 illustrate the NMSE performance of the EPA-DR, OMP, EPA-AML, and the proposed network. In Figure 6, it is observed that EPA-DR exhibits the most inferior performance, primarily because its fixed-threshold mechanism fails to accommodate the path energy dispersion induced by fractional Doppler shifts. While all methods experience varying degrees of performance degradation due to the introduction of IDoI, as shown in Figure 7, the impact of EPA-AML is particularly pronounced, which is derived based on the assumption of the absence of IDoI. In contrast, our proposed method consistently achieves superior NMSE performance. This phenomenon underscores the advantages of data-driven adaptive strategies in modeling fractional Doppler effects in AFDM channel estimation. Furthermore, as shown in Figure 7, the comparative results involving a CNN-based scheme and an ablation study without the ViT blocks further validate the architectural effectiveness.

To further evaluate various channel estimation algorithms, the Minimum Mean Square Error (MMSE) equalizer [23] is employed for signal detection. As shown in Figure 8, our proposed method achieves the best BER performance compared with other methods in the presence of IDoI, which is very close to the ideal CSI. Although the BER curves of all evaluated methods are similar in the low-SNR region, our approach achieves a 6 dB improvement over the OMP algorithm at 25 dB SNR.

To verify the robustness and generalization of the model, we conducted experiments under different parameter settings. As illustrated in Figure 9a, we evaluate the NMSE performance with varying numbers of propagation paths. Although the training set comprises conditions with one to five paths, the proposed model demonstrates superior performance even in an unseen six-path scenario, which shows its remarkable generalization across diverse channel conditions. Furthermore, Figure 9b investigates the impact of pilot power bias. The results indicate that our model maintains a significant performance, exhibiting consistent resilience even when the power deviation reaches 2 dB.

5.2. Complexity Analysis

Unlike traditional methods, our approach directly outputs the effective channel matrix in an end-to-end manner. Therefore, the complexity of our model can be considered in two parts, localized feature extraction and channel reconstruction. In the feature extraction stage, the Transformer blocks process the effective region of the received signal of size Q, with $N_p\propto Q$ tokens and transform block feature dimension dim D shown in Table 2. The linear layer within the Transformer blocks and the feed-forward network contribute a complexity of $\mathcal{O}\left(Q{D}^2\right)$. The self-attention mechanism has complexity $\mathcal{O}\left(Q^{2}D\right)$. Consequently, the feature extraction stage complexity is $\mathcal{O}(Q{D}^2+Q^{2}D)$. In the following, the reconstruction stage employs transposed convolutions to upsample these features into the final $N\times N$ estimation matrix. These convolution operations result in a complexity of $\mathcal{O}\left(N^2\right)$.

Table 4. Complexity and computational efficiency comparison.

Method	Params	Latency (ms)	Theoretical Complexity
Proposed	246.7 k	19.3/9.7 †	$\mathcal{O}(Q{D}^2+Q^{2}D)+\mathcal{O}\left(N^2\right)$
EPA-DR	-	8.4	$\mathcal{O}\left(Q\right)+\mathcal{O}\left(N\right)$
OMP	-	12.3	$\mathcal{O}(T^3+T{N}_{O}Q+T^{2}Q)+\mathcal{O}\left(N^2\right)$
EPA-AML	-	18.9	$\mathcal{O}\left(Q{V}^P\right)+\mathcal{O}\left(N^2\right)$

^† The latency values represent one example execution time, respectively. ^† Proposed model implemented in Python 3.10.9 (GPU/CPU) and the traditional baselines in MATLAB R2023a (CPU). ^† The latency values for the Proposed method represent CPU and GPU (RTX 2080) for one example execution time, respectively.

demonstrates the complexity of other methods from the perspectives of feature extraction and reconstruction, and execution time of one sample. $N_O$ represents the column dimension of the measurement matrix, T is the number of iterations, and V represents the Doppler resolution. Our proposed methods have 246.7 k parameters and 109M FLOPs, which maintain a competitive scale compared to deep learning-based OTFS estimation frameworks [20]. Although the computational overhead of the proposed data-driven approach exceeds that of conventional methods, dedicated hardware effectively bridges this latency gap. This ensures that the execution time remains comparable to traditional methods. However, practical deployment must be evaluated against strict mobile device operational boundaries, including limited battery capacity, thermal constraints, and restricted memory bandwidth. And mobile CPUs struggle with the massive parallel computations required by our model, leading to latency bottlenecks. In contrast, GPUs are inherently optimized for such tasks. Therefore, dedicated hardware—such as mobile GPUs or Neural Processing Units (NPUs) are needed.

6. Conclusions

In this paper, we propose a deep learning network based on the Transformer architecture and the diagonal reconstruction property of the AFDM effective channel matrix to improve AFDM channel estimation performance under fractional Doppler shifts. By integrating ViT blocks with 1D feature processing module, the proposed network effectively captures dispersed path energy patterns without requiring any prior knowledge of the number of paths or distinct path delays, achieving up to 4 dB and 6 dB improvements in NMSE and BER respectively over existing traditional methods, making it well-suited for high-speed scenarios that require highly accurate estimation of the effective channel matrix with dedicated hardware. The model’s rationality, feasibility, and robustness are validated through comparative experiments, ablation studies, and evaluations under different parameter settings. However, this study has some limitations. First, the current results lack validation on a measured channel, so the model’s performance under real hardware imperfections remains uncertain. Second, the model sacrifices sensing capability, as it does not explicitly estimate path parameters (such as delay and Doppler shifts). Instead, the network implicitly models these parameters to directly reconstruct the effective channel matrix.

In future work, we will focus on two main directions. First, we plan to incorporate lightweight attention mechanisms and model compression methods, such as network pruning, to reduce the model’s parameters and accelerate inference speed. Second, since the feature extraction stage implicitly models path parameters, we aim to explore an end-to-end communication paradigm. By designing a unified network for joint channel estimation and signal detection, we can potentially eliminate the need for traditional standalone detectors.