TransAlpha: Lightweight Design Empowers Stock Return Forecasting

1 Introduction

Stock return forecasting has advanced with AI. A solid forecasting method and stable prediction results can impact portfolio optimization and trading strategies. Practical stock forecasting methods can be categorized into time series forecasting models that use temporal sequences of individual stocks and multi-factor models that learn universal factor-return relationships across stocks and time via cross-sectional data. In recent years, Transformer [Vaswani et al., 2017] approaches rise as an alternative for time series forecasting. Despite their success in processing language tokens, Transformers fail to replicate breakthroughs in time series forecasting. Influential work such as DLinear [Zeng et al., 2023] is skeptical about whether complex Transformer architectures can outperform linear models. In particular, stock forecasting fundamentally differs from other time series forecasting tasks in the following aspects of challenges: (1) Weak signal-to-noise ratio: predictive features (alpha) signals are subtle, prune to noise and overfitting [Fama, 1970]. Signals are diluted if model has multiple compositions and residual connection blocks, which can decrease predictive power; (2) High dimensionality of cross-sectional data: quantitative trading strategies contain hundreds of alpha factors, introducing inter-dependencies [Li et al., 2024]; (3) Non-stationarity and volatility: market micro-structures shift rapidly; trained models cannot generalize across time. As a result, leading quantitative companies still rely on more robust and interpretable linear models to combine with existing trading strategies for stable profit, but invest heavily in Transformer and Large Language Model (LLM) research to explore possibilities improving model's predictive power.

Our work comes from a realistic production environment of a quantitative company: we demonstrate Transformer model theories can be transferred to a profitable pipeline using proprietary alpha factors in A-Share market. We propose TransAlpha, a lightweight transformer-based architecture for stock return forecasting. Our main contributions are three-fold:

• A systematic investigation that incorporates state-ofthe-art time series model architectures into multi-factor framework to examine their effectiveness for stock return forecasting.
• A domain-tailored Transformer architecture with three novel components: cross-sectional denoising, temporal attention gating, and smart pooling, all validated by real trading data.
• An industry-aligned pipeline that enables for training and back-testing custom models under a realistic setup with custom strategies.

These components address universal challenges such as noise in high-dimensional data, inefficient attention, and information loss when pooling. We balance academic reproducibility and proprietary privacy by providing the processed data files with anonymous alpha factors.

The remainder of this paper is organized as follows: Section 2 reviews related work in time series forecasting methods and past stock forecasting attempts. Section 3 presents the TransAlpha architecture. Section 4 presents research questions and experimental results. Section 5 concludes.

2 RelatedWork

2.1 Times Series Forecasting

Traditional time series methods are maximum likelihoodbased statistical models, such as ARIMA, and GARCH [Bollerslev, 1986]. Both methods and their variants assume stationarity, which are usually unrealistic, especially in intraday price movements. Later, a shift to machine learning introduced tree-based models such as XGBoost [Chen and Guestrin, 2016], LightGBM [Ke et al., 2017]. They can capture non-linearity without distribution or stationarity assumptions. Further, LSTMs [Hochreiter and Schmidhuber, 1997] addressed the vanishing gradient problem, and DeepAR [Salinas et al., 2020] pioneered in probabilistic forecasting. In addition, Convolutional Neural Network (CNN) models such as TCN [Bai et al., 2018] and TimesNet [Wu et al., 2022] use convolutional layers to capture local temporal patterns. Modern TCN [donghao and wang xue, 2024] refines TCN with much larger effective receptive fields. Later, Transformer [Vaswani et al., 2017] replaced LSTM's recurrent layers with self-attention, enabling long-range dependency modeling. Transformer adaptations for time series faced high computational complexity up to O(L2 ), where L is sequence length. Many variants are introduced to tackle this issue and attempt novel structures to solve long-sequence forecasting problem. Informer [Zhou et al., 2021] proposes probability sparse attention to approximate the full attention matrix. Autoformer [Wu et al., 2021] decomposes series to trend and seasonal components. PatchTST [Nie et al., 2023] frames time series as tokens. Pyraformer [Liu et al., 2022] proposes pyramidal attention to hierarchically compress temporal information. Crossformer [Zhang and Yan, 2023] establishes a Hierarchical Encoder-Decoder (HED) to use the information at different scales; iTransformer [Liu et al., 2024] simply applies the attention and feed-forward network on the inverted dimensions. New direction involves adapting LLM techniques [Zhang et al., 2024]. Time-MoE [Shi et al., 2025] used a sparse Mixture-of-Experts (MoE) design based on scaling laws [Kaplan et al., 2020] with 2.4 billion parameters but only activate a subset of experts.

2.2 Stock Return Prediction

The fluctuation of stock prices is influenced by a complex array of factors. Therefore, there is a need to utilize machine learning techniques [Zou et al., 2023]. Models must balance between predictive power and noise robustness. Successful adaptations can date back to support vector machines (SVMs) methods [Cao and Tay, 2003]. LSTMs and GRUs [Zhu et al., 2024] were also adapted. Graph WaveNet, [Park et al., 2025]) incorporates Graph Neural Networks (GNNs) and hypergraphs to capture higher-order stock relationships. Mamba [Gu and Dao, 2024] integrates selective State-Space models (SSMs) into a simplified end-to-end neural network architecture. MASTER [Li et al., 2024] introduced marketguided attention to prioritize macroeconomic signals over price fluctuations; MIGA [Yu et al., 2024] employed a MoE framework to partition experts by temporal frequency. Despite these advances, several gaps remain in the literature: (1) few works validate their models on high-dimensional alpha data from real production environment; (2) there are limited studies systematically studying performance of Transformer models in stock forecasting. (3) successful methods and insights are typically non-disclosed for proprietary concerns.

3 Methodology

3.1 ProblemFormulation

We formalize stock return forecasting as a supervised learning task. Let S = {s1, s2, . . . , sN} denote the A-share stock universe, where N represents the total number of stocks. Let T denote the set of 15-minute cross sections within the observation period. For each cross section t ∈ T, the alpha factor matrix at time t is denoted as Xt ∈ RN×D , where D = 400 is the dimension of alpha factors, and each row xt,s ∈ RD corresponds to the factor value vector of stock s at time t. The true return vector at time t is Yt ∈ RN , where each element yt,s represents the 15-minute forward raw return of stock s at time t.

We use a look-back window of fixed length L = 16, corresponding to 16 consecutive 15-minute cross sections. Then, we construct a sequence of factor matrices {Xt-L+1,Xt-L+2, . . . ,Xt}, where each matrix inside describes the cross-sectional distribution of factor values across all stocks at a t. The forecasting model maps this temporal sequence of cross-sectional factor data to a predicted return vector ˆYt+1 ∈ RN for the next 15-minute cross section:

where f denotes the model, and Θ is the set of trainable parameters. The observation period T is partitioned into nonoverlapping subsets of train, validation, and test segments.

3.2 Model

We propose TransAlpha for intraday cross-sectional stock return forecasting in the A-share market. The model addresses three challenges: (1) idiosyncratic noise in alpha factor signals, (2) time-step relevance heterogeneity in intraday trends, and (3) signal dilution induced by over-engineering. We formalize the model following section 3.1.

Let X = {Xt-L+1,Xt-L+2, . . . ,Xt} denote the input to TransAlpha. X is a temporal sequence of cross-sectional alpha factor matrices, where X ∈ RN×D . The output is a predicted 15-minute forward return vector ˆYt+1 ∈ RN . To present the three core modules of the model, we define the intermediate variables first. Let Xdenoise ∈ RN×L×D be the denoised factor, Z ∈ RN×L×d be the projected feature, and Zgated ∈ RN×L×d be the time-weighted feature, where d is the model's hidden dimension. We use d = 64 for all baseline models. Let Zpool ∈ RN×2d denote the aggregated features, and we define an output head that maps the aggregated features to predicted returns. The end-to-end workflow is demonstrated in Figure 1

Cross-Sectional Denoising The Cross-Section Denoising (CSD) module separates stock-specific alpha signals from market-wide systematic beta signals, with adaptive scaling to suppress noise while preserving predictive components. For an input X ∈ RN×L×D , let β be the market-wide signal, and α be the stock-specific residual. They are defined and obtained by

where β ∈ R1×L×D and α ∈ RN×L×D . We introduce a fixed scalar hyperparameter τ = 10-3 to partition the stockspecific alpha signal into noisy alphas (|α| ≤ τ) and nonnoisy, predictive alphas (|α| > τ). Further, let sα, sβ ∈ R1×1×D be factor-wise trainable scaling tensors. sα is applied exclusively to the non-noisy components α. During training, sα is optimized to amplify factors with high predictive power and dampen factors with weak signals, without affecting noisy components. Similarly, sβ is applied uniformly to the market-wide systematic signal β. Training optimizes sβ to down-weight overly dominant systematic trends and upweight stable, predictive market signals, ensuring a balanced combination of idiosyncratic and systematic components for return forecasting. We then combine both components and obtain the denoised input:

where 1(·) is the indicator function and ⊙ denotes elementwise multiplication.

Temporal Attention Gating The Temporal Attention Gating (TAG) module uses relevance to return predictions to weight time steps so that model can capture intraday trend with minimal complexity. The denoised sequence Xdenoise is projected to model dimension d to yield Z via a feed-forward block.

where Linear(X) =WX+ b is a linear layer with trainable parameters W, b. We then transformed Z to Zgated = Z⊙wt. The attention time-step weights wt are computed as

where wt sums to 1 to ensure signal weights are normalized. TAG performs element-wise multiplication of attention weights with the entire temporal sequence to generate Zgated, while traditional approaches only applies gating to the final hidden state. Meanwhile, TAG replaces the complex multi-head attention with a single-head scalar-scoring attention layer and removes positional encoding for further simplification.

Smart Pooling and Output Head The Smart Pooling (SP) component aggregates time-weighted features into fixeddimensional vectors by combining the most recent time step and max pooled peak signals with extreme volatility. Let Zgated ∈ RN×L×d be output from TAG module, we define Zlast as the most recent time step, Zmax as the max-pooled feature, and Zpool as the concatenated features, and denote

ˆYt+1 as the predicted returns. We present Zlast and Zmax using standard tensor indexing, where the first dimension is individual stocks, second dimension is intraday time steps and third dimension is model's hidden feature.

where MLP denotes a two-layer perceptron, and Concat[· , · ] denotes feature concatenation along the last dimension. The pooling used direct concatenation instead of average or sum in traditional pooling methods that lead to information loss.

3.3 Training

We design a multi-component loss function specifically adapted to intraday cross-sectional stock return forecasting. The loss is designed to balance regression performance, linear correlations between predicted and actual returns, relative rank accuracy, and regularization to penalize parameter size.

where λ1, . . . , λ5 are non-negative weighting coefficients. Let ˆYt+1 = {ˆ yt+1,s}s∈S and Yt+1 = {yt+1,s}s∈S be the vectors of predicted and true forward returns for time t + 1, respectively, and ¯ˆ yt+1, ¯ yt+1 be their cross-sectional means. We elaborate on each loss component as follows. First, the quantile ranking loss Lrank penalizes misalignment between predicted and true rankings of top and bottom stocks by enforcing a minimum predicted return gap between the top 20% and bottom 20% of stocks, sorted by true returns, with a margin m = 0.1, formulated as:

where ˆ ytop t+1,s and ˆ ybottom t+1,s are predicted returns for top/bottom quantile stocks, and E[·] denotes the expectation over training samples. Next, the Huber loss Lhuber robustifies training against extreme returns by balancing mean squared error (MSE) for small prediction errors and mean absolute error (MAE) for outliers. Let ϵt+1,s = ˆ yt+1,s - yt+1,s denote the prediction error for stock s, the loss with threshold δ = 1.0 is:

The correlation loss term Lcorr minimizes 1 - ρpearson, where ρpearson is the pearson correlation coefficient. We denote auxiliary terms for cross-covariance and the product of standard deviations:

such that Lcorr = 1 - Ccov/Cstd. The product correlation loss Lprod jointly maximize both correlations. Given a crosssectional vector of values V = {vs}s∈S, the rank operator maps V to a rank vector by rank(V) = {rs}s∈S , where rs is the ordinal rank of vs in V (1 smallest, N largest). The operator converts continuous return values into ordinal ranks. Consequently, ρpearson = Ccov/Cstd and

Then, the product loss is Lprod = 1 - ρpearson · ρspearman.

Finally, the regularization loss Lreg = λreg∥θ∥22 (with λreg = 10-5 ) prevents overfitting to training data noise, where ∥ · ∥2 denotes the ℓ2 norm. The weighting coefficients for the total loss function are chosen to prioritize ranking and correlation objectives. We set small weights for huber and regularization components. The final utilized weights are determined empirically: λ1 = 0.4, λ2 = 0.05, λ3 = 0.4, λ4 = 0.1, λ5 = 0.05.

3.4 Inference

After obtaining the trained model, we map the test file's crosssectional sequence, with input batch size B, to return predictions in the next time period. For each cross section t, we maintain a sliding look-back window of length L and apply two preliminary filters before feeding the batch.

Look-back Completeness Filter The completeness filter discards sequences which do not contain enough valid time steps (< L) in a look-back window. Formally, the filter retains only sequences satisfying:

where ∅ represent an empty factor matrix without valid stocks at time τ.

NaNElimination Filter For each time step τ in a complete look-back window, the NaN elimination filter removes rows with missing or non-finite factor values, typically caused by trading halts or data feed errors. Let Xτ (i, j) denote the j-th factor value of the i-th stock at time τ. The cleaned crosssectional matrix ˜Xτ ∈ R ˜Nt×D , where ˜Nt ≤ Nt is the number of valid stocks at time τ) is defined as,

where Xτ(i, :) denotes the i-th row of Xτ , and Rfin = R \ {NaN,±∞} denotes the set of finite real numbers. The filter ensures all elements of ˜Xτ are valid real numbers, maintaining numerical stability of forward pass. The filtered sequence ˜Xt = { ˜Xt-L+1, . . . , ˜Xt} is then reshaped into a

Algorithm 1 TransAlpha Forward Pass

• 1: Step 1: Cross-Sectional Denoising
• 2: β = 1B ∑Bi=1 Xi ∈ R1×L×D
• 3: α =X-β ∈ RB×L×D
• 4: αscaled = α(sα ⊙1(|α| > τ) +1(|α| ≤ τ))
• 5: βscaled = β ⊙ sβ
• 6: Xdenoise = αscaled +βscaled ∈ RB×L×D
• 7: Step 2: Input Projection toModel Dimension
• 8: Z1 = Linear(Xdenoise) ∈ RB×L×dmodel
• 9: Z2 = LayerNorm(Z1) ∈ RB×L×dmodel
• 10: Z3 = ReLU(Z2) ∈ RB×L×dmodel
• 11: Step 3: Temporal Attention Gating
• 12: wt = Softmax(Linear(Z3)) ∈ RB×L×1
• 13: Zgated = Z3 ⊙wt ∈ RB×L×dmodel
• 14: Step 4: Smart Temporal Pooling
• 15: Zlast = Zgated[:,-1, :] ∈ RB×dmodel
• 16: Zmax =maxLt=1 Zgated[:, t, :] ∈ RB×dmodel
• 17: Zpool = Concat(Zlast,Zmax) ∈ RB×2dmodel
• 18: Step 5: Output Head
• 19: ˆY =MLP(Zpool) ∈ RB×1
• 20: Return: ˆY

three-dimensional tensor X compatible with the TransAlpha model input size, which is then feeded into the forward pass workflow in Algorithm 1.

3.5 Convergence Analysis

For simplicity, we sketch TransAlpha's convergence to the true return Yt+1 under three mild assumptions. (1) Multicomponent loss L is Lipschitz continuous. (2) Model parameters (sα, sβ,W) are bounded via Lreg in a compact parameter space Θ. (3) Training uses mini-batch Adamwith decreasing learning rate ηk ∝ 1/√k. The expected loss satisfies

where C > 0 is a constant, Θ∗ is optimal parameter. As epochs k → ∞, ∥ ˆYt+1(Θk) - Yt+1∥2 → 0.

3.6 Complexity Analysis

Let h be the number of attention heads in baseline models, K be TCN kernel size, G be crossformer's segment size. TransAlpha's time complexity is derived by calculating floatpoint operations. The input projection is the dominant term.

where TCSD is CSD's mean aggregation and element-wise operations on input tensors, TProj is input projection's linear mapping of D-dimensional alpha factors to d-dimensional model features. TTAG is TAG's linear layer and element-wise gating on N × L × d features, and TMLP is the output MLP's time complexity which contains two linear layers operated on 2d-dimensional pooled features. Table 1 summaries the time complexity of TransAlpha and baseline models.

TransAlpha's space complexity is dominated by intermediate activation tensors (N × L × d), resulting in an asymptotic complexity of O(NLd), identical to all baseline models. Trainable parameters O(Dd+d2 ) are negligible in comparison to activation memory usage. In terms of time efficiency, TransAlpha ranks the best, highlighting its lightweight characteristic. Autoformer has the same nominal O(NLd) time complexity, but its practical efficiency is lower due to Fast Fourier Transform (FFT) during decomposition. Crossformer introduces an extra L/G overhead factor. Informer and Pyraformer add a logL factor. iTransformer introduces an L2 /h factor, and TCN adds a kernel size factor K.

4 Experiments

We formulate three core research questions (RQs).

• RQ1: Are Transformer models effective in general? How is TransAlpha's predictive power across different market segments?
• RQ2: Can the models transfer their predictive power into profit in back-test?
• RQ3: Are the proposed architectures effective?
• RQ4: Do TransAlpha's prediction patterns exhibit distinct characteristics?

4.1 Experimental Set up

Data Set Details We evaluate on an intraday dataset directly coming from production environment, covering the entire A-share market (Shanghai and Shenzhen Stock Exchanges), with four key market segments (CSI 300, CSI 500, CSI 1000, and CSI 2000) for subset modeling to assess generalization ability across market capitalization. The dataset contains 400 proprietary alpha factors in 15-minute frequency cross-sectionals stored in HDF5 format and are normalized without extreme values. The observation period is non-overlapping, consisted of training (2021-01-04 to 202309-28), validation (2023-10-09 to 2023-12-29), and testing (2024-01-02 to 2024-12-31) sets to avoid look-ahead bias and ensure out-of-sample evaluation. There are Nd = 242 trade days, each from 09:45 am to 15:00 pm daily (Nt = 16 cross sections per day, total of T = Nd ×Nt = 3856). Transaction costs are applied only on position liquidation at 0.1% per trade, inclusive of fees.

Baseline Models We compare TransAlpha against 3 categories of baselines. Linear models include DLinear [Zeng et al., 2023], RLinear [Li et al., 2023], and TiDE [Das et al., 2023]. Non-Transformer DL models include GRU [Chung et al., 2014], DeepAR [Salinas et al., 2020], and TCN [Bai et al., 2018], TimesNet [Wu et al., 2022], Modern TCN [donghao and wang xue, 2024]. Transformer models include Informer [Zhou et al., 2021], Autoformer [Wu et al., 2021], PatchTST [Nie et al., 2023], Pyraformer [Liu et al., 2022], Crossformer [Zhang and Yan, 2023], Fedformer [Zhou et al., 2022] and iTransformer [Liu et al., 2024].

Implementation Details Our model is implemented in PyTorch 2.7.0 with CUDA11.8. All experiments are conducted on a server equipped with an AMD EPYC 7T83 64-Core Processor (256 total CPUs across 2 NUMA nodes, 64 physical cores per socket, 128 siblings per node) and 1.5 TiB system memory. For training, we utilized a cluster of four NVIDIA GeForce RTX 4090 GPUs (24564 MiB VRAM per GPU). We use grid search to perform hyperparameter optimization and obtain optimal values based on validation set. We trained the model separately for all market subsets and adopted a naive long-short 50 strategy to back-test.

Evaluation Metrics Stock price forecasting's success is measured by trading profitability rather than point prediction accuracy. Traditional forecasting metrics such as MSE fail to capture the relative performance rankings that drive quant trading strategies. We use the Information Coefficient (IC), Rank IC, Information Coefficient Information Ratio (ICIR) and Rank ICIR to quantify predictive power. We evaluate a model's profitability power using the Annualized Return (AR), Information Ratio (IR), Sharp Ratio (SR), Maximum Drawdown (MDD) and Turn over Rate (TR).

4.2 RQ1Model Predictive Power

We experiment on both the full market data and all CSI indices subsets to evaluate TransAlpha's predictive power under different market segments, via these standard metrics. See Table 3's first four columns. Overall, TransAlpha's predictive power is consistent across all subsets and constantly among the top models. Notably, subsets with more stocks do not always mean performance increase.

4.3 RQ2 Backtest Result

Overall, TransAlpha offers strong AR, IR and SR, with notable lead in SR. Most importantly, strong predictive power does not necessarily transfer to profitability, a sign of overfitting, with iTransformer as a typical example. MDD < 0.05 suggests worst peak-to-trough loss is less than 5% at any point. TR < 0.20 suggests it only needs to rebalance the portfolio by less than 20% regularly, consuming low trading costs and minimal operational hassle. TransAlpha maintains all lower than the empirical thresholds, offering stable predictions.

5 Conclusions

This paper tackles challenges of forecasting noisy crosssection data in a real stock market context. We introduce TransAlpha, a lightweight Transformer variant. Extensive experiments suggest it outperforms on core trading metrics and its core components are effective to handle noise. Beyond finance, CSD has the potential to denoise sensor data in manufacturing; TAG can prioritize time steps in patient monitoring; Smart Pooling can aggregate sparse signals in traffic flow. The experiments are exclusively conducted on A-Share market, using proprietary alpha factors from a single quantitative firm. Also, TransAlpha relies on a fixed lookback window and static loss function weights. In future work, we will extend to American stock markets and include MoE components with tailored experts to scale up parameter size, incorporating methods to adapt dynamic market conditions. Finally, we will test the architecture in different forecasting tasks, such as the benchmark data sets in long-range time series forecasting.

Figure 2: Pairwise correlation heatmaps of model predictions: first row and column are TransAlpha, with the rest models sorted in alphabetical order.

4.4 RQ3 Ablation Study

We include an ablation study to test the effectiveness of core components in TransAlpha by removing CSD and TAG portion respectively. Table 2 shows an evident drop in backtest metrics especially IR and SR after removing either core components, which indicates the components are effective in boosting profitability. Specifically, CSD can reduce the over-fitting issue, which can be generalized to other highdimensional forecasting tasks. TAG also verifies its efficiency gain over standard multi-head attention O(NLD) vs O(NL2D).

4.5 RQ4 Pattern distinction

The model's generated predictions offer distinct patterns from heatmaps and dendrograms of hierarchical clustering. See Figure 2, 3. On large-cap subsets (CSI 300 and CSI 500), the correlations are much lower, indicating the model learned distinct pattern compared to other baselines.