deepaksamuel-cuk/trackfit-llm-small

trackfit-llm-small

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)

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)

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.