| --- |
| license: apache-2.0 |
| library_name: pytorch |
| pipeline_tag: feature-extraction |
| tags: |
| - time-series |
| - finance |
| - stocks |
| - mamba |
| - self-supervised |
| - masked-reconstruction |
| - embeddings |
| --- |
| |
| # stock-embed-mamba3 |
|
|
| A 3.3M-parameter Mamba-3 encoder with a linear decoder, pretrained by masked-patch |
| reconstruction on 1-minute US equity bars: windows of up to 10,000 bars, the most recent |
| part of a 30-trading-day span. It produces a 256-dimensional embedding per window and can |
| fill in masked 5-bar patches of its nine derived input channels. Running it needs a CUDA |
| GPU and a source build of `mamba_ssm`; the companion Dockerfile does that build. The |
| data-preparation helpers run anywhere. |
|
|
| **Read [the limitations](#what-the-model-does-not-have) first.** The model is a research artifact released with a documented |
| null result. It has no demonstrated forecasting or trading value. It is not generative: it |
| cannot sample or roll forward. It can only fill in masked patches given context on both |
| sides. |
|
|
| Code, build recipe and tests: https://github.com/HammySession/stock-embed-mamba3 |
| (note the different spelling of the two account names: `HamSession` here, `HammySession` |
| on GitHub). |
|
|
| ## Model details |
|
|
| | | | |
| |---|---| |
| | Architecture | `MambaEncoder` (input projection, 8 pre-norm Mamba-3 blocks, LayerNorm) plus a `Linear(256 → 45)` decoder (45 = 5 bars × 9 channels) | |
| | Parameters | 3,299,245 in 87 tensors | |
| | d_model / d_state / expand | 256 / 16 / 2 | |
| | Mamba-3 block settings | headdim 64, 1 group, rope fraction 0.5, chunk size 64, no MIMO (Mamba-3 constructor options, see the `Mamba3` signature); all recorded in `config.json` and passed explicitly by the loader | |
| | Input | `(batch, seq_len, 9)` float32, `seq_len` divisible by 5, at most 10,000 bars (see Sequence length) | |
| | Patching | 5 consecutive bars form one token | |
| | Output | `(batch, seq_len/5, 256)` tokens; `embed()` mean-pools them to one vector per window | |
| | Objective | MSE on masked patches, random patch masking at ratio 0.15, per-channel loss weights `[1, 1, 1, 0.001, 0.1, 1, 1, 1, 1]` | |
| | Precision | bf16 autocast in training; weights stored fp32 | |
| | Files | `model.safetensors`, `config.json`, `modeling_stock_mamba.py`, `adjust_ohlcv.py`, `features.py`, `preprocessing_contract.md`, `export_receipt.json`, `LICENSE` | |
| | Dependencies | `torch`, `safetensors`, `mamba_ssm` with the `Mamba3` class, built from state-spaces/mamba commit `e9594ce` (reports 2.3.2.post1), `causal_conv1d` from commit `4f6ae4e` (1.6.2.post1), the companion Dockerfile being the supported build path; `pandas` and `numpy` for data preparation | |
| | Hardware | A CUDA GPU is required. The Mamba-3 kernels have no CPU path, and the module fails at import time on a CPU-only machine | |
|
|
| The nine input channels, in order: `log_return, hl_range, oc_body, log_vol_ratio, |
| session_flag, tod_sin, tod_cos, dow_sin, dow_cos`. Exact formulas, clipping bounds and |
| timezone conventions are in `preprocessing_contract.md`. Two helpers take you from your own |
| **unadjusted** minute bars to model input without guessing. `adjust_ohlcv.py` applies split |
| and dividend adjustment under the training convention and the pipeline's raw-bar quality |
| checks. `features.py` builds the nine channels. Both are pandas and numpy only. The authors |
| checked both against the training code with zero numerical difference on synthetic bars. |
| No feature data ships with this model. |
|
|
| ## Sequence length |
|
|
| Training windows were 30 trading days per security, bars from 04:00 to 20:00 Eastern with |
| pre-market and post-market included, capped at 10,000 bars by keeping the **most recent** |
| 10,000. A liquid name prints far more than 10,000 bars in 30 such days, so for most of the |
| universe the model saw the last 10,000 bars of the span, roughly the last ten sessions. |
| Bar counts per window were not recorded. Pass at most 10,000 bars, trimmed from the front. Shorter inputs run but were not evaluated. Longer |
| inputs are out of distribution. |
|
|
| The `log_vol_ratio` channel divides each bar's volume by a 20-day rolling mean that |
| excludes the current day. In training that mean was computed on each security's full |
| history, so every window started warm. To match, feed `make_features` at least 20 trading |
| days of bars before the window you want, then slice the lead-in off the result. |
|
|
| ## What the model has (authors' measurements) |
|
|
| Measured by the authors on the frozen encoder, US equities, with downstream probes fit on |
| a 2025 development window and scored on a 2026 out-of-sample window. The 2025 window |
| overlaps the validation period used to pick the checkpoint. Only the 2026 numbers are |
| fully out of sample. |
|
|
| - **A stable per-stock identity embedding.** A per-stock fixed effect explains 99.5% of the |
| mean-pooled window embedding, and its lag-one daily autocorrelation is 0.83, both |
| measured on a 200-stock sample. That makes it suitable as a per-stock fingerprint for |
| nearest-neighbour retrieval. Whether its clusters carry economic meaning was not shown: |
| in the worked examples a per-ticker summary-statistics baseline recovers sector |
| membership better at every k, and k-means clusters match sectors at an adjusted Rand |
| index of 0.07. It is not suitable for detecting what changed this week. |
| - **Reconstruction error on the three price channels** (`log_return`, `hl_range`, |
| `oc_body`), with no baseline comparison; see the next section. Validation MSE at the best |
| epoch (epoch 16, zero-based, the last of the 17 run): |
|
|
| | Channel | Validation MSE | |
| |---|---| |
| | log_return | 7.0e-05 | |
| | hl_range | 3.9e-05 | |
| | oc_body | 3.1e-05 | |
| | log_vol_ratio | 2.15 | |
| | session_flag | 7.8e-03 | |
| | tod_sin / tod_cos | 5.3e-02 / 1.9e-02 | |
| | dow_sin / dow_cos | 5.4e-02 / 5.4e-02 | |
| | total (weighted mean over channels) | 0.02591 | |
|
|
| ## What the model does not have |
|
|
| - **No cross-sectional return signal**, meaning no ability to rank stocks against each |
| other on a given day. A small MLP-classifier probe on the frozen embedding, tuned over |
| regularisation and hidden sizes and scored on forward returns at one, three and five |
| trading days, gives a Spearman rank correlation between prediction and realised forward |
| return of about 0.0045, with top-versus-bottom-quintile accuracy of 50.5%, which is |
| chance. Used as an extra input to a gradient-boosted classifier of large next-day |
| moves, the embedding never beat two controls: the same embedding shuffled across stocks |
| within each day, and one constant vector per stock. The authors stopped work |
| on combining this embedding with other models. |
| - **No forecasting.** Training masked 15% of patches at random with context on both sides. |
| Masking the trailing patches to "forecast" is out of distribution and was never |
| evaluated. |
| - **No usable volume reconstruction.** `log_vol_ratio` was trained at loss weight 0.001 and |
| its validation MSE of 2.15 is about four orders of magnitude above the price channels. |
| Treat its reconstructions as untrained. As an input it still shapes the embedding: in |
| the worked examples, feeding it cold with no lead-in cut nearest-neighbour sector |
| accuracy from 23% to about 15%. |
| - **No baseline comparison.** The reconstruction objective was never benchmarked against |
| zero, last-value or linear-interpolation fills. Low MSE on `log_return` at 1-minute |
| resolution is probably close to what predicting zero would score on a channel whose mean is near |
| zero. |
|
|
| ## Training data and provenance |
|
|
| - Source: 1-minute open, high, low, close and volume (OHLCV) bars for roughly 8,300 |
| US-listed securities, stocks and ETFs, licensed from EODHD (EOD Historical Data) and |
| adjusted for splits and dividends. |
| - Date ranges: training bars dated up to and including 2023-12-31. Validation bars |
| 2024-01-01 to 2025-08-01, used for checkpoint selection. The checkpoint has seen no data |
| after 2025-08-01. |
| - Sampling: windows of 30 trading days at a 5-day stride, at most 10,000 bars per window. |
| Each epoch draws a random 15% of training windows and 20% of validation windows. |
| - Optimizer: AdamW, learning rate 9.7e-4, weight decay 8.4e-4, betas (0.9, 0.95), cosine |
| schedule over 30 planned epochs with 2 warmup epochs, batch 32 windows per GPU with 3 |
| gradient-accumulation steps on 8 GPUs (effective batch 768 windows), gradient clip 1.0, |
| seed 42. |
| - Compute: 8 × RTX 5090 under distributed data parallel, about 13.5 hours. |
| - Stopping: early stopping (patience 8, minimum improvement 0.001) fired after epoch 16 of |
| the planned 30. Validation loss was still improving every epoch, by less than 0.001, |
| and the learning rate was at 4.4e-4, about 45% of peak. The checkpoint is the last |
| epoch run, not the end of a completed schedule, which matters if you use it as a warm |
| start: expect the loss to keep falling under continued training. |
|
|
| **Nothing from the training data ships here.** No bars, feature panels, per-ticker tables |
| or identifiers. The weights store no bars and cannot reproduce the licensed data unless |
| it is supplied as input. |
|
|
| ## Usage |
|
|
| Requirements: a CUDA GPU, `torch`, `safetensors`, `huggingface_hub`, `pandas`, `numpy`, |
| and `mamba_ssm` with the `Mamba3` class. The PyPI release of `mamba_ssm` is source only |
| and needs the CUDA toolkit to build. The Dockerfile in the companion repo reproduces the |
| build used to export and validate the weights (the training run itself used an unpinned |
| upstream build), for compute capability 12.0 (Blackwell) by default, with a build |
| argument for other architectures. Torch is an unpinned CUDA 12.8 nightly because nightlies |
| are pruned from the index; `export_receipt.json` records the exact versions used. |
|
|
| ```python |
| import sys |
| import pandas as pd, torch |
| from huggingface_hub import snapshot_download |
| |
| repo = snapshot_download("HamSession/stock-embed-mamba3") |
| sys.path.insert(0, repo) # the helpers live inside the snapshot |
| from adjust_ohlcv import adjust_ohlcv |
| from features import make_features |
| from modeling_stock_mamba import MaskedReconstructionModel |
| |
| bars = ... # one security, UNADJUSTED 1-minute bars, columns: timestamp (bar start; |
| # tz-aware, or naive taken as Eastern), open, high, low, close, volume |
| splits = pd.DataFrame({"date": [...], "ratio": [...]}) # ex-dates; empty frame if none |
| dividends = pd.DataFrame({"date": [...], "amount": [...]}) # ex-dates; empty frame if none |
| |
| feats = make_features(adjust_ohlcv(bars, splits, dividends)) # (seq_len, 9) float32 |
| feats = feats[-10_000:] # keep the most recent bars |
| feats = feats[len(feats) % 5 :] # seq_len divisible by 5, trimmed from the front |
| x = torch.from_numpy(feats).unsqueeze(0).cuda() # (1, seq_len, 9) |
| |
| model = MaskedReconstructionModel.from_pretrained(repo, device="cuda") |
| emb = model.embed(x) # (1, 256) |
| tokens = model.encoder(x) # (1, seq_len // 5, 256) |
| |
| n_patches = x.shape[1] // 5 |
| patch_mask = torch.zeros(1, n_patches, dtype=torch.bool, device="cuda") |
| patch_mask[:, n_patches // 2 : n_patches // 2 + 10] = True # mask ten 5-bar patches |
| filled = model.inpaint(x, patch_mask) # (1, seq_len, 9); visible bars unchanged |
| ``` |
|
|
| `embed` and `inpaint` take equal-length windows only. There is no padding mask, so batch |
| securities of the same `seq_len` or call the model once per security. Inference runs in |
| fp32 as shown, as did the export round-trip. On a CPU-only machine `adjust_ohlcv` and |
| `make_features` work (numpy and pandas only) but the model does not, and there is no CPU |
| fallback. A 10,000-bar window in fp32 peaks at about 320 MB of GPU memory and takes about |
| 5 ms per `embed` call on an RTX 5090 after warm-up. |
|
|
| ## Where to get bars and worked examples |
|
|
| No data ships with the model. Free 1-minute bars for the last 30 days come from Yahoo |
| Finance through `yfinance` (pass `auto_adjust=False`, apply its dividends only because |
| its prices are already split-adjusted; no extended-hours volume). Longer |
| histories come from Alpaca (free with an account, since 2016), Polygon.io (now Massive), |
| EODHD (the training source) or Databento. The GitHub repo's README compares them. |
|
|
| The GitHub repo's `examples/` directory goes from Yahoo data to a clustering plot and a |
| leave-one-out sector classifier on 66 large US stocks. On the authors' run the embedding |
| put a stock's nearest neighbour in the same sector 23% of the time against 8% chance, |
| and per-ticker means and standard deviations of the same nine channels did better at |
| every k. Treat the embedding as a stock fingerprint, not a sector detector. |
|
|
| ## Intended use |
|
|
| - Research on self-supervised representations of intraday price data. |
| - A documented negative baseline for whether masked reconstruction on minute bars yields |
| cross-sectional signal, meaning signal that ranks stocks against each other on a given |
| day. |
| - Stock-identity embeddings for retrieval and clustering. |
| - A warm-start backbone for further pretraining. |
|
|
| ## Out of scope |
|
|
| Live or automated trading, portfolio construction, any routing of outputs to order |
| execution, price or return forecasting, and any use as the sole basis for a financial |
| decision. |
|
|
| ## License |
|
|
| Apache-2.0 for these weights and for the companion code. The data vendor's permission |
| covers the release of the weights on the condition that no vendor data is distributed, |
| and none is. |
|
|
| ## Disclaimer |
|
|
| **Not investment advice.** Released for research and educational purposes only. Outputs |
| are embeddings and reconstructions of derived features, not return forecasts, price |
| targets or trading signals. **No demonstrated edge.** In the authors' own evaluation this |
| model carried no cross-sectional return signal and produced no strategy that was |
| profitable after transaction costs in the authors' backtests. **Training data:** US equity 1-minute bars licensed from EODHD, released with the |
| vendor's permission on the condition that no vendor data is distributed. Users must obtain |
| their own data license to run inference. Not affiliated with or endorsed by any data vendor |
| or exchange. |
|
|
| ## Integrity |
|
|
| `export_receipt.json` records the SHA-256 of the source checkpoint and of |
| `model.safetensors`, the library versions used for the export, and a round-trip check. The |
| staged weights reproduce the original checkpoint's encoder tokens and decoder output |
| exactly (max abs diff 0.0). |
|
|
| ## Issues and contact |
|
|
| Open an issue on the companion repository: |
| https://github.com/HammySession/stock-embed-mamba3/issues |
|
|
