## 2 Two-Stage Time-Domain Equalization Method
### Long Short-Term Memory  
In our proposed two-stage time-domain equalization module, we use an LSTM network to deal with nonlinear noise and other distortions in terahertz communication systems. LSTM networks have been widely used in time series processing research, and their excellent performance has been fully verified and recognized.  
Figure 2 illustrates the structure of the LSTM. As shown in the figure, at time step \(n\), the LSTM module consists of an input vector \(x_{n}\), current cell state \(C_{n}\), candidate cell state \(\bar{C}_{n}\), hidden state \(h_{n}\), and the forget gate \(f_{n}\), input gate \(i_{n}\), and output gate \(o_{n}\). The core working mechanism of the LSTM is to forget irrelevant information while retaining useful information. At each time step, the hidden state \(h_{n}\) is output. During the establishment  
Figure 1: Two-stage time-domain equalization structure. Architecture of the proposed two-stage time-domain equalization method. (**a**) THz communication system model and dataset construction. Two datasets under different frequencies and transmission distances are collected as training and validation data sources. (**b**) Training process of two-stage equalizer. Both the transmitted and received signals are upsampled before being processed by the first-stage long short-term memory (LSTM)-based equalizer, making it more effective in capturing nonlinear distortions. The output is then downsampled and refined by the second-stage LSTM equalizer. (**c**) In the inference stage, the trained Stage-1 and Stage-2 models are cascaded to process validation data and generate equalization results.  
of the two-stage time-domain equalization model, the input sequence \(X=\{x_{1},\ldots,x_{N}\}\) produces the output sequence \(H=\{h_{1},\ldots,h_{N}\}\) through the following equations and is processed iteratively.  
Forget gate:  
\[f_{n}=\sigma(W_{f}\cdot[h_{n-1},x_{n}]+b_{f}). \tag{1}\]  
Memory gate:  
\[i_{n}=\sigma(W_{i}\cdot[h_{n-1},x_{n}]+b_{i}). \tag{2}\]  
Temporary state of the cell:  
\[\tilde{C}_{n}=\tanh(W_{c}\cdot[h_{n-1},x_{n}]+b_{c}). \tag{3}\]  
Current state of the cell:  
\[C_{n}=f_{n}*C_{n-1}+i_{n}*\tilde{C}_{n}. \tag{4}\]  
Output gate:  
\[o_{n}=\sigma(W_{o}\cdot[h_{n-1},x_{n}]+b_{o}). \tag{5}\]  
State of hidden layer:  
\[h_{n}=o_{n}*\tanh(C_{n}). \tag{6}\]  
Here, \(W\) represents the weight matrix, \(b\) denotes the bias vector, and \(\sigma\) is the sigmoid activation function. \(h_{n-1}\) and \(C_{n-1}\) are the hidden state and cell state from the previous time step \(n-1\), while \(h_{n}\) and \(C_{n}\) are the newly generated states at the current time step, which will be passed on to the next time step \(n+1\). Therefore, LSTM not only accumulates information over a long period but also has the ability to reset old states through the forget gate, allowing it to restart the accumulation of information when needed.  
It should be noted that although we used an LSTM network to solve the nonlinear noise problem in terahertz communication systems in this study, the choice of network architecture is not the focus of this paper. The reason for choosing LSTM is that it performs well in most time series tasks [36; 37; 38; 39; 40; 41; 42]. We use the network model as a tool to address the distortion problem in terahertz communication systems. The innovation of this paper lies in the two-stage time-domain equalization training strategy, so this paper will not discuss the choice of network architecture in detail.