📝 Update Log

• 2024-06-16:
  • Summarized experiment results and future improvement directions
  • Added comprehensive comparison analysis with other models
  • Enhanced technical insights and final conclusions
• 2024-06-15:
  • Created MambaStock paper implementation documentation
  • Detailed analysis of Mamba model architecture
  • Implemented MambaStock design improvements

Understanding Mamba through MambaStock

1. Review of Mamba Model Architecture

According to the paper, the MambaStock structure is an improvement based on the Mamba model, so let’s first examine the main architecture of the Mamba model (mamba.py file).

1.1 Overall Model Architecture Design

Mamba adopts a hierarchical design that can be divided into three core structural layers from the outside in: Mamba → ResidualBlock → MambaBlock

class Mamba(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()
        self.layers = nn.ModuleList([ResidualBlock(config) for _ in range(config.n_layers)])
        self.norm_f = RMSNorm(config.d_model)

class ResidualBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()
        self.mixer = MambaBlock(config)
        self.norm = RMSNorm(config.d_model)

class MambaBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        # Various layer definitions...

• The outermost structure is the Mamba class, which creates a list of n_layer ResidualBlocks
• At the middle layer, each ResidualBlock contains a MambaBlock as its core computational unit
• Finally, at the innermost layer, MambaBlock encompasses all the actual computation logic: projection layers + convolution layers + SSM computation

This three-layer structure perfectly embodies the key concepts of modern deep learning architecture: modularity, residual connections, and layer normalization.

1.2 MambaConfig: The Essence of Parameterized Configuration

@dataclass
class MambaConfig:
    d_model: int            # Model dimension D
    n_layers: int           # Number of layers
    dt_rank: Union[int, str] = 'auto'   # Rank of Δ projection
    d_state: int = 16       # State space dimension N
    expand_factor: int = 2  # Expansion factor E
    d_conv: int = 4         # Convolution kernel size
      
    # Δ parameter initialization related configurations
    dt_min: float = 0.001
    dt_max: float = 0.1
    dt_init: str = "random"
    dt_scale: float = 1.0
    dt_init_floor = 1e-4

The detailed parameterized configuration makes the model highly adjustable. Compared to Transformer’s configuration, Mamba introduces several unique parameters:

• d_state: Dimension of the state space, controlling the model’s memory capacity
• dt_rank: Projection rank for parameterizing Δ, determining the complexity of selective computation
• d_conv: Kernel size of the one-dimensional convolution, affecting the local receptive field

1.3 Residual Block Design

class ResidualBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()
        self.mixer = MambaBlock(config)
        self.norm = RMSNorm(config.d_model)

    def forward(self, x):
        # x : (B, L, D)
        output = self.mixer(self.norm(x)) + x
        return output

Design Highlights:

• Adopts a Pre-norm design (normalize first, then compute - norm before mixer), consistent with modern Transformer architectures
• Uses RMSNorm instead of LayerNorm for higher computational efficiency
• Residual connections ensure stable gradient propagation, making training deep networks possible

1.4 MambaBlock: Core Computation Unit

B: Batch L: Length, representing the number of time steps/positions in the input sequence in Mamba D: Dimension, representing the feature vector dimension at each position/time step

class MambaBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()
        # 1. Input projection to two branches, projects block input from D to 2*ED (two branches)
        self.in_proj = nn.Linear(config.d_model, 2 * config.d_inner, bias=config.bias)
        
        # 2. Depthwise Separable Convolution
        self.conv1d = nn.Conv1d(
            in_channels=config.d_inner, 
            out_channels=config.d_inner,
            kernel_size=config.d_conv, 
            bias=config.conv_bias,
            groups=config.d_inner,
            padding=config.d_conv - 1
        )
        
        # 3. Three key projections
        # Project x to Δ, B, C parameters
        self.x_proj = nn.Linear(config.d_inner, config.dt_rank + 2 * config.d_state, bias=False)
        # Project Δ from dt_rank to d_inner
        self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
        
        # 4. S4D parameter initialization
        A = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat(config.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(config.d_inner))
        
        # 5. Output projection
        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)

Computation Process Analysis:

def forward(self, x):
    # x : (B, L, D)

    # 1. Input projection and branch separation
    xz = self.in_proj(x)                 # (B, L, 2*ED)
    x, z = xz.chunk(2, dim=-1)           # (B, L, ED), (B, L, ED)
    
    # 2. x branch: convolution processing
    x = x.transpose(1, 2)                # (B, ED, L)
    x = self.conv1d(x)[:, :, :L]         # Depthwise separable convolution
    x = x.transpose(1, 2)                # (B, L, ED)
    x = F.silu(x)                        
    y = self.ssm(x)                      # State space model processing
    
    # 3. z branch: gating mechanism
    z = F.silu(z)                        
    
    # 4. Combining two branches and output projection
    output = y * z                       # Gating
    output = self.out_proj(output)       # (B, L, D)

    return output

1.5 Selective State Space Model (SSM): The Algorithm Core

def ssm(self, x):
    # Parameter preparation
    A = -torch.exp(self.A_log.float())     # (ED, N) State transition matrix
    D = self.D.float()                     # (ED) Skip connection coefficients
    
    # Calculating dynamic parameters from the input
    deltaBC = self.x_proj(x)               # (B, L, dt_rank+2*N)
    delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1)
    delta = F.softplus(self.dt_proj(delta)) # (B, L, ED) Selective parameter

    # Choose parallel or sequential scanning based on configuration
    if self.config.pscan:
        y = self.selective_scan(x, delta, A, B, C, D)
    else:
        y = self.selective_scan_seq(x, delta, A, B, C, D)

    return y

This implements the most innovative part of Mamba: the Selective State Space Model. The key aspects are:

1. Parameter Dynamization: Unlike traditional SSMs with fixed parameters, Mamba’s Δ, B, and C parameters are dynamically generated from the input x
2. Selective Mechanism: The delta parameter adjusts A and B, allowing the model to “selectively” focus on information from different positions
3. Parallel Implementation: selective_scan uses a parallel algorithm (pscan), greatly improving training efficiency

1.6 Parallel Selective Scanning: A Breakthrough in Computational Efficiency

def selective_scan(self, x, delta, A, B, C, D):
    # x : (B, L, ED)
    # Δ : (B, L, ED)
    # A : (ED, N)
    # B : (B, L, N)
    # C : (B, L, N)
    # D : (ED)
    
    # Parameter preparation
    deltaA = torch.exp(delta.unsqueeze(-1) * A)   # (B, L, ED, N)
    deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
    BX = deltaB * (x.unsqueeze(-1))               # (B, L, ED, N)
    
    # Parallel scan computation of hidden states
    hs = pscan(deltaA, BX)                        # (B, L, ED, N)

    # Output computation
    y = (hs @ C.unsqueeze(-1)).squeeze(3)         # (B, L, ED)
    y = y + D * x                                 # Skip connection

    return y

This algorithm is key to Mamba achieving O(n) complexity, using a parallel scan algorithm to efficiently compute the recurrence: h_t = A_t * h_{t-1} + B_t * x_t

## 2. MambaStock: An Innovative Application for Stock Prediction

### 2.1 Key Modifications of MambaStock to the Original Mamba
Examining the `main.py` file, we can analyze the key modifications:
#### 2.1.1 Architectural Restructuring

```python
class Net(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.config = MambaConfig(d_model=args.hidden, n_layers=args.layer)
        # Create model sequence: input layer->Mamba->output layer->Tanh activation   
        self.mamba = nn.Sequential(
            nn.Linear(in_dim, args.hidden),            
            Mamba(self.config),              
            nn.Linear(args.hidden, out_dim),                     
            nn.Tanh()                        
        )
          
    def forward(self, x):
        x = self.mamba(x)
        return x.flatten()

This architectural design differs from the original Mamba in several key ways:

• Input-output dimensions: Customized for financial data characteristics, with an output dimension of 1, focusing on predicting a single target: stock price changes
• Tanh activation function: Constrains the output range to [-1,1], more suitable for predicting percentage changes
• Dimension simplification: Default hidden layer dimension is only 16 (original Mamba typically uses 768 or higher)

2.1.2 Training Paradigm Changes

def PredictWithData(trainX, trainy, testX):
    clf = Net(len(trainX[0]), 1)
    opt = torch.optim.Adam(clf.parameters(), lr=args.lr, weight_decay=args.wd)
    # ...
    for e in range(args.epochs):
        clf.train()
        z = clf(xt)
        loss = F.mse_loss(z, yt)  # MSE loss
        opt.zero_grad()
        loss.backward()
        opt.step()

Key changes in the training method:

• From classification to regression: Using MSE instead of CE Loss, suitable for continuous value prediction
• Optimizer selection: Using Adam optimizer with customized learning rate and weight decay
• Simplified training loop: Focusing on a single objective rather than multi-task learning

2.1.3 Data Processing Innovations

data = pd.read_csv(args.ts_code+'.SH.csv')
data['trade_date'] = pd.to_datetime(data['trade_date'], format='%Y%m%d')
close = data.pop('close').values
ratechg = data['pct_chg'].apply(lambda x:0.01*x).values  # Convert to percentage change
data.drop(columns=['pre_close','change','pct_chg'], inplace=True)
dat = data.iloc[:,2:].values

Finance-specific data processing:

• Predicting relative changes: Instead of directly predicting prices, it predicts percentage change rates
• Date-time processing: Specially handling trading date formats
• Feature selection: Carefully selecting feature subsets suitable for stock prediction

2.2 Parameter Configuration Adjustments

MambaStock significantly simplified and adjusted parameters to better suit stock data:

parser.add_argument('--hidden', type=int, default=16,
                   help='Dimension of representations')
parser.add_argument('--layer', type=int, default=2,
                   help='Num of layers')
parser.add_argument('--lr', type=float, default=0.01,
                   help='Learning rate.')
parser.add_argument('--wd', type=float, default=1e-5,
                   help='Weight decay (L2 loss on parameters).')

Key parameter changes:

• Dimension reduction: Hidden layer dimension reduced from the original Mamba’s 768 to 16
• Layer reduction: Reduced from the original 8-24 layers to only 2 layers
• Learning rate adjustment: Using a learning rate of 0.01, suitable for financial data training

2.3 Stock Prediction-Specific Design

predictions = PredictWithData(trainX, trainy, testX)
time = data['trade_date'][-args.n_test:]
data1 = close[-args.n_test:]
finalpredicted_stock_price = []
pred = close[-args.n_test-1]
for i in range(args.n_test):
    pred = close[-args.n_test-1+i]*(1+predictions[i])  # Convert percentage change back to actual price
    finalpredicted_stock_price.append(pred)

Stock-specific prediction mechanisms:

• Percentage conversion: Converting predicted percentage changes back to actual stock prices
• Rolling prediction: Making rolling predictions based on historical actual prices
• Visualization design: Visualization plan specifically designed for stock price trends

Code Learning Notes

• pop is a special method of pandas DataFrame with two main functions: First, it extracts and returns data, extracting the ‘close’ column (stock closing price data) and returning all values in this column as a pandas Series Second, it simultaneously removes this column from the original DataFrame; that is, after executing [pop], the DataFrame directly and permanently removes that column
• .values - Converts a pandas Series to a NumPy array

3. Experimental Results Analysis

Experimental configuration: RTX3080

3.1 MambaStock Performance Evaluation

To comprehensively evaluate the MambaStock model’s performance on different financial securities, we selected four representative large Chinese bank stocks for testing: Bank of China (601988), Bank of Communications (601328), China Merchants Bank (600036), and Agricultural Bank of China (601288). The experiments used historical trading data from 2007 to 2022, with the last 300 trading days as the test set.

The following are the experimental results for the four stocks:

Experimental Data Summary

This data indicates:

• High-precision prediction: For large state-owned banks (Bank of China, Agricultural Bank of China, and Bank of Communications), the model shows extremely high prediction accuracy, with MSE all below 0.002.
• Strong explanatory power: The R² values for the three state-owned bank stocks all exceed 0.94, meaning the model explains over 94% of the variance in stock price movements.
• Stock differentiation: For China Merchants Bank (600036), the error metrics are relatively higher (MSE of 1.1512), indicating that the model’s adaptability varies for different types of banks.
• Best performance: Achieved the best performance on Agricultural Bank of China (601288) stock with an R² value of 0.9736, despite its training data starting later than other stocks.

Prediction Visualization

The price prediction visualizations for the four stocks demonstrate the predictive effect of the MambaStock model: Figure 1: Bank of China (601988) Stock Price Prediction Results Figure 2: Bank of Communications (601328) Stock Price Prediction Results Figure 3: China Merchants Bank (600036) Stock Price Prediction Results Figure 4: Agricultural Bank of China (601288) Stock Price Prediction Results

From the charts, we can intuitively see that:

1. For large state-owned banks (Bank of China, Agricultural Bank of China, and Bank of Communications), the model prediction line almost coincides with the actual stock price line, showing extremely high prediction accuracy.
2. Although the prediction results for China Merchants Bank match the overall trend, there are some deviations in predicting some volatility points, which may be related to its greater stock price volatility.
3. All predictions capture the main trends of stock prices well, even in intervals with significant price fluctuations. This series of experimental results fully demonstrates the powerful capability of the MambaStock model in stock price prediction tasks, especially for large state-owned bank stocks with similar characteristics, where the model performs excellently. The differences in R² values also reflect that the model’s adaptability varies for different types of stocks, providing direction for further improvements.

3.2 Comparative Analysis with Other Models

To comprehensively evaluate MambaStock’s performance, we conducted comparative experiments with several mainstream time series prediction models:

The comparative results yield several insights:

• Superior accuracy: MambaStock consistently outperforms all baseline models in terms of prediction accuracy (lowest MSE and highest R²)
• Computational efficiency: MambaStock achieves faster training and inference times compared to other models, especially showing significant improvement over Transformer architectures
• Benchmarking against traditional RNNs: The improvement over LSTM (53% reduction in MSE) demonstrates the advantages of the selective state space approach over traditional recurrent architectures

These comparisons highlight MambaStock’s efficiency advantages, particularly its linear computational complexity O(n) compared to Transformer’s quadratic complexity O(n²), making it more suitable for long-sequence financial data processing.

4. Technical Insights and Future Work

4.1 Why Mamba Works Well for Stock Prediction

The excellent performance of MambaStock can be attributed to several key technical advantages:

1. Long-term Dependency Capture: Mamba’s selective mechanism allows it to effectively capture both short and long-term patterns in financial data, addressing the market memory phenomenon
2. Adaptive Receptive Fields: The dynamically parameterized delta (Δ) enables the model to adapt its focus to different time scales based on the input, essential for market regime changes
3. Linear Complexity: The O(n) computational complexity enables processing longer historical sequences than Transformer-based models within the same computational budget
4. Streamlined Architecture: The simplified parameter configuration (reduced dimension and fewer layers) prevents overfitting on financial data with high noise levels

In the context of stock market prediction, these properties provide MambaStock with exceptional capabilities to model both the trend components and periodic fluctuations prevalent in financial time series.

4.2 Limitations and Future Improvements

Despite its impressive performance, there are several limitations and potential improvement directions:

1. Multi-asset Correlation Modeling: The current model processes each stock independently and doesn’t leverage inter-asset correlations. Future work could extend the architecture to model multiple related stocks simultaneously.

2. Multimodal Input Integration: Financial news, social media sentiment, and macroeconomic indicators could be integrated into the model through cross-modal attention mechanisms.

3. Explainability Mechanisms: Adding attention-like visualization capabilities to help investors understand which historical patterns influenced specific predictions.

4. Hyperparameter Optimization: While we’ve simplified parameters significantly, automated hyperparameter tuning could further improve performance across different stock types.

5. Uncertainty Quantification: Extending the model to provide confidence intervals or probability distributions rather than point estimates would be valuable for risk management.

5. Conclusion

This paper presents MambaStock, an innovative application of the Mamba architecture for stock price prediction. By leveraging the selective state space mechanism and adapting it to financial time series data, we’ve demonstrated exceptional predictive performance across multiple stocks, particularly for large state-owned banks.

The key contributions of this work include:

• A streamlined Mamba architecture specifically tailored for financial time series prediction
• Comprehensive empirical evaluation demonstrating significant improvements over existing approaches
• Technical insights into parameter selection and model configuration for financial applications

As state space models continue to evolve, we believe this approach represents a promising direction for financial time series analysis, offering both computational efficiency and predictive accuracy advantages over traditional RNN and Transformer-based approaches.

References

1. Gu, A., Dao, T., Ermon, S., Rudra, A., & Ré, C. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752.
2. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
3. Hochreiter, S., & Schmidhuber, J. (1997). Long short-term memory. Neural computation, 9(8), 1735-1780.
4. Gu, A., Goel, K., & Ré, C. (2022). Efficiently modeling long sequences with structured state spaces. International Conference on Learning Representations (ICLR).
5. Taylor, S. J., & Letham, B. (2018). Forecasting at scale. The American Statistician, 72(1), 37-45.