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 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_ratiowas 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_returnat 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.
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
- Downloads last month
- 27