trackfit-llm-large

A 174M-parameter Llama-style transformer, trained from scratch, that reconstructs straight-line particle tracks in a 12-layer strip detector. It reads the list of hit strips in one detector view and writes a short "program" that regenerates the event: the slope and intercept of each track (up to 3 tracks), which input hits are noise, and which hits the tracks should have produced but are missing.

input : [21,54,86,119,151,184,217,249,282,314,347,380]
output: gen_evt(slp=[0.6], icpt=[21.01], add=[], rem=[])

Detector and encoding

• 12 layers (0–11), 32 strips per layer (0–31), one view (x or y) at a time.
• Each hit is one integer: hit = layer * 32 + strip, so hits are in [0, 383]. Decode with strip = hit % 32, layer = hit // 32.
• A track is the line strip = slp * layer + icpt (slope in strips per layer, intercept = strip at layer 0). It produces a hit at round(slp * layer + icpt) in every layer where that value is inside 0–31.

The output call gen_evt(slp=[...], icpt=[...], add=[...], rem=[...]) means:

Regenerating the hits from the call (tracks − rem + add) and comparing with the input gives two scores: similarity S = |input ∩ regenerated| / |input| and dissimilarity D = |regenerated − input| / |regenerated|. S = 1, D = 0 means the call reproduces the event exactly.

Slopes and intercepts are single vocabulary tokens taken from a fixed grid (allowed.csv, 41,497 (slope, intercept) pairs, |slope| ≤ 8.17 strips/layer), so the model can only output values on that grid.

Files

Quick start

pip install torch "transformers>=4.45" huggingface_hub

import json, sys, torch
from huggingface_hub import snapshot_download
from transformers import LlamaForCausalLM

path = snapshot_download("deepaksamuel-cuk/trackfit-llm-large")
sys.path.insert(0, path)                      # makes trackfit.py importable
from trackfit import TrackCallTokenizer, predict, score, parse_call

device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = TrackCallTokenizer.from_pretrained(path)
model = LlamaForCausalLM.from_pretrained(path, dtype=torch.bfloat16 if device == "cuda" else torch.float32).to(device).eval()

hits = [21, 54, 86, 119, 151, 184, 217, 249, 282, 314, 347, 380]   # one event, one view
call = predict(model, tokenizer, [hits])[0]
print(call)                   # gen_evt (slp=[0.6], icpt=[21.01], add=[], rem=[])
print(parse_call(call))       # {'slp': [0.6], 'icpt': [21.01], 'add': [], 'rem': []}
print(score(hits, call))      # (1.0, 0.0)  -> S, D

(With transformers < 5, pass torch_dtype= instead of dtype=.) The model runs on CPU as well; a GPU is only needed for large batches.

Run the examples

examples = [json.loads(line) for line in open(f"{path}/examples.jsonl")]
calls = predict(model, tokenizer, [e["input_hits"] for e in examples])   # batched greedy decoding
for e, call in zip(examples, calls):
    S, D = score(e["input_hits"], call)
    print(f"{e['name']}  ({e['description']})")
    print(f"  input : {e['input_hits']}")
    print(f"  model : {call}   S={S:.3f} D={D:.3f}")
    if "target_call" in e:
        print(f"  truth : {e['target_call']}")

Expected output (identical on GPU in bfloat16 and on CPU in float32):

On noisy events the predicted slope/intercept can differ slightly from the true values (for example icpt=9.64 instead of 9.89) while still reproducing every hit exactly (S = 1, D = 0): several nearby lines on the grid give the same strips in every layer, and the add/rem lists adjust accordingly.

Using your own data

1. Split each event into its x-view and y-view hits and encode each as layer * 32 + strip (0–383). Run the two views separately.
2. Pass a list of events (each a list of ints) to predict; it sorts the hits, left-pads the batch and decodes greedily. Batches of a few hundred to a few thousand events work well on a GPU.
3. Parse the calls with parse_call. For a single track per view, the slopes m_x, m_y (strips/layer) give the zenith angle θ = arctan((pitch / layer_gap) · sqrt(m_x² + m_y²)) and azimuth φ = atan2(m_y, m_x); for the detector used here, pitch = 3.0 cm and layer gap = 16.8 cm.

Plot one event:

import matplotlib.pyplot as plt, numpy as np
from trackfit import decode

p = parse_call(call)
strips, layers = zip(*(decode(h) for h in hits))
plt.scatter(strips, layers, color="k", label="hits")
lay = np.array([-0.5, 11.5])
for s, i in zip(p["slp"], p["icpt"]):
    plt.plot(s * lay + i, lay, label=f"slp={s}, icpt={i}")
plt.xlim(-0.5, 31.5); plt.ylim(-0.5, 11.5); plt.xlabel("strip"); plt.ylabel("layer"); plt.legend(); plt.show()

Training

• Architecture: LlamaForCausalLM, hidden size 1024, 12 layers, 16 attention heads, intermediate size 2816, 174,179,328 parameters, trained from scratch.
• Data: synthetic events generated on the fly (120,000 new events per epoch), one third each 1-, 2- and 3-track, with tracks drawn uniformly from allowed.csv. Per track: noise class clean : add : rem : add+rem = 1 : 1 : 1 : 8; 1–5 noise hits and up to 7 missing hits per track, at least 5 hits kept on every track. Maximum per event: 15 noise hits and 21 missing hits (3 tracks).
• Objective: next-token prediction on the call, weighted cross-entropy, warmup–stable–decay learning-rate schedule.
• Checkpoint: epoch 320, selected by the best mean similarity on 10,000 real cosmic-muon events (x and y views).

Results

Synthetic test set (20,000 events, never seen in training):

Real cosmic-muon data (703,694 events from one run, x and y views):

The zenith angle from the model's x/y slopes agrees with an independent RANSAC fit to a median of 0.51°.

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_* (number of predicted tracks), model_slopes_*, model_icpts_* (all tracks), model_slope_*, model_icpt_* (first track), n_add_pred_*, n_rem_pred_*, notrack_* (a predicted track has no hit inside the detector): the model's call and its scores.
• ols_slope_*, ols_icpt_*: 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 (the model's first track in each view).

Test file columns:

• event, n_tracks, n_hits, input_hits, target_call, true_slopes, true_icpts: the event and its truth.
• predicted_target, parsed, similarity, dissimilarity, perfect, pred_mult, mult_ok (correct number of tracks), exact (all slopes and intercepts equal the truth), model_slopes, model_icpts, model_slope, model_icpt, n_add_pred, n_rem_pred, notrack: 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 with the settings above.

OLS and RANSAC fit a single straight line to all hits of an event. On 2- and 3-track test events they therefore describe at most one of the tracks; compare them with the model on the n_tracks == 1 rows.

Limitations

• Slope limit. Slopes are capped at |slope| = 8.17 strips/layer (θ ≈ 55.6° from one view). Real tracks steeper than this are clipped to the limit, and on some short, sparse events the model falls back to a near-limit slope through one or two hits and labels the rest as noise. Both pile up at θ ≈ 52–58°. Such events have low S / high D in that view and can be flagged.
• Split single muons. On real data a muon that fires two neighbouring strips per layer is sometimes predicted as two nearly parallel tracks. About 3–5% of real events get a 2-track prediction in one view, but only a few tens show two clearly separated tracks in both views.
• Grid values. Slopes and intercepts are restricted to the allowed.csv grid; the model is a pattern-recognition step, not a precision fit. A least-squares refit on the hits the model assigns to each track gives continuous parameters.
• Geometry. Trained for a 12-layer, 32-strip view only; other geometries need retraining.