A 31.9M-parameter Llama-architecture causal language model trained from scratch to reconstruct particle-track
parameters from detector hit patterns, treating track fitting as a language translation problem: an array of
encoded hit positions is translated into a short domain-specific "call" that regenerates it.
Task
A detector stack has 12 layers, each with 32 strips. A charged particle (e.g. a cosmic muon) crosses the stack
roughly in a straight line, hitting strip round(slp * layer + icpt) in every layer it geometrically passes
through (slp = slope in strips/layer, icpt = intercept in strips). Each hit is encoded as a single integer
layer * 32 + strip in [0, 383].
Real data is noisy:
- add — spurious hits not on the track (detector noise).
- rem — genuine track hits that failed to register (detector inefficiency).
Given the (possibly noisy) sorted list of encoded input hits, the model outputs:
gen_evt(slp=[<slope>], icpt=[<intercept>], add=[<noise hit codes>], rem=[<missing hit codes>])
i.e. in one generation pass it must jointly recover the true straight-line track parameters and classify
which input hits are noise and which true-track hits are missing.
Architecture
LlamaForCausalLM — 8 layers, hidden size 512, 8 attention heads, ~31.95M parameters — trained from scratch with
a custom fixed-vocabulary tokenizer (TrackCallTokenizer, 4569 tokens: the 384 hit codes, gen_evt/slp/icpt/
add/rem/syntax tokens, and numeric literals for slp/icpt). custom_tokenizer.py in this repo has the full
implementation.
This checkpoint is best_by_full_20k from epoch 777 of continued training (run
small_transformer_cosmic_aligned_cont1000). Validation score at checkpoint time: perfect=0.8993 good=0.9091
S=0.9857 D=0.0122.
Results
Synthetic test set (19,270 events, known ground truth)
| metric |
all |
clean |
add |
rem |
add+rem |
| parsed % |
100.00 |
100.00 |
100.00 |
100.00 |
100.00 |
| mean similarity S |
0.997 |
1.000 |
0.992 |
1.000 |
0.997 |
| perfect (S=1, D=0) % |
97.44 |
99.98 |
92.85 |
100.00 |
97.50 |
| exact call % |
62.57 |
100.00 |
80.00 |
42.43 |
28.86 |
| model |Δslope| median |
0.000 |
0.000 |
0.000 |
0.010 |
0.030 |
| OLS-fit |Δslope| median |
0.035 |
0.014 |
0.055 |
0.029 |
0.086 |
| model |Δintercept| median |
0.000 |
0.000 |
0.000 |
0.060 |
0.180 |
| OLS-fit |Δintercept| median |
0.196 |
0.081 |
0.305 |
0.162 |
0.478 |
The model recovers exact track parameters (median error = 0) even on noisy events, beating an independent
per-hit ordinary-least-squares fit by roughly an order of magnitude in dispersion (Gaussian-fit σ: 0.049 vs. 0.504
strips/layer for slope, 0.271 vs. 1.871 strips for intercept) and lands closer to the truth than the fit in ~83%
of individual events. The gap is driven by noisy events: the fit weights every input hit equally, so a single
injected noise hit can pull the fitted line far from the true track, while the model implicitly classifies which
hits are noise before "fitting" the rest.
Real cosmic-muon events (100,000 events, no ground truth — compared against two independent fits)
| metric |
value |
| parsed % |
99.98 |
| mean similarity S |
0.982 |
| perfect (S=1, D=0) % |
88.87 |
| |Δslope| vs. reference (per-layer-mean) fit, median |
0.067 |
| |Δintercept| vs. reference fit, median |
0.404 |
Usage
from transformers import LlamaForCausalLM
from custom_tokenizer import TrackCallTokenizer # included in this repo
tokenizer = TrackCallTokenizer.from_pretrained("deepaksamuel-cuk/trackfit-llm-small")
model = LlamaForCausalLM.from_pretrained("deepaksamuel-cuk/trackfit-llm-small")
hits = [29, 30, 122, 123, 153, 154, 184] # encoded hit codes (layer*32 + strip), sorted, deduped
prompt = "<s>[" + ",".join(str(h) for h in hits) + "]<CALL>"
input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
out = model.generate(input_ids, max_new_tokens=56, do_sample=False,
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True))
# gen_evt (slp=[-1.09], icpt=[12.98], add=[...], rem=[...])
TrackCallTokenizer (custom_tokenizer.py, included here) is a from-scratch fixed-vocabulary tokenizer, not a
standard BPE/WordPiece tokenizer — AutoTokenizer will not auto-detect it; import the class directly as above.
Limitations
- Single-track only. Trained on single-particle events; behavior on multi-track events is untested.
- Fixed geometry. The 12-layer × 32-strip layout and the
add/rem noise model are specific to this
detector; the model will not generalize to a different strip/layer count without retraining.
- No physical units.
slp/icpt are in strips/layer and strips, not calibrated to any physical geometry —
converting to a physical angle requires the detector's actual strip pitch and layer spacing.
