TimeTron-v2-33M: 33.5M latent-distilled TS foundation model (GIFT 0.8794 nMASE / 0.6122 nCRPS)

.gitattributes

@@ -1,35 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,123 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: time-series-forecasting
+tags:
+  - time-series
+  - forecasting
+  - time-series-forecasting
+  - knowledge-distillation
+  - foundation-model
+  - probabilistic-forecasting
+  - gift-eval
+---
+
+# TimeTron-v2-33M
+
+A **33.5M-parameter** time-series foundation model, distilled from **Chronos-2** by
+**latent-space knowledge distillation** — the student matches the teacher's internal
+representations rather than its forecasts.
+
+**GIFT-Eval: 0.8794 normalized MASE / 0.6122 normalized CRPS** across all 97 configurations
+(official Salesforce harness). On probabilistic accuracy it beats models up to **21× larger**.
+
+| model | params | nMASE | nCRPS |
+|---|---|---|---|
+| **TimeTron-v2-33M** | **33.5M** | **0.8794** | **0.6122** |
+| Chronos-Small | 46M | 0.8917 | 0.6630 |
+| Moirai-Base | 91M | 0.9010 | 0.6095 |
+| Chronos-Base | 200M | 0.8758 | 0.6521 |
+| Chronos-Large | 710M | 0.8696 | 0.6473 |
+| *Chronos-2 (teacher)* | *~120M* | *0.6978* | *0.4854* |
+
+Trained on **8.25B points** (7.5B latent distillation + 750M output distillation) on a single
+rented RTX 4090.
+
+## Usage
+
+```python
+import numpy as np
+from modeling_timetron import TimeTron
+
+model = TimeTron.from_pretrained("CorteriIntelligence/TimeTron-v2-33M")
+
+context = np.random.randn(4, 512).astype("float32")     # (batch, history)
+quantiles = model.predict(context, prediction_length=128)   # (4, 128, 21), input's own scale
+median = model.predict_median(context, prediction_length=128)  # (4, 128)
+```
+
+- **21 quantile levels**: `[0.01, 0.05, 0.1, 0.15 … 0.9, 0.95, 0.99]` — index 10 is the median
+- **384 steps decoded natively**; longer horizons extend autoregressively in whole chunks
+- **Any context length** — cropped to a multiple of 32 and left-padded as needed, up to 2048
+- Output is in the **input's own scale**; no manual normalization required
+
+Requires `torch`, `numpy`, and `safetensors`. `modeling_timetron.py` is self-contained.
+
+## Model details
+
+| | |
+|---|---|
+| parameters | 33.52M |
+| hidden / layers / heads | 512 / 12 / 16 |
+| patch length | 32 |
+| max context | 2048 (trained at ≤512) |
+| native horizon | 384 |
+| attention | bidirectional, RoPE + QK-norm |
+| normalization | causal patch-norm, asinh-compressed, μ-anchored |
+| output | 21 quantiles per future patch |
+| teacher | Chronos-2 (~120M) |
+
+Forecasts are produced in `z = asinh((y − μ)/σ)` space and inverted as `ŷ = μ + σ·sinh(z)`,
+which is why the model handles wide dynamic ranges without manual scaling.
+
+### Training
+
+Latent KD at three taps (student layers 4/8/12 → teacher layers 4/8/final) with SmoothL1
+against LayerNorm'd, pooled teacher representations, plus a pinball loss on 21 quantiles and a
+masked-view consistency term. The final checkpoint is a weight average of one latent-distilled
+checkpoint and two output-distilled variants (`W_OUT` 0.2 and 0.5) — averaging over
+*objective-diverse* checkpoints was the single most reliable source of gain in the project.
+
+## Limitations
+
+- **Test-data leakage: declared `Yes`.** Training pools contain series that are GIFT-Eval test
+  datasets (m4 family, LOOP_SEATTLE, SZ_TAXI, hierarchical_sales, restaurant,
+  temperature_rain). Long series used a 10% tail holdout and short series a last-window
+  holdout, but partial overlap with official test horizons remains possible. We declare it
+  rather than argue the edge case.
+- **Univariate.** The architecture contains a group-attention branch for in-context learning
+  across related series, and `predict(..., group_ids=...)` exposes it — but **on this
+  checkpoint it is untrained and passing `group_ids` is a no-op**. Training it on a frozen
+  backbone measured **3.4% worse** on GIFT's 43 multivariate configurations; the capability
+  turned out to be inseparable from the backbone it co-adapts with.
+- **Known-future covariates** are supported by the architecture (`future_values`) with zero
+  additional parameters, but are only lightly trained (~5% of rows).
+- **Trained at context ≤512.** Longer contexts are accepted but untested.
+- **Horizons beyond 384** use autoregressive chunk extension, not a native long-horizon head.
+
+## Evaluation
+
+Numbers above are the official GIFT-Eval harness (gluonts metric engine, official windowing
+and seasonality), aggregated as the geometric mean of per-configuration
+`model_MASE / seasonal_naive_MASE`. Submission artifacts are in `gifteval_submission/`.
+
+On **IEX** Indian electricity spot prices (private data, clean for every model compared),
+TimeTron scores **1.189 MASE / 862 CRPS** versus TimesFM-2.5's 1.232 / 908 — better on both at
+1/7 the size.
+
+## Citation
+
+```bibtex
+@misc{srivastava2026timetron,
+  title  = {TimeTron: Latent-Space Distillation of a Time-Series Foundation Model at 33M Parameters},
+  author = {Srivastava, Aditya},
+  year   = {2026},
+  note   = {Corteri Intelligence}
+}
+```
+
+## Acknowledgements
+
+Distilled from [amazon/chronos-2](https://huggingface.co/amazon/chronos-2). Evaluated with
+[GIFT-Eval](https://huggingface.co/datasets/Salesforce/GiftEval). Pretraining corpora: LOTSA,
+Time-300B, GIFT-Eval-Pretrain.

config.json

@@ -0,0 +1,48 @@
+{
+  "patch_length": 32,
+  "max_context": 2048,
+  "quantile_levels": [
+    0.01,
+    0.05,
+    0.1,
+    0.15,
+    0.2,
+    0.25,
+    0.3,
+    0.35,
+    0.4,
+    0.45,
+    0.5,
+    0.55,
+    0.6,
+    0.65,
+    0.7,
+    0.75,
+    0.8,
+    0.85,
+    0.9,
+    0.95,
+    0.99
+  ],
+  "teacher_dim": 768,
+  "hidden_size": 512,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 16,
+  "head_dim": 32,
+  "intermediate_size": 1280,
+  "attention_dropout": 0.0,
+  "group_layers": [
+    5,
+    11
+  ],
+  "kd_taps": [
+    3,
+    7,
+    11
+  ],
+  "max_future_patches": 12,
+  "architectures": [
+    "TimeTron"
+  ],
+  "model_type": "timetron"
+}

gifteval_submission/all_results.csv

@@ -0,0 +1,98 @@
+dataset,model,eval_metrics/MSE[mean],eval_metrics/MSE[0.5],eval_metrics/MAE[0.5],eval_metrics/MASE[0.5],eval_metrics/MAPE[0.5],eval_metrics/sMAPE[0.5],eval_metrics/MSIS,eval_metrics/RMSE[mean],eval_metrics/NRMSE[mean],eval_metrics/ND[0.5],eval_metrics/mean_weighted_sum_quantile_loss,domain,num_variates
+m4_yearly/short,TimeTron-v2-33M,3594296.0624474045,3594296.0624474045,915.4947640085894,3.4231836274618535,0.17810498389235804,0.15368990387408332,31.170788996269206,1895.8628807082553,0.30401040089771975,0.1468038290417079,0.11831008302449424,-,1
+m4_quarterly/short,TimeTron-v2-33M,1973282.3586666666,1973282.3586666666,597.1541953125,1.2458785657638598,0.12443442296981812,0.10779740524291992,9.521712062686712,1404.7356899668587,0.2351280179087623,0.09995309675163348,0.0800003569693292,-,1
+m4_monthly/short,TimeTron-v2-33M,1908314.2554444445,1908314.2554444445,561.6873668981482,0.9557448071829153,0.15690245896798594,0.13056823190053304,7.934129578024344,1381.4174805048779,0.28710597717780223,0.11673791783985618,0.09357452600595556,-,1
+m4_weekly/short,TimeTron-v2-33M,433262.9955003214,433262.9955003214,330.8395650310692,2.480185797433929,0.07958839208890213,0.08251528089092966,19.24135863018563,658.2271610168646,0.1199187879677542,0.06027384161573412,0.047081977600511166,-,1
+m4_daily/short,TimeTron-v2-33M,367110.8017168542,367110.8017168542,183.26163254292135,3.3867674046981024,0.04095082452392372,0.03180903770242821,26.69460591322561,605.8966922808328,0.09359429472407577,0.028308857708514115,0.022078856026623284,-,1
+m4_hourly/short,TimeTron-v2-33M,685281.5974235105,685281.5974235105,210.77577747584542,1.1031086366965892,0.11508565709210825,0.10825852840994866,24.196257021037926,827.8173696072766,0.1130153431301189,0.028775546019590148,0.025145577127351135,-,1
+electricity/15T/short,TimeTron-v2-33M,163404.54274774774,163404.54274774774,61.09176625844594,1.1563833098585172,0.19761363097197637,0.19085472313133445,11.367411721119591,404.23327763526316,0.7968234703250702,0.12042391334803951,0.0980929782651497,-,1
+electricity/15T/medium,TimeTron-v2-33M,294249.6370810811,294249.6370810811,63.85813302364865,1.0237698293784556,0.18383462717458773,0.16589991774000562,10.073302640228203,542.4478196850653,0.9346922168820937,0.1100339936041762,0.08957952794482943,-,1
+electricity/15T/long,TimeTron-v2-33M,518713.4781621622,518713.4781621622,77.11960810810811,1.1372180626950643,0.19313267492605388,0.17414501879809496,11.51863970330012,720.2176602681735,1.1362316249577773,0.12166563314228529,0.0994076143807144,-,1
+electricity/H/short,TimeTron-v2-33M,1863850.5232432433,1863850.5232432433,195.29604940878377,1.0531645319905507,0.2071355585941179,0.14245030583562077,11.002882542024434,1365.2291101654855,0.6444856363045942,0.09219368216938824,0.07594032492622362,-,1
+electricity/H/medium,TimeTron-v2-33M,8526295.972612612,8526295.972612612,296.62565878378376,1.290791834567778,0.3714921230350138,0.16095980532534487,12.791498600819262,2919.9821870368683,1.1397652734927797,0.11578276970646689,0.09493813854282841,-,1
+electricity/H/long,TimeTron-v2-33M,14490457.035723723,14490457.035723723,362.5508648648649,1.4928027759808222,0.4299999475456491,0.18107912209084084,17.69269712835727,3806.6332940964676,1.50133504528955,0.14298995386968408,0.12057306061017156,-,1
+electricity/D/short,TimeTron-v2-33M,1887535838.4755316,1887535838.4755316,5021.312864864864,1.5952281502252055,0.5084715135332953,0.10675854712134009,12.853891913920549,43445.78044500446,0.7139584565578992,0.08251684619709619,0.06499693054125838,-,1
+electricity/W/short,TimeTron-v2-33M,89720899100.5982,89720899100.5982,34812.19639639639,1.7566177439545647,0.24587304596443935,0.11112677084433066,13.001777909154436,299534.47063835274,0.6834388152975958,0.07942994411413172,0.06352301038428682,-,1
+solar/10T/short,TimeTron-v2-33M,31.00848587667275,31.00848587667275,2.868869040830292,1.2593578374526235,6.612096596915852,1.5710907799194038,8.032930284337072,5.568526364907752,1.619883716517713,0.8345536932983104,0.6430746612258154,-,1
+solar/10T/medium,TimeTron-v2-33M,51.26058808891838,51.26058808891838,3.8030058511667773,1.6492573662459555,3.236224299761621,1.6290907678621542,13.471872687664497,7.1596499976548005,1.5887547069407746,0.8439020690318162,0.6972475910796475,-,1
+solar/10T/long,TimeTron-v2-33M,58.21633116636253,58.21633116636253,4.165829777156833,1.8140158841554672,2.01724133506689,1.685170741041592,11.931839350309946,7.629962723785912,1.649094210892775,0.9003773698209266,0.6987511382871364,-,1
+solar/H/short,TimeTron-v2-33M,735.1190453323089,735.1190453323089,11.909288361185812,0.917716481731851,5.787753402215246,1.4176168423453708,7.109408682997376,27.113078861175264,0.9993971446194256,0.4389803475864836,0.341090753933536,-,1
+solar/H/medium,TimeTron-v2-33M,776.3046836982968,776.3046836982968,13.254836716849148,1.0129880353501837,6.741881029693307,1.417995480725365,9.268599022927534,27.862244771344194,1.0036630874613721,0.47747015548120747,0.38215533668975454,-,1
+solar/H/long,TimeTron-v2-33M,685.0357664233577,685.0357664233577,13.762947333738849,1.0480795370654432,7.43866723652702,1.4178152562347932,8.286784708555352,26.173187930081383,0.9074721248154772,0.4771864663194915,0.37786528266317465,-,1
+solar/D/short,TimeTron-v2-33M,153507.20778588808,153507.20778588808,258.97734184914844,1.0119912773093018,1.1992364516803529,0.4241537349357512,5.843921754961797,391.7999588895947,0.566017491815129,0.3741340501562192,0.28898256167581393,-,1
+solar/W/short,TimeTron-v2-33M,3981635.270072993,3981635.270072993,1571.5380930656934,1.7341981981769463,0.36490653379120097,0.28844262561658873,16.376238023731513,1995.4035356471113,0.4073117491760472,0.3207902151861179,0.27086939869406473,-,1
+hospital/short,TimeTron-v2-33M,3696.574532811821,3696.574532811821,19.3916028357236,0.7952632608462912,0.19962461654749294,0.1794748849218071,5.303300800459719,60.79946161613457,0.22071616416305384,0.07039602129862446,0.05589563089631618,-,1
+covid_deaths/short,TimeTron-v2-33M,518170.53032581456,518170.53032581456,117.59174498746867,36.48802357941842,0.0989387894061914,0.37541412124060153,696.8008990667008,719.8406284211906,0.2708435395908337,0.04424446631877708,0.037168792781997355,-,1
+us_births/D/short,TimeTron-v2-33M,286170.9866666667,286170.9866666667,373.7851822916667,0.5500547775149256,0.03625814437866211,0.0357220999399821,4.665746882646273,534.9495178674963,0.05014717643511761,0.03503932774803556,0.027926156250498005,-,1
+us_births/M/short,TimeTron-v2-33M,123431370.66666667,123431370.66666667,9832.345052083334,1.1144189334005354,0.030072328945000965,0.030647508800029755,6.475356301731352,11109.967176669186,0.034507482303935216,0.030539196695689153,0.022935175096698842,-,1
+us_births/W/short,TimeTron-v2-33M,3787263.4285714286,3787263.4285714286,1545.3914620535713,1.4050179544851435,0.020756800259862627,0.020955807396343777,8.0164218744744,1946.0892653142682,0.026417970953975998,0.020978537564905,0.015888432448803228,-,1
+saugeenday/D/short,TimeTron-v2-33M,1125.3322916666666,1125.3322916666666,13.250651041666666,2.9439087446346712,0.3210107930501302,0.3564278157552083,28.627443692395936,33.545972808470864,1.086562052111618,0.4291917444114977,0.3710848471072012,-,1
+saugeenday/M/short,TimeTron-v2-33M,498.0833798363095,498.0833798363095,13.423924037388392,0.7864016251801428,0.3948113577706473,0.381299200512114,5.637896915705379,22.317781696134354,0.6700467592359634,0.40302647099733035,0.32475039903101705,-,1
+saugeenday/W/short,TimeTron-v2-33M,1168.4376953125,1168.4376953125,17.248260498046875,1.3991118215013922,0.44062271118164065,0.4730045795440674,11.409062990882681,34.18241792665493,1.0318247086623504,0.5206530854990096,0.42115961814366754,-,1
+temperature_rain_with_missing/short,TimeTron-v2-33M,189.75097770980202,189.75097770980202,5.924746265797922,1.3810640379397443,18.04182227835526,1.5441743857327364,18.673439114099658,13.7750128025277,1.6216831887238623,0.6974992730994705,0.5597701566827407,-,1
+kdd_cup_2018_with_missing/H/short,TimeTron-v2-33M,4259.035931336197,4259.035931336197,22.183533911748057,0.9330613507973277,0.904529729892268,0.48666915032836133,7.5492096724842614,65.26128968489817,1.366036990727736,0.4643415423572821,0.3740733885446834,-,1
+kdd_cup_2018_with_missing/H/medium,TimeTron-v2-33M,5181.063323577293,5181.063323577293,25.800004826826626,1.0693786593853554,1.001936864458849,0.5725511122544352,13.02182780897339,71.97960352472978,1.5066634629076956,0.5400408270110403,0.4485085555496209,-,1
+kdd_cup_2018_with_missing/H/long,TimeTron-v2-33M,3987.3993707559716,3987.3993707559716,23.717143384442323,1.0111767484474288,1.0197800990360104,0.5997961841095377,15.277030597942177,63.145857906564004,1.4816569922308815,0.5565000222706372,0.47802813297063845,-,1
+kdd_cup_2018_with_missing/D/short,TimeTron-v2-33M,2858.4746326218497,2858.4746326218497,21.054603863288783,1.1868401460737816,0.5611588905304542,0.45768057088029457,9.43914393719036,53.46470455002861,1.1976661155692032,0.47164546846410893,0.3760212369881984,-,1
+car_parts_with_missing/short,TimeTron-v2-33M,1.3371987472702438,1.3371987472702438,0.47970586730867837,0.8760970166837794,0.8232178409906193,1.8807302824415437,15.088876264544554,1.1563731003747206,2.772864073437567,1.150285461367449,0.9847379481678661,-,1
+restaurant/short,TimeTron-v2-33M,146.14415802745867,146.14415802745867,7.306213224992995,0.697099558956002,0.6934012394052956,0.3999791863675633,4.62975270427469,12.089009803431324,0.5563642411081943,0.33624886094013273,0.263262385601512,-,1
+hierarchical_sales/D/short,TimeTron-v2-33M,30.129841876808598,30.129841876808598,2.333401747881356,0.7503448683624596,0.6404265107711385,1.0862491602935098,7.202072463023026,5.489065665193722,1.683984558346157,0.715861815384469,0.583622752136203,-,1
+hierarchical_sales/W/short,TimeTron-v2-33M,483.6865730932203,483.6865730932203,9.09037651449947,0.7353939183747501,0.5282679919091651,0.4681441096936242,6.691019418601793,21.99287550760974,1.009703672069916,0.41734363221454884,0.3532952066098958,-,1
+LOOP_SEATTLE/5T/short,TimeTron-v2-33M,44.87091027920214,44.87091027920214,3.6567248774509804,0.5777993204241603,0.10276269233632752,0.0783165741754144,4.521999058360721,6.698575242482699,0.11486578013914665,0.06270476043006352,0.04974600871519414,-,1
+LOOP_SEATTLE/5T/medium,TimeTron-v2-33M,140.09096604102166,140.09096604102166,6.981449041376419,1.0877889603000435,0.25860870833736455,0.15073082230039425,13.17282271721858,11.836002958812644,0.21058680812172076,0.1242143208990372,0.1036475416391804,-,1
+LOOP_SEATTLE/5T/long,TimeTron-v2-33M,153.02013186561175,153.02013186561175,7.414257610939113,1.1548136782559684,0.2810150651086458,0.15716604272374154,14.667447215113286,12.370130632520084,0.21870643941533616,0.13108559086139604,0.11003073929030986,-,1
+LOOP_SEATTLE/H/short,TimeTron-v2-33M,71.8999774251806,71.8999774251806,5.065310448322334,1.0202959899270712,0.15539988380545258,0.1146176434386839,8.490333468866238,8.479385439121199,0.1500937450734167,0.0896611459175514,0.07227682702429591,-,1
+LOOP_SEATTLE/H/medium,TimeTron-v2-33M,82.6188338493292,82.6188338493292,5.770919198271414,1.1690365541956667,0.16793463581978843,0.12728610358617132,11.607466021262942,9.089490296453878,0.16065570539054488,0.10200033932725991,0.08396297938789848,-,1
+LOOP_SEATTLE/H/long,TimeTron-v2-33M,90.46640866873065,90.46640866873065,5.808141447368421,1.1930645398856634,0.20059537726823184,0.1298614806394049,12.928726191176958,9.511383110185955,0.16998115011047368,0.10379926366026627,0.08585839249944315,-,1
+LOOP_SEATTLE/D/short,TimeTron-v2-33M,17.96810178018576,17.96810178018576,2.887948876096491,0.8743414119001611,0.05329589213872227,0.0532498417747033,7.613969183937755,4.238879778925767,0.07575445088110062,0.051611508863510194,0.042425785769131905,-,1
+SZ_TAXI/15T/short,TimeTron-v2-33M,17.465227995745575,17.465227995745575,2.8107664972288995,0.5608097051373252,1347438975473.752,0.40223191742204195,4.120360863303886,4.179142016699789,0.39076765201153413,0.262818688641306,0.2074312155513682,-,1
+SZ_TAXI/15T/medium,TimeTron-v2-33M,17.036907385149572,17.036907385149572,2.8210557725694443,0.5492299800126319,12427031760132.814,0.41180620764055825,6.265977546058238,4.127578876914355,0.38414010999418646,0.2625463272999124,0.21516472919583218,-,1
+SZ_TAXI/15T/long,TimeTron-v2-33M,16.019083867521367,16.019083867521367,2.734242565883191,0.515921741945987,8100707226050.443,0.4004795867832977,7.951909501922769,4.002384772547658,0.37045360503017744,0.25307662134483827,0.21695880951351135,-,1
+SZ_TAXI/H/short,TimeTron-v2-33M,7.280379857772436,7.280379857772436,1.8652708917601495,0.5683504179155325,0.9885649721846621,0.298435308994391,4.179242750420337,2.6982179040567567,0.2513556481627869,0.17376149394481585,0.13779760935597798,-,1
+M_DENSE/H/short,TimeTron-v2-33M,54860.36,54860.36,108.45819444444444,0.9227931542878034,0.35844712602851964,0.2624811299641927,10.428949198517122,234.22288530372091,0.41545303392668964,0.19237781089463618,0.15794283276784504,-,1
+M_DENSE/H/medium,TimeTron-v2-33M,61251.61777777778,61251.61777777778,118.90834201388888,0.9804366936927182,0.41260201134720575,0.2795029873318142,9.357737301349488,247.49064179838754,0.4306024496088221,0.20688549263133296,0.1676531668236521,-,1
+M_DENSE/H/long,TimeTron-v2-33M,68924.61827160494,68924.61827160494,129.64962962962963,1.0663441665405762,0.46777280256354686,0.29832169897762345,9.897417928535283,262.5349848526953,0.45448686850785247,0.224442674589197,0.18147393114763583,-,1
+M_DENSE/D/short,TimeTron-v2-33M,10296.266666666666,10296.266666666666,52.76764467592593,0.7695067169128392,0.13674998018280324,0.11043270534939235,8.126480407768868,101.470521170765,0.17566048160282852,0.0913485982911326,0.07517903037508732,-,1
+ett1/15T/short,TimeTron-v2-33M,5.917625790550595,5.917625790550595,1.2054476056780135,0.775313089475921,0.5100240519957006,0.2584114801316034,6.078214665925738,2.4326170661554185,0.45986588356987596,0.22787977441859086,0.1813064205080266,-,7
+ett1/15T/medium,TimeTron-v2-33M,12.08630580357143,12.08630580357143,1.8101517547123016,1.1515004868310985,0.8091579541839798,0.4297645786830357,11.575991215071925,3.4765364665959466,0.6626312879329936,0.34501671422236974,0.2826016117883024,-,7
+ett1/15T/long,TimeTron-v2-33M,11.20203869047619,11.20203869047619,1.8216586061507936,1.194893964850682,0.891450765050001,0.435021003844246,13.496783815849204,3.3469446799246905,0.6379309652605284,0.34720993148386087,0.2885558241029272,-,7
+ett1/H/short,TimeTron-v2-33M,99.32315848214286,99.32315848214286,5.089588564918155,0.8505728650343709,0.4962696983700707,0.26831755865187873,6.108706601500973,9.966100465184105,0.4651586817159045,0.23755192069398193,0.1878610611502167,-,7
+ett1/H/medium,TimeTron-v2-33M,141.67676711309525,141.67676711309525,7.130074055989583,1.3149084571201797,4011442887069.745,0.4623036702473958,13.687392746323486,11.902805010294642,0.5701161444933562,0.34151364550096613,0.27618217324753486,-,7
+ett1/H/long,TimeTron-v2-33M,163.6448578042328,163.6448578042328,7.635671916335979,1.4118754317721867,630978605448.0175,0.5010712477265211,19.886371608029233,12.792374986851847,0.6125991317596777,0.3656557903561202,0.3048880888607935,-,7
+ett1/D/short,TimeTron-v2-33M,60128.0253968254,60128.0253968254,160.4562251984127,1.8551646321145863,1.2022953093998017,0.5045923142206101,10.507672791598937,245.21016576974412,0.6440526340985293,0.42144359787099683,0.31864046625948955,-,7
+ett1/W/short,TimeTron-v2-33M,1091585.5714285714,1091585.5714285714,776.6309291294643,1.4253502332603116,0.7225370407104492,0.4580290658133371,7.849473192064395,1044.7897259394215,0.4157467687301662,0.30903998313264675,0.23815349942193847,-,7
+ett2/15T/short,TimeTron-v2-33M,9.178926013764881,9.178926013764881,1.9383880615234375,0.8392014895126303,0.11996391857777632,0.13934214455740793,6.5139096052630805,3.0296742421859286,0.14382941420903225,0.09202217701053683,0.07369685746006278,-,7
+ett2/15T/medium,TimeTron-v2-33M,13.818829365079365,13.818829365079365,2.3433723958333332,0.9820298491127656,0.155606453598294,0.18363314189608135,10.609390043849288,3.7173686076416157,0.19810530202639062,0.12488255678561565,0.10203793801749834,-,7
+ett2/15T/long,TimeTron-v2-33M,13.403065476190477,13.403065476190477,2.326241164434524,0.9764523064657172,0.1533723890105101,0.1745628138950893,10.654336040311875,3.661019731740117,0.1951023684307373,0.1239696024545954,0.10207835901799667,-,7
+ett2/H/short,TimeTron-v2-33M,112.21834077380953,112.21834077380953,6.641959635416667,0.7655689452719162,0.12441594022488449,0.11462356022426061,5.397166436430102,10.593315853584727,0.13435271990094289,0.08423852878780462,0.06681312313047659,-,7
+ett2/H/medium,TimeTron-v2-33M,255.79183407738094,255.79183407738094,9.802804129464286,1.0575185100318665,0.17742592120953204,0.17193174816313245,11.294579492534332,15.99349349196044,0.21552260983496177,0.13209908959197075,0.10930294497295005,-,7
+ett2/H/long,TimeTron-v2-33M,271.2439153439153,271.2439153439153,10.156983713624339,1.1009622187511445,0.189720385805821,0.18205826370804398,17.48313352304711,16.469484367882174,0.22553691256682565,0.1390920745660805,0.12122977268560248,-,7
+ett2/D/short,TimeTron-v2-33M,110018.89523809524,110018.89523809524,195.96488095238095,1.2583994408667627,0.4202055915953621,0.13003619966052826,13.125234001646534,331.69096345558654,0.18983548975651104,0.1121558717280305,0.09204130838282955,-,7
+ett2/W/short,TimeTron-v2-33M,7022237.714285715,7022237.714285715,1770.156529017857,1.141450760753593,0.1705326693398612,0.1983142750603812,8.899844991492236,2649.9505116672867,0.22229022171147653,0.1484889946310067,0.11673145907121207,-,7
+jena_weather/10T/short,TimeTron-v2-33M,1238.1079365079365,1238.1079365079365,8.135577876984128,0.33229472781442576,0.42045986491143233,0.5884432353670634,2.6148980093312026,35.18675797097448,0.21855386731621268,0.05053213510985654,0.03928677467098002,-,21
+jena_weather/10T/medium,TimeTron-v2-33M,2141.3916305916305,2141.3916305916305,12.292389294733045,0.6646123478151146,1.0742200969677758,0.6972859172077922,8.066477157393967,46.27517293961883,0.28398994524120563,0.07543818295935632,0.0666566483462141,-,21
+jena_weather/10T/long,TimeTron-v2-33M,2295.7854497354497,2295.7854497354497,13.063463128306879,0.6677037697991172,0.8745041023535293,0.6822967044890873,9.20122181438605,47.91435536178536,0.2931961564035924,0.07993757089330292,0.06884325158711575,-,21
+jena_weather/H/short,TimeTron-v2-33M,1199.701963241437,1199.701963241437,8.984033162332915,0.5411756544463394,1.5899653235869151,0.6141557303088451,4.813406627946494,34.63671409417234,0.21232082759622908,0.05507154492173466,0.04466766941557371,-,21
+jena_weather/H/medium,TimeTron-v2-33M,1633.3838293650795,1633.3838293650795,11.545520213293651,0.8423404648390338,2.8865114449375113,0.6809031653025793,9.070069493842523,40.415143564820845,0.24730743122391843,0.07064908581888237,0.058096831004421805,-,21
+jena_weather/H/long,TimeTron-v2-33M,1196.515873015873,1196.515873015873,12.148456101190476,1.0240500618077049,2.882397952759871,0.657022724454365,16.184287220447,34.59069055419208,0.20827428105029874,0.07314716531555115,0.06183056710404284,-,21
+jena_weather/D/short,TimeTron-v2-33M,471.8342261904762,471.8342261904762,10.69546130952381,1.1044140362257397,0.6010409644009095,0.4685462103949653,7.327962382098359,21.721745468319902,0.13080041026605907,0.0644041580042822,0.051402151070832756,-,21
+bitbrains_fast_storage/5T/short,TimeTron-v2-33M,1857179.6618400447,1857179.6618400447,173.284798710912,0.7735176210732811,1.9193502390981823,0.7844275553589883,15.586743984302542,1362.7837913036847,4.278804119742369,0.544071418624432,0.4175159038020963,-,2
+bitbrains_fast_storage/5T/medium,TimeTron-v2-33M,3328777.795826032,3328777.795826032,283.1186354697321,1.075164345736633,4.128901214134457,0.8305642721538105,25.751099718839953,1824.4938464752443,5.543316530787815,0.8601926584764452,0.6841183786820589,-,2
+bitbrains_fast_storage/5T/long,TimeTron-v2-33M,4298508.070369263,4298508.070369263,371.9540902735552,0.9944457072181663,5.703163649105187,0.8427418743036303,22.812084193233666,2073.2843679460043,5.479071587717764,0.98296361052114,0.7793536751062298,-,2
+bitbrains_fast_storage/H/short,TimeTron-v2-33M,3189831.0076515093,3189831.0076515093,336.4742050741588,1.166380663555333,4.324445060044518,0.6063532365354214,19.936113829448857,1786.0098005474408,5.090707955987181,0.9590607578023669,0.6832619544848516,-,2
+bitbrains_rnd/5T/short,TimeTron-v2-33M,1863694.2786425638,1863694.2786425638,137.64595141625762,1.8174156722897612,1.283698269665347,0.7178146995356822,56.335901043128345,1365.1718861163833,5.585889232551913,0.563207494777613,0.45632781133772493,-,2
+bitbrains_rnd/5T/medium,TimeTron-v2-33M,2225068.469509134,2225068.469509134,166.59690113761093,4.621169289065555,1.0616097361912498,0.7737502938154637,167.02357055287717,1491.666339872672,6.276199866200281,0.7009579962222026,0.6375884086196214,-,2
+bitbrains_rnd/5T/long,TimeTron-v2-33M,2526342.553647849,2526342.553647849,190.6982169619402,3.533227292535709,1.8235743499590866,0.7396551721067658,123.04844364053936,1589.4472478342427,6.088814574418669,0.7305219373185541,0.6479502546127784,-,2
+bitbrains_rnd/H/short,TimeTron-v2-33M,1845780.2303397127,1845780.2303397127,171.3636513788841,5.940762081040158,2.1772018881073696,0.6309933032146434,195.4189426963251,1358.5949471198958,5.945876466236105,0.749971214056103,0.6671548054565138,-,2
+bizitobs_application/short,TimeTron-v2-33M,3846447.7866666666,3846447.7866666666,802.4288888888889,2.7888631318616226,0.039902966817220054,0.039473953247070315,45.04911043953967,1961.2362903706087,0.07563704443410081,0.030946474844510446,0.02712425175052398,-,2
+bizitobs_application/medium,TimeTron-v2-33M,21395602.773333333,21395602.773333333,2128.1989583333334,6.954895414358934,0.09381882349650066,0.0958718999226888,136.6775957490632,4625.538106354042,0.18163126766444423,0.08356810942995752,0.07643752398640181,-,2
+bizitobs_application/long,TimeTron-v2-33M,20505351.964444444,20505351.964444444,2293.850972222222,7.430197088147029,0.09910538567437066,0.10250606113009983,178.18038105920755,4528.283556099866,0.1746377966921116,0.08846466321423808,0.08411196698530009,-,2
+bizitobs_service/short,TimeTron-v2-33M,17581.621164021162,17581.621164021162,39.59574074074074,1.7647892474529427,0.31517190600198414,0.11039424189814814,27.404425159391323,132.59570567714914,0.09823113774379423,0.029333790584754743,0.025524187753620836,-,2
+bizitobs_service/medium,TimeTron-v2-33M,140078.9942857143,140078.9942857143,107.75006944444445,3.8227884474889127,0.4408571273561508,0.15915798611111112,75.51363819266368,374.27128434561246,0.2816853883423936,0.08109524140609913,0.07299651699820946,-,2
+bizitobs_service/long,TimeTron-v2-33M,209060.09735449735,209060.09735449735,125.97153439153439,3.752276917693752,0.19191491505456348,0.16670787347056878,88.06615924411318,457.23090157435485,0.3387312691907938,0.09332373988597772,0.08848903912963689,-,2
+bizitobs_l2c/5T/short,TimeTron-v2-33M,22.087946428571428,22.087946428571428,2.826266624813988,0.2945102086164606,0.16528732125067702,0.21728590102422804,3.0689940270310605,4.699781529876834,0.1614095210159594,0.09706543576006887,0.07911019299634953,-,7
+bizitobs_l2c/5T/medium,TimeTron-v2-33M,191.80221088435374,191.80221088435374,10.06801990327381,0.9621880447316018,0.6271967291070425,0.9746027317176871,19.152805108628495,13.849267521582279,0.7286751339341283,0.5297259035567636,0.46650916891333793,-,7
+bizitobs_l2c/5T/long,TimeTron-v2-33M,364.4816666666667,364.4816666666667,13.914151785714285,1.294419370846612,1.0347643080435767,1.2360301184275793,33.96116123183925,19.09140295176514,1.0719537340429284,0.78125882107357,0.7242053312735177,-,7
+bizitobs_l2c/H/short,TimeTron-v2-33M,211.21521577380952,211.21521577380952,9.249025375124008,0.897339827475659,0.7445551371402878,0.8421546088324653,5.28505659256937,14.533245190727689,0.7833791360644242,0.4985461548825487,0.39268200667278613,-,7
+bizitobs_l2c/H/medium,TimeTron-v2-33M,272.71651785714283,272.71651785714283,12.181212797619047,1.1821341226043072,1.3437424342595352,1.0496454148065477,12.348168686730247,16.51413085382161,0.9999557067232343,0.737590937095577,0.5831507966902421,-,7
+bizitobs_l2c/H/long,TimeTron-v2-33M,241.733060515873,241.733060515873,11.579703000992064,1.1563342876318414,1.4209144342336193,0.9612372504340277,11.577100308789067,15.547767058837517,0.9497024855728949,0.7073216803815635,0.5734519320448719,-,7

gifteval_submission/config.json

@@ -0,0 +1,9 @@
+{
+  "model": "TimeTron-v2-33M",
+  "model_type": "pretrained",
+  "model_dtype": "float32",
+  "model_link": "https://huggingface.co/CorteriIntelligence/TimeTron-v2-33M",
+  "org": "Corteri Intelligence",
+  "testdata_leakage": "Yes",
+  "replication_code_available": "No"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd9a9f840864f8877890020c72ea2481676b0c770e82a85057548e882ce89373
+size 184445376

modeling_timetron.py

@@ -0,0 +1,428 @@
+"""TimeTron — a 33.5M-parameter time-series foundation model.
+
+Distilled from Chronos-2 by latent-space knowledge distillation: the student matches the
+teacher's internal representations rather than its forecasts.
+
+    from modeling_timetron import TimeTron
+    model = TimeTron.from_pretrained("timetron-v2-33m")
+    q = model.predict(context, prediction_length=128)   # (B, H, 21) in the input's own scale
+
+Self-contained: torch + numpy only.
+"""
+import json
+import math
+import os
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+QUANTILE_LEVELS_21 = [0.01, 0.05] + [round(0.1 + 0.05 * i, 2) for i in range(17)] + [0.95, 0.99]
+
+
+# ---------------------------------------------------------------- input path
+def patchify(x, mask, patch_size=32):
+    """x,(B,L) mask,(B,L) 1=padding -> x_patched,(B,N,P) m_patched,(B,N,P) patch_mask,(B,N)."""
+    B, L = x.shape
+    P = patch_size
+    mask = mask.to(torch.long)
+    N = (L + P - 1) // P
+    pad = N * P - L
+    x_patched = F.pad(x, (pad, 0), value=0.0).reshape(B, N, P)
+    m_patched = F.pad(mask, (pad, 0), value=1).reshape(B, N, P)
+    patch_mask = m_patched.amin(dim=-1)
+    return x_patched, m_patched, patch_mask
+
+
+class CausalPatchNormV2(nn.Module):
+    """Cumulative per-patch stats + σ-relative floor + asinh squash (+ training dither)."""
+
+    def __init__(self, sigma_min=1e-3, dither=0.01):
+        super().__init__()
+        self.sigma_min = sigma_min
+        self.dither = dither
+
+    def forward(self, x, mask):
+        valid = 1.0 - mask.float()
+        x_valid = x * valid
+        count = valid.sum(-1).cumsum(-1)
+        safe = count.clamp(min=1.0)
+        S1 = x_valid.sum(-1).cumsum(-1)
+        S2 = (x_valid ** 2).sum(-1).cumsum(-1)
+        mu = S1 / safe
+        var = (S2 / safe - mu ** 2).clamp(min=0.0)
+        cum_abs = x_valid.abs().sum(-1).cumsum(-1) / safe
+        floor = torch.maximum(torch.full_like(cum_abs, self.sigma_min), 0.05 * cum_abs)
+        sigma = torch.maximum((var + 1e-8).sqrt(), floor)
+        xn = torch.asinh((x - mu.unsqueeze(-1)) / sigma.unsqueeze(-1))   # smooth, unbounded-safe
+        if self.training and self.dither > 0:
+            xn = xn + self.dither * torch.randn_like(xn)
+        return xn, mu, sigma
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-4):                 # v1 backward-amplifier fix baked in
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        rms = (x.pow(2).mean(-1, keepdim=True) + self.eps).sqrt()
+        return x / rms * self.weight
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(hidden_dim, output_dim, bias=False)
+        self.skip = nn.Linear(input_dim, output_dim, bias=False)
+        self.act = nn.SiLU()
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x))) + self.skip(x)
+
+
+# ---------------------------------------------------------------- RoPE (v1, unchanged)
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim, max_position_embeddings=16384, rope_theta=10000.0):
+        super().__init__()
+        i = torch.arange(0, head_dim, 2).float()
+        inv_freq = 1.0 / (rope_theta ** (i / head_dim))
+        positions = torch.arange(max_position_embeddings).float()
+        angles = torch.outer(positions, inv_freq)
+        self.register_buffer("cos_cached", torch.cat([angles.cos(), angles.cos()], dim=-1))
+        self.register_buffer("sin_cached", torch.cat([angles.sin(), angles.sin()], dim=-1))
+
+    def forward(self):
+        return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    cos = cos[position_ids].unsqueeze(1)
+    sin = sin[position_ids].unsqueeze(1)
+    return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
+
+
+# ---------------------------------------------------------------- attention
+class Attention(nn.Module):
+    """v1 attention (QK-norm, per-dim scale, RoPE) with a causal switch — v2 runs bidirectional."""
+
+    def __init__(self, hidden_size, num_heads, head_dim, causal=False, dropout=0.0):
+        super().__init__()
+        assert num_heads * head_dim == hidden_size
+        self.h, self.hd, self.hidden = num_heads, head_dim, hidden_size
+        self.causal = causal
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+        self.scale = nn.Parameter(torch.ones(head_dim))
+        self.rotary = RotaryEmbedding(head_dim=head_dim)
+        self.dropout = dropout
+
+    def forward(self, x, token_mask, position_ids):
+        B, N, _ = x.shape
+        q = self.q_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
+        k = self.k_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
+        v = self.v_proj(x).view(B, N, self.h, self.hd).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        cos, sin = self.rotary()
+        q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
+        q = q * self.scale
+        blocked = token_mask.bool()[:, None, None, :]
+        if self.causal:
+            blocked = blocked | torch.ones(N, N, dtype=torch.bool, device=x.device).triu(1)
+        blocked = blocked & ~blocked.all(dim=-1, keepdim=True)
+        attn_mask = torch.zeros(blocked.shape, dtype=q.dtype, device=x.device).masked_fill(blocked, float("-inf"))
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
+                                             dropout_p=self.dropout if self.training else 0.0, scale=1.0)
+        return self.o_proj(out.transpose(1, 2).contiguous().view(B, N, self.hidden))
+
+
+class GroupAttention(nn.Module):
+    """Chronos-2-style group attention: attends ACROSS series of the same group at each
+    token position. o_proj zero-initialized -> exact identity until the multivariate phase
+    turns it on. Skipped entirely (no compute) when group_ids is None."""
+
+    def __init__(self, hidden_size, num_heads, head_dim):
+        super().__init__()
+        self.h, self.hd, self.hidden = num_heads, head_dim, hidden_size
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        nn.init.zeros_(self.o_proj.weight)                    # neutral start (drift-gate lesson)
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+    def forward(self, x, group_ids):
+        if group_ids is None:
+            return torch.zeros_like(x)
+        B, N, d = x.shape
+        y = x.transpose(0, 1)                                  # (N, B, d): attend across batch per position
+        q = self.q_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
+        k = self.k_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
+        v = self.v_proj(y).view(N, B, self.h, self.hd).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        same = group_ids[:, None] == group_ids[None, :]        # (B, B)
+        attn_mask = torch.zeros(B, B, dtype=q.dtype, device=x.device).masked_fill(~same, float("-inf"))
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        return self.o_proj(out.transpose(1, 2).contiguous().view(N, B, d)).transpose(0, 1)
+
+
+class MLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.fc2(F.silu(self.fc1(x)))
+
+
+class V2Layer(nn.Module):
+    def __init__(self, cfg, with_group=False):
+        super().__init__()
+        self.input_layernorm = RMSNorm(cfg.hidden_size)
+        self.self_attn = Attention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim,
+                                   causal=False, dropout=cfg.attention_dropout)
+        self.group_attn = GroupAttention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim) if with_group else None
+        self.group_norm = RMSNorm(cfg.hidden_size) if with_group else None
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size)
+        self.mixer = MLP(cfg.hidden_size, cfg.intermediate_size)
+
+    def forward(self, x, token_mask, position_ids, group_ids=None):
+        x = x + self.self_attn(self.input_layernorm(x), token_mask, position_ids)
+        if self.group_attn is not None:
+            x = x + self.group_attn(self.group_norm(x), group_ids)
+        x = x + self.mixer(self.post_attention_layernorm(x))
+        return x
+
+
+# ---------------------------------------------------------------- the model
+@dataclass
+class StudentV2Config:
+    patch_length: int = 32
+    max_context: int = 2048                 # time-index normalization constant C
+    quantile_levels: list = field(default_factory=lambda: list(QUANTILE_LEVELS_21))
+    teacher_dim: int = 768                  # Chronos-2 latents
+    hidden_size: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 16
+    head_dim: int = 32
+    intermediate_size: int = 1280
+    attention_dropout: float = 0.0
+    group_layers: tuple = (5, 11)           # 0-indexed layer positions carrying group attention
+    kd_taps: tuple = (3, 7, 11)             # 0-indexed layers tapped for latent distillation
+    max_future_patches: int = 12            # up to 384-step native horizon
+
+
+class StudentV2(nn.Module):
+    def __init__(self, cfg: StudentV2Config = StudentV2Config()):
+        super().__init__()
+        self.cfg = cfg
+        P = cfg.patch_length
+        self.norm_layer = CausalPatchNormV2()
+        self.input_embedding = ResidualBlock(3 * P, cfg.hidden_size, cfg.hidden_size)   # [vals, tidx, mask]
+        self.reg_token = nn.Parameter(torch.zeros(1, 1, cfg.hidden_size))
+        nn.init.trunc_normal_(self.reg_token, std=0.02)
+        self.layers = nn.ModuleList(
+            V2Layer(cfg, with_group=(i in cfg.group_layers)) for i in range(cfg.num_hidden_layers))
+        self.final_norm = RMSNorm(cfg.hidden_size)
+        nq = len(cfg.quantile_levels)
+        self.quantile_head = ResidualBlock(cfg.hidden_size, 2 * cfg.hidden_size, P * nq)
+        self.latent_projs = nn.ModuleList(
+            nn.Linear(cfg.hidden_size, cfg.teacher_dim) for _ in cfg.kd_taps)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, cfg.hidden_size))
+        nn.init.trunc_normal_(self.mask_token, std=0.02)
+
+    def _embed_context(self, x, mask):
+        P = self.cfg.patch_length
+        x_patched, m_patched, patch_mask = patchify(x, mask, P)
+        B, Nc, _ = x_patched.shape
+        xn, mu, sigma = self.norm_layer(x_patched, m_patched)
+        xn = xn * (1.0 - m_patched.float())
+        t = torch.arange(-Nc * P + 1, 1, device=x.device, dtype=xn.dtype) / self.cfg.max_context
+        tidx = t.view(1, Nc, P).expand(B, Nc, P)
+        obs = 1.0 - m_patched.float()
+        h = self.input_embedding(torch.cat([xn, tidx, obs], dim=-1))
+        return h, patch_mask, mu, sigma, Nc
+
+    def forward(self, x, mask, k_future=4, aug_mask=None, group_ids=None,
+                future_values=None, future_observed=None):
+        """x (B, L) L%32==0; k_future future patches decoded natively.
+
+        future_values (B, k_future*P): KNOWN future values, for covariate-informed tasks. They
+        reuse the context's own [values, tidx, observed] embedding path and the context anchor
+        (mu, sigma), so no new parameters are involved. future_observed (B, k_future*P) marks
+        where the value is known; patches with nothing known keep the mask token exactly as
+        before, so future_values=None reproduces the univariate/multivariate path bit-for-bit.
+
+        Returns quantiles (B, k_future*P, n_q) in asinh-normalized anchor space."""
+        cfg = self.cfg
+        P = cfg.patch_length
+        B = x.shape[0]
+        h_ctx, patch_mask, mu, sigma, Nc = self._embed_context(x, mask)
+        if aug_mask is not None:
+            h_ctx = torch.where(aug_mask.unsqueeze(-1), self.mask_token.to(h_ctx.dtype), h_ctx)
+        # future tokens: value channel (known covariates, else 0), future time index, observed flag
+        tf = torch.arange(1, k_future * P + 1, device=x.device, dtype=h_ctx.dtype) / cfg.max_context
+        if future_values is None:
+            fv = torch.zeros(B, k_future, P, device=x.device, dtype=h_ctx.dtype)
+            fo = torch.zeros(B, k_future, P, device=x.device, dtype=h_ctx.dtype)
+        else:
+            fo = (torch.ones_like(future_values) if future_observed is None
+                  else future_observed).to(h_ctx.dtype).view(B, k_future, P)
+            fv = torch.asinh((future_values - mu[:, -1:]) / sigma[:, -1:]).clamp(-4.0, 4.0)
+            fv = fv.to(h_ctx.dtype).view(B, k_future, P) * fo          # unknown slots stay 0
+        fut_feats = torch.cat([fv, tf.view(1, k_future, P).expand(B, k_future, P), fo], dim=-1)
+        # mask token only where the future is genuinely unknown
+        h_fut = self.input_embedding(fut_feats) + \
+            self.mask_token.to(h_ctx.dtype) * (1.0 - fo.amax(dim=-1, keepdim=True))
+        seq = torch.cat([h_ctx, self.reg_token.expand(B, 1, -1).to(h_ctx.dtype), h_fut], dim=1)
+        N = seq.shape[1]
+        token_mask = torch.cat([patch_mask,
+                                torch.zeros(B, 1 + k_future, dtype=patch_mask.dtype, device=x.device)], dim=1)
+        pos = torch.arange(N, device=x.device).unsqueeze(0).expand(B, N)
+        taps = {}
+        for i, layer in enumerate(self.layers):
+            seq = layer(seq, token_mask, pos, group_ids)
+            if i in cfg.kd_taps:
+                taps[i] = seq[:, :Nc]
+        seq = self.final_norm(seq)
+        taps[cfg.kd_taps[-1]] = seq[:, :Nc]                      # final tap post-norm
+        latents = [proj(taps[t]) for t, proj in zip(cfg.kd_taps, self.latent_projs)]
+        h_future = seq[:, Nc + 1:]
+        nq = len(cfg.quantile_levels)
+        q = self.quantile_head(h_future).view(B, k_future * P, nq)
+        return {"quantiles": q, "hidden": seq[:, :Nc], "reg": seq[:, Nc],
+                "latents": latents, "mu": mu, "sigma": sigma, "patch_mask": patch_mask}
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+# --------------------------------------------------------------------- public API
+QUANTILE_LEVELS = QUANTILE_LEVELS_21
+
+
+class TimeTron(nn.Module):
+    """Thin wrapper over StudentV2 adding load/save and a batched `predict`."""
+
+    def __init__(self, config: StudentV2Config = None):
+        super().__init__()
+        self.config = config or StudentV2Config()
+        self.model = StudentV2(self.config)
+        self.median_index = self.config.quantile_levels.index(0.5)
+        self.native_horizon = self.config.max_future_patches * self.config.patch_length
+        self.eval()   # inference must be deterministic: CausalPatchNormV2 dithers in train mode
+
+    # ---- persistence -------------------------------------------------
+    @classmethod
+    def from_pretrained(cls, path_or_repo: str, device: str = "cpu"):
+        """`path_or_repo` may be a local directory or a Hugging Face repo id."""
+        d = path_or_repo
+        if not os.path.isdir(d):
+            from huggingface_hub import snapshot_download
+            d = snapshot_download(path_or_repo)
+        with open(os.path.join(d, "config.json")) as f:
+            raw = json.load(f)
+        cfg = StudentV2Config(**{k: v for k, v in raw.items()
+                                 if k in StudentV2Config.__dataclass_fields__})
+        cfg.group_layers = tuple(cfg.group_layers)
+        cfg.kd_taps = tuple(cfg.kd_taps)
+        self = cls(cfg)
+        w = os.path.join(d, "model.safetensors")
+        if os.path.exists(w):
+            from safetensors.torch import load_file
+            state = load_file(w)
+        else:
+            state = torch.load(os.path.join(d, "pytorch_model.bin"),
+                               map_location="cpu", weights_only=False)
+            state = state.get("model", state)
+        missing, unexpected = self.model.load_state_dict(state, strict=False)
+        assert not missing and not unexpected, (missing, unexpected)
+        return self.to(device).eval()
+
+    def save_pretrained(self, directory: str):
+        os.makedirs(directory, exist_ok=True)
+        cfg = {k: (list(v) if isinstance(v, tuple) else v)
+               for k, v in self.config.__dict__.items()}
+        cfg["architectures"] = ["TimeTron"]
+        cfg["model_type"] = "timetron"
+        with open(os.path.join(directory, "config.json"), "w") as f:
+            json.dump(cfg, f, indent=2)
+        from safetensors.torch import save_file
+        save_file({k: v.contiguous() for k, v in self.model.state_dict().items()},
+                  os.path.join(directory, "model.safetensors"))
+
+    # ---- inference ---------------------------------------------------
+    @torch.no_grad()
+    def predict(self, context, prediction_length: int = 128, group_ids=None):
+        """context: (B, L) tensor/array of raw values, L a multiple of 32 (it is cropped
+        and left-padded for you). Returns (B, prediction_length, 21) in the input's scale.
+
+        Horizons up to `native_horizon` (384) decode in one pass; beyond that the median
+        path is fed back autoregressively in whole chunks.
+
+        `group_ids`: optional (B,) integer tensor. Rows sharing an id attend to each other
+        (in-context learning across related series). Leaving it None is bit-identical to a
+        model without the group-attention branch. NOTE: on the released checkpoint this
+        branch is untrained -- see the model card.
+        """
+        was_training = self.training
+        self.eval()   # dither is a *training* augmentation; never let it reach a forecast
+        try:
+            return self._predict(context, prediction_length, group_ids)
+        finally:
+            if was_training:
+                self.train()
+
+    @torch.no_grad()
+    def _predict(self, context, prediction_length, group_ids):
+        dev = next(self.parameters()).device
+        x = torch.as_tensor(context, dtype=torch.float32, device=dev)
+        if x.dim() == 1:
+            x = x[None, :]
+        L = max(32, min(self.config.max_context, (x.shape[1] // 32) * 32))
+        x = x[:, -L:] if x.shape[1] >= L else F.pad(x, (L - x.shape[1], 0), mode="replicate")
+        g = None if group_ids is None else torch.as_tensor(group_ids, dtype=torch.long, device=dev)
+
+        outs, done = [], 0
+        cur = x
+        while done < prediction_length:
+            need = min(prediction_length - done, self.native_horizon)
+            kf = (need + 31) // 32
+            o = self.model(cur, torch.zeros_like(cur, dtype=torch.long),
+                           k_future=kf, group_ids=g)
+            y = (o["mu"][:, -1:].unsqueeze(-1)
+                 + o["sigma"][:, -1:].unsqueeze(-1)
+                 * torch.sinh(o["quantiles"].clamp(-4.0, 4.0)))[:, :need]
+            outs.append(y)
+            done += need
+            if done < prediction_length:
+                cur = torch.cat([cur, y[:, :, self.median_index]], dim=1)
+                keep = max(32, min(self.config.max_context, (cur.shape[1] // 32) * 32))
+                cur = cur[:, -keep:]
+        return torch.cat(outs, dim=1)
+
+    @torch.no_grad()
+    def predict_median(self, context, prediction_length: int = 128, group_ids=None):
+        return self.predict(context, prediction_length, group_ids)[..., self.median_index]
+
+    def num_params(self):
+        return sum(p.numel() for p in self.model.parameters())