trackfit-llm-small

trackfit-llm-small

A 37.3M-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.

This is the extended version of the model: its slope/intercept vocabulary (extended-allowed.csv) covers the full range of straight tracks seen in real cosmic-muon data, including steep tracks that cross only part of the detector.

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, defined at layer 0). 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.

Track vocabulary: extended-allowed.csv

The model can only emit (slp, icpt) pairs from a fixed list, extended-allowed.csv (columns s, i, 41,497 rows). Each row is one representative line for a distinct clean hit pattern (0.01 precision), so no two rows produce the same set of hits.

• Slopes: [-8.14, 8.17] strips/layer.
• Intercepts (at layer 0): [-64.92, 95.92] strips.

The earlier version of this model used a list whose intercepts spanned only about [-5, 36]. A steep real track that crosses only the upper layers extrapolates to a layer-0 intercept far outside that range, so the old model could not represent most steep tracks and under-counted large zenith angles. The extended list keeps every row of the old one except the empty-track placeholder (-0.04, -0.76), a line just outside the detector. It adds one row for every other distinct clean hit pattern with |slope| ≤ 8.2 strips/layer and at least 4 hits inside the detector. Every track in the training data is drawn from this file.

Architecture

LlamaForCausalLM: 8 layers, hidden size 512, intermediate size 1536, 8 attention heads, 256 positions (RoPE), 37,278,208 parameters. Trained from scratch with a custom fixed-vocabulary tokenizer (TrackCallTokenizer, 9,772 tokens: the 384 hit codes, gen_evt/slp/icpt/add/rem/syntax tokens, and one numeric literal for every slope/intercept value in extended-allowed.csv). custom_tokenizer.py in this repo has the full implementation.

Training

• Trained from random initialisation, one GPU, batch 256, LR 6e-4 with 500 warm-up steps and cosine decay.
• Weighted cross-entropy: slp/icpt tokens weighted ×5, EOS ×3.
• Every epoch has 184,656 freshly generated single-track events, with tracks drawn uniformly from extended-allowed.csv. The noise classes are mixed clean : add : rem : add+rem = 1 : 1 : 1 : 8.
• This checkpoint is epoch 1740. It was selected as the best checkpoint on real data: the mean similarity S over the x and y views of the first 10,000 events of INORUN_20160718_211418.csv is 0.9555.

Results

Synthetic test set (20,000 events, known ground truth)

A new single-track test set drawn from extended-allowed.csv: 5,000 events per generated class, add length 2–5, rem length 5–8, at least 5 hits left after removal, and no input shared with the training/validation data. The columns below group events by the noise actually present in the truth call.

The model recovers exact track parameters (median error 0) even on noisy events. It beats an independent per-hit ordinary-least-squares fit by roughly a factor of 4–5 in dispersion. The Gaussian-fit σ of the error from truth, clipped to ±1.2 and ±6 strips, is:

• Slope: 0.062 for the model vs 0.323 strips/layer for the fit.
• Intercept: 0.381 vs 1.630 strips.

The model lands closer to the truth than the fit in ~80% of individual events. The fit weights every input hit equally, so a single noise hit can pull the fitted line far from the true track. The model implicitly classifies which hits are noise before "fitting" the rest.

Many exact-call misses are not tracking errors. Hit positions are rounded to integer strips, so different (slp, icpt) pairs with different add/rem lists can regenerate exactly the same hits. 40.9% of the perfect (S=1, D=0) events use such an equivalent but different call.

Real cosmic-muon events (703,694 events, no ground truth)

These are all events of INORUN_20160718_211418.csv. The model is run independently on the x view and the y view, and compared with a RANSAC straight-line fit (residual threshold 1 strip, 100 trials) on the same hits.

The zenith angle combines the two views: θ = arctan(3.0/16.8 · √(m_x² + m_y²)), using a 3 cm strip pitch and a 16.8 cm layer gap. Its distribution agrees with RANSAC:

• Per-event agreement: the median |Δθ| between model and RANSAC is 0.60°, and the median azimuth |Δφ| is 1.96°.
• Compared with a plain OLS fit: OLS gets only 10.5% of events above 30°. A frequently firing noise strip, far from the track, pulls every OLS line toward it and flattens steep tracks.

Result files

Per-event results behind the numbers above are in results/. Load them with pandas.read_csv; hit lists are stored as JSON-style strings.

Cosmic file columns (_x / _y = detector view):

• filename, eventid, x_hits, y_hits, n_hits_*, n_layers_*: the event and its hits.
• predicted_target_*, parsed_*, similarity_*, dissimilarity_*, perfect_*, pred_mult_*, model_slope_*, model_icpt_*, n_add_pred_*, n_rem_pred_*: the model's call and its scores.
• ols_slope_*, ols_icpt_*, ols_slope_err_*, ols_icpt_err_*, ols_rms_*: ordinary least squares over every hit.
• ransac_slope_*, ransac_icpt_*, ransac_n_inliers_*: RANSAC (residual threshold 1 strip, 100 trials).
• zenith_deg_model/ols/ransac, azimuth_deg_model/ols/ransac: angles from each method's x and y slopes.

Test file columns:

• event, n_hits, input_hits, noise_class, true_slope, true_icpt, true_add, true_rem: the event and its truth.
• predicted_target, parsed, similarity, dissimilarity, perfect, good (S > 0.9 and D < 0.1), exact_call, pred_slope, pred_icpt, pred_add, pred_rem, predicted_hits: the model's call and its scores.
• ols_*: ordinary least squares over every hit. lm_*: least squares over the mean strip per layer. ransac_*: RANSAC with the settings above.
• model_dslope, model_dicpt, ols_dslope, ols_dicpt, lm_dslope, lm_dicpt, ransac_dslope, ransac_dicpt: estimate minus truth.

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 = [7, 40, 74, 107, 141, 175, 208, 242, 275, 309, 343, 376]   # 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.58], icpt=[6.75], 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, so import the class directly as above. Single-track completions are at most 47 tokens, so max_new_tokens=56 is enough.

To load the list of allowed track parameters:

from huggingface_hub import hf_hub_download
import pandas as pd

allowed = pd.read_csv(hf_hub_download("deepaksamuel-cuk/trackfit-llm-small", "extended-allowed.csv"))   # columns s, i

Previous version

The earlier 31.9M-parameter model (4,569-token vocabulary, intercepts limited to about [-5, 36]) is still available at revision d02d44e92d3dff100f71eba616f6ce286f3347d7. Load it with from_pretrained("deepaksamuel-cuk/trackfit-llm-small", revision="d02d44e92d3dff100f71eba616f6ce286f3347d7").

Limitations

• Single-track only. Trained on single-particle events; behaviour on multi-track events is untested.
• One view at a time. Each call fits one 2-D strip view. The 3-D direction (zenith/azimuth) comes from combining separate x-view and y-view predictions.
• Fixed geometry. The 12-layer × 32-strip layout and the add/rem noise model are specific to this detector. The model will not generalise to a different strip/layer count without retraining.
• Fixed parameter list. Outputs are restricted to the (slp, icpt) pairs in extended-allowed.csv, i.e. |slope| ≤ ~8.2 strips/layer and tracks with at least 4 hits inside the detector.
• No physical units. slp/icpt are in strips/layer and strips. Converting them to a physical angle requires the detector's strip pitch and layer spacing (3 cm and 16.8 cm for the numbers above).