How Far Can Foundation Models Be Distilled for Efficient Jet Tagging?

Jay Agarwal; Siddharth Khare; Dhruv Kumar

cs.LGOct 21, 2026

How Far Can Foundation Models Be Distilled for Efficient Jet Tagging?

Authors: Jay Agarwal, Siddharth Khare, Dhruv Kumar

Organizations: BITS Pilani, KEK · BITS Pilani

Abstract

Jet tagging is one of the most important applications of machine learning at the Large Hadron Collider. Recent models build in more physics knowledge or rely on large-scale pretraining, but the experiments place strict limits on their inference cost. Knowledge distillation can transfer their performance to small students that are cheap to run. We ask which properties of a teacher help the student, how small a student can get and how close it can come to the state of the art. On a jet tagging benchmark, we distil six teachers that differ in architecture, symmetry and pretraining, including two foundation models, into five small students. Every teacher improves every student, and the gain depends mainly on the strength of the teacher. We then shrink our two strongest distilled students, Transformer and L-GATr-slim, in width and in depth. With 20 thousand parameters, a distilled L-GATr-slim outperforms every published jet tagger trained from scratch at a fraction of the inference cost of every teacher. Even with two thousand parameters, it matches ParT, a model with a thousand times as many parameters. We also study what the students learn from their teachers. Our benchmark of inference cost shows that students of the same size can differ by almost an order of magnitude in speed. We release the code and the weights of all trained models.

Figures & tables

Teacher	Params	AUC	$1/\epsilon_{B}$ ( $\epsilon_{S}{=}0.5$ )	$1/\epsilon_{B}$ ( $\epsilon_{S}{=}0.3$ )
ParT	2.1M	0.98571 $\pm$ 0.00011	394 $\pm$ 14	1,586 $\pm$ 19
MIParT	0.7M	0.98632 $\pm$ 0.00005	467 $\pm$ 5	1,871 $\pm$ 69
L-GATr ∗	1.1M	0.98680 $\pm$ 0.00005	528 $\pm$ 14	2,189 $\pm$ 176
L-GATr-slim ∗	2.0M	0.98711 $\pm$ 0.00008	541 $\pm$ 5	2,321 $\pm$ 46
OmniLearned-S †	2.7M	0.98746 $\pm$ 0.00001	546 $\pm$ 11	2,403 $\pm$ 100
ParT-f.t. †	2.1M	0.98769 $\pm$ 0.00005	644 $\pm$ 16	2,556 $\pm$ 0.2

Table 1: Teacher performance on the top-tagging test set: AUC and background rejection

1/\epsilon_{B}

at two signal efficiencies

\epsilon_{S}

. We mark Lorentz-equivariant models with an asterisk and pretrained models with a dagger, and give means and standard deviations over three runs.

Student	Params	AUC	$1/\epsilon_{B}$ ( $\epsilon_{S}{=}0.5$ )	$1/\epsilon_{B}$ ( $\epsilon_{S}{=}0.3$ )
Transformer	188k	0.98512 $\pm$ 0.00013	393 $\pm$ 5	1,612 $\pm$ 121
LorentzNet ∗	203k	0.98556 $\pm$ 0.00010	413 $\pm$ 16	1,714 $\pm$ 83
ParT	205k	0.98565 $\pm$ 0.00004	384 $\pm$ 7	1,603 $\pm$ 76
MIParT	200k	0.98580 $\pm$ 0.00019	410 $\pm$ 5	1,767 $\pm$ 83
L-GATr-slim ∗	202k	0.98622 $\pm$ 0.00016	425 $\pm$ 20	1,704 $\pm$ 80

Table 2: Students trained without a teacher: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

, with means and standard deviations over three runs.

Figure 1: Background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

(top) and

\epsilon_{S}=0.3

(bottom) of every student, trained without a teacher and distilled from each teacher. Black lines and grey bands show the teachers. We show means and standard deviations over three runs.

	Student
Teacher	Transformer	LorentzNet ∗	ParT	MIParT	L-GATr-slim ∗
No teacher	0.98512 $\pm$ 0.00013	0.98556 $\pm$ 0.00010	0.98565 $\pm$ 0.00004	0.98580 $\pm$ 0.00019	0.98622 $\pm$ 0.00016
ParT	0.98617 $\pm$ 0.00010	0.98630 $\pm$ 0.00010	0.98607 $\pm$ 0.00001	0.98624 $\pm$ 0.00005	0.98643 $\pm$ 0.00005
MIParT	0.98646 $\pm$ 0.00007	0.98655 $\pm$ 0.00010	0.98622 $\pm$ 0.00008	0.98631 $\pm$ 0.00002	0.98660 $\pm$ 0.00001
L-GATr ∗	0.98680 $\pm$ 0.00011	0.98671 $\pm$ 0.00018	0.98659 $\pm$ 0.00007	0.98675 $\pm$ 0.00007	0.98693 $\pm$ 0.00001
L-GATr-slim ∗	0.98691 $\pm$ 0.00006	0.98685 $\pm$ 0.00014	0.98671 $\pm$ 0.00011	0.98690 $\pm$ 0.00011	0.98712 $\pm$ 0.00007

Table 3: AUC of every student (columns), trained without a teacher and distilled from each teacher (rows), with means and standard deviations over three runs.

Figure 2: Gain of every student from distillation against the AUC of its teacher, in AUC (left) and relative background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

(right). Black ticks show the mean over the five students. All values are means over three runs.

Figure 3: Width study: background rejection

1/\epsilon_{B}

\epsilon_{S}=0.3

of ten-block students of decreasing width, trained without a teacher (dotted, hollow) and distilled from three teachers (solid, filled). Darker lines mark stronger teachers, and grey dashed lines show the teachers. Error bars show standard deviations over three runs.

Student	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
		No teacher			ParT-f.t. †
			$1/\epsilon_{B}$			$1/\epsilon_{B}$
Transformer	1.2k	0.97422 $\pm$ 0.00052	115 $\pm$ 4	398 $\pm$ 27	0.97401 $\pm$ 0.00049	112 $\pm$ 5	377 $\pm$ 11
	16k	0.98388 $\pm$ 0.00012	318 $\pm$ 14	1,246 $\pm$ 57	0.98624 $\pm$ 0.00019	488 $\pm$ 8	1,975 $\pm$ 67
	84k	0.98495 $\pm$ 0.00012	373 $\pm$ 8	1,431 $\pm$ 46	0.98688 $\pm$ 0.00009	557 $\pm$ 18	2,347 $\pm$ 57
L-GATr-slim ∗	1.8k	0.98491 $\pm$ 0.00018	358 $\pm$ 5	1,449 $\pm$ 86	0.98567 $\pm$ 0.00004	416 $\pm$ 22	1,872 $\pm$ 153

Table 4: Width study: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of ten-block students of three widths, trained without a teacher and distilled from ParT-f.t., with means and standard deviations over three runs. Table 12 gives the other two teachers.

Figure 4: Depth study: background rejection

1/\epsilon_{B}

\epsilon_{S}=0.3

of students of fixed width with a varying number of blocks, trained without a teacher (dotted, hollow) and distilled from three teachers (solid, filled), one run each. Darker lines mark stronger teachers, and grey dashed lines show the teachers.

Figure 5: What Transformer and L-GATr-slim learn from their teacher, distilled (solid, filled) and trained without a teacher (dotted, hollow). (a) KL divergence from ParT-f.t. to the student against the number of student parameters in the width study. (b) Percentage of the background jets each teacher accepts at

\epsilon_{S}=0.5

that the student also accepts. (c, d) Probability the student assigns to the true class minus that of ParT-f.t., for signal (c) and background (d) jets, in ten groups of equal size ordered by the probability ParT-f.t. assigns to the true class. All values are means over three runs.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teachers
ParT	2.1M	0.98571	394 $\pm$ 14	1,586 $\pm$ 19	197	18,198	1,099	7.9
MIParT	0.7M	0.98632	467 $\pm$ 5	1,871 $\pm$ 69	76.0	7,225	5,132	8.9
L-GATr ∗	1.1M	0.98680	528 $\pm$ 14	2,189 $\pm$ 176	1,868	4,263	707	50.5
L-GATr-slim ∗	2.0M	0.98711	541 $\pm$ 5	2,321 $\pm$ 46	309	16,546	339	8.1

Table 5: Performance and inference cost on real test jets in the fastest 32-bit setting: operations per jet in millions (median over 300 jets, one multiply-accumulate counted as two), GPU throughput and peak memory at a batch of 512 jets and CPU latency for one jet on one core. Teachers are shown with their own performance and students after distillation from ParT-f.t., with standard deviations of the background rejection over three runs. Deep Sets + GNN [ 5 ] is distilled from OmniLearned-L, and its performance is as published. Appendix G gives the same table without a teacher and for every other teacher.

Figure 6: Background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

(left) and

\epsilon_{S}=0.3

(right) against (a) GPU time per jet and (b) GPU peak memory at a batch of 512 jets, (c) CPU latency for one jet and (d) operations per jet, in the fastest 32-bit setting. Students distilled from ParT-f.t. are filled and teachers hollow, and (d) gives the sizes of the Transformer and L-GATr-slim students. Error bars show standard deviations over three runs.

Appendix figures & tables28 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Width	Heads	Blocks	Read-out	Params
Teachers
ParT	128	8	8 + 2	class token	2,141,005
MIParT	64	8	5 + 5 + 2	class token	656,949
L-GATr ∗	16 mv, 32 s	8	12	global token	1,077,474
L-GATr-slim ∗	32 v, 96 s	8	12	global token	2,030,817
OmniLearned-S †	128	8	8 + 2	4 class tokens	2,708,149

Appendix

Table 6: Configuration of every model. Width counts hidden channels, as multivector (mv), vector (v) and scalar (s) channels for the L-GATr family; blocks counts the blocks of each model, with a sum for blocks of different kinds.

Model	Optimizer	Learning rate	Weight decay	Batch	Steps	Schedule
Transformer	AdamW	$10^{-3}$	0.1	256	200,000	cosine
LorentzNet ∗	AdamW	$10^{-3}$	$10^{-2}$	128	336,000	cosine
ParT	Ranger	$10^{-3}$	$10^{-2}$	512	48,000	flat, then decay
MIParT	Ranger	$10^{-3}$	$10^{-2}$	256	72,000	flat, then decay
L-GATr-slim ∗	Lion	$3\times 10^{-5}$	2	128	200,000	cosine
Transformer 84k	AdamW	$10^{-3}$	0.1	256	200,000	cosine

Appendix

Table 7: Training recipes. Steps are optimizer updates.

	This work				Published
			$1/\epsilon_{B}$				$1/\epsilon_{B}$
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
ParT	2.1M	0.98571 $\pm$ 0.00011	394 $\pm$ 14	1,586 $\pm$ 19	2.1M	0.9858	413 $\pm$ 16	1,602 $\pm$ 81
MIParT	0.7M	0.98632 $\pm$ 0.00005	467 $\pm$ 5	1,871 $\pm$ 69	0.7M	0.9868	505 $\pm$ 8	2,010 $\pm$ 97
L-GATr ∗	1.1M	0.98680 $\pm$ 0.00005	528 $\pm$ 14	2,189 $\pm$ 176	1.1M	0.9870 $\pm$ 0.0001	540 $\pm$ 20	2,240 $\pm$ 70
L-GATr-slim ∗	2.0M	0.98711 $\pm$ 0.00008	541 $\pm$ 5	2,321 $\pm$ 46	2.0M	0.9869 $\pm$ 0.0001	546 $\pm$ 7	2,264 $\pm$ 93

Appendix

Table 8: Performance of the teachers and of LorentzNet trained without a teacher, with means and standard deviations over three runs, and the published performance of the same models [ 35 , 42 , 37 , 33 , 6 , 21 ] .

	Single run			Ensemble of three runs
		$1/\epsilon_{B}$			$1/\epsilon_{B}$
Teacher	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
ParT	0.98571 $\pm$ 0.00011	394 $\pm$ 14	1,586 $\pm$ 19	0.98630	448	1,960
MIParT	0.98632 $\pm$ 0.00005	467 $\pm$ 5	1,871 $\pm$ 69	0.98663	499	1,972
L-GATr ∗	0.98680 $\pm$ 0.00005	528 $\pm$ 14	2,189 $\pm$ 176	0.98719	554	2,462
L-GATr-slim ∗	0.98711 $\pm$ 0.00008	541 $\pm$ 5	2,321 $\pm$ 46	0.98748	589	2,490

Appendix

Table 9: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of every teacher as a single run, with the mean and standard deviation over three runs, and as the ensemble of its three runs, which averages their logits.

Teacher	Transformer	LorentzNet ∗	ParT	MIParT	L-GATr-slim ∗
	Student
AUC
L-GATr-slim ∗	0.98691 $\pm$ 0.00006	0.98685 $\pm$ 0.00014	0.98671 $\pm$ 0.00011	0.98690 $\pm$ 0.00011	0.98712 $\pm$ 0.00007
L-GATr-slim ∗ ens.	0.98701 $\pm$ 0.00003	0.98711 $\pm$ 0.00009	0.98678 $\pm$ 0.00010	0.98698 $\pm$ 0.00010	0.98719 $\pm$ 0.00003
OmniLearned-S †	0.98705 $\pm$ 0.00007	0.98713 $\pm$ 0.00018	0.98677 $\pm$ 0.00010	0.98705 $\pm$ 0.00010	0.98727 $\pm$ 0.00002
$1/\epsilon_{B}$ at $\epsilon_{S}=0.5$

Appendix

Table 10: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of every student of Section 4 distilled from L-GATr-slim, from the ensemble of its three runs (L-GATr-slim ens.) and from OmniLearned-S, with means and standard deviations over three runs.

$1/\epsilon_{B}$ at $\epsilon_{S}=0.5$
	Student
Teacher	Transformer	LorentzNet ∗	ParT	MIParT	L-GATr-slim ∗
No teacher	393 $\pm$ 5	413 $\pm$ 16	384 $\pm$ 7	410 $\pm$ 5	425 $\pm$ 20
ParT	441 $\pm$ 20	446 $\pm$ 25	423 $\pm$ 12	441 $\pm$ 12	458 $\pm$ 12
MIParT	482 $\pm$ 9	488 $\pm$ 2	448 $\pm$ 10	474 $\pm$ 11	487 $\pm$ 9
L-GATr ∗	537 $\pm$ 14	527 $\pm$ 9	512 $\pm$ 8	523 $\pm$ 9	533 $\pm$ 24

Appendix

Table 11: Background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

(top) and

\epsilon_{S}=0.3

(bottom) of every student (columns), trained without a teacher and distilled from each teacher (rows). Means and standard deviations over three runs.

Figure 7: Width study: background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

of ten-block students of decreasing width, trained without a teacher (dotted, hollow) and distilled from three teachers (solid, filled, darker for a stronger teacher). Error bars show the standard deviation over three runs, and grey dashed lines show the rejection of the teachers.

Figure 8: Depth study: background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

of students of fixed width with a varying number of blocks, trained without a teacher (dotted, hollow) and distilled from three teachers (solid, filled, darker for a stronger teacher), one run each. Grey dashed lines show the rejection of the teachers.

Student	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
		L-GATr ∗			OmniLearned-S †
			$1/\epsilon_{B}$			$1/\epsilon_{B}$
Transformer	1.2k	0.97391 $\pm$ 0.00031	113 $\pm$ 5	391 $\pm$ 26	0.97360 $\pm$ 0.00038	111 $\pm$ 3	377 $\pm$ 17
	16k	0.98590 $\pm$ 0.00009	457 $\pm$ 9	1,876 $\pm$ 142	0.98615 $\pm$ 0.00021	468 $\pm$ 20	1,976 $\pm$ 120
	84k	0.98658 $\pm$ 0.00011	503 $\pm$ 19	2,230 $\pm$ 133	0.98669 $\pm$ 0.00006	530 $\pm$ 6	2,152 $\pm$ 183
L-GATr-slim ∗	1.8k	0.98560 $\pm$ 0.00012	403 $\pm$ 2	1,773 $\pm$ 67	0.98543 $\pm$ 0.00046	404 $\pm$ 16	1,731 $\pm$ 167

Appendix

Table 12: Width study: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of ten-block students of three widths, distilled from L-GATr and OmniLearned-S. Means and standard deviations over three runs.

			No teacher			ParT-f.t. †
				$1/\epsilon_{B}$			$1/\epsilon_{B}$
Student	Blocks	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
Transformer	0	433	0.86367	10	16	0.86364	10	16
	1	19k	0.98143	222	869	0.98275	257	1,128
	2	38k	0.98360	300	1,224	0.98577	435	1,771
	3	57k	0.98459	346	1,355	0.98648	492	2,103

Appendix

Table 13: Depth study: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of students of fixed width with a varying number of blocks, trained without a teacher and distilled from ParT-f.t., one run each.

			L-GATr ∗			OmniLearned-S †
				$1/\epsilon_{B}$			$1/\epsilon_{B}$
Student	Blocks	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$
Transformer	0	433	0.86357	10	16	0.86364	10	16
	1	19k	0.98295	273	1,091	0.98261	257	1,020
	2	38k	0.98549	408	1,836	0.98557	425	1,711
	3	57k	0.98641	502	2,148	0.98647	506	2,040

Appendix

Table 14: Depth study: AUC and background rejection

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

of students of fixed width with a varying number of blocks, distilled from L-GATr and OmniLearned-S, one run each.

Figure 9: KL divergence from the teacher’s predicted probabilities to the student’s against the number of student parameters, for the three teachers of the scaling studies, in the width study (top, means over three runs) and the depth study (bottom, one run, three at the depth of Section 4 ). Transformer and L-GATr-slim are shown distilled (solid, filled) and trained without a teacher (dotted, hollow).

Figure 10: Of the background jets each teacher accepts at

\epsilon_{S}=0.5

(left) and

\epsilon_{S}=0.3

(right), the percentage every student of Section 4 also accepts at the same signal efficiency, distilled (filled) and trained without a teacher (hollow). Means over three runs.

Figure 11: For each teacher, the percentage of the background jets it accepts and another teacher rejects that the student also accepts at

\epsilon_{S}=0.5

(left) and

\epsilon_{S}=0.3

(right), when distilled from this teacher, distilled from the other teacher or trained without a teacher. Means over the five other teachers, the five students of Section 4 and three runs.

Figure 12: Background jets Transformer and L-GATr-slim accept at

\epsilon_{S}=0.5

among those each teacher rejects (left) and accepts (right), with distillation (solid, filled) and without (dotted, hollow). Means over three runs.

Background jets the teacher rejects
	Transformer		LorentzNet ∗		ParT		MIParT		L-GATr-slim ∗
Teacher	without	with	without	with	without	with	without	with	without	with
ParT	185	87	178	76	179	105	163	87	145	71
MIParT	201	74	183	69	186	91	150	60	156	61
L-GATr ∗	248	77	222	85	243	105	217	87	191	73
L-GATr-slim ∗	253	88	224	91	255	111	229	95	200	74

Appendix

Table 15: Background jets every student of Section 4 accepts at

\epsilon_{S}=0.5

, without and with distillation, among the jets each teacher rejects (top) and accepts (bottom). Means over three runs.

Figure 13: Probability every student of Section 4 assigns to the true class minus that assigned by ParT-f.t., for signal (left) and background (right) jets in ten groups of equal size and increasing ParT-f.t. confidence, distilled (solid, filled) and trained without a teacher (dotted, hollow). Means over three runs.

Figure 14: Probability every student of Section 4 assigns to the true class minus that assigned by OmniLearned-S, for signal (left) and background (right) jets in ten groups of equal size and increasing OmniLearned-S confidence, distilled (solid, filled) and trained without a teacher (dotted, hollow). Means over three runs.

Figure 15: Probability every student of Section 4 assigns to the true class minus that assigned by L-GATr, for signal (left) and background (right) jets in ten groups of equal size and increasing L-GATr confidence, distilled (solid, filled) and trained without a teacher (dotted, hollow). Means over three runs.

Figure 16: Operations per jet, GPU throughput and peak memory at a batch of 512 jets and CPU latency for one jet on one core against the number of parameters, for every model of Table 5 , on real test jets in the fastest 32-bit setting.

Model	Particle lists	Attention	GPU	CPU
			Fastest setting
ParT	padded	built-in	padded, compiled, built-in	padded, compiled, built-in
MIParT	padded	built-in	padded, compiled, built-in	padded, compiled, built-in
L-GATr ∗	packed, padded	xFormers, Flex	packed, compiled, xFormers	padded, compiled, xFormers
L-GATr-slim ∗	packed, padded	xFormers, Flex	packed, compiled, xFormers	packed, compiled, xFormers
OmniLearned-S †	padded	built-in	padded, compiled, built-in	padded, compiled, built-in

Appendix

Table 16: Execution settings timed for every model: particle lists and attention kernels (Flex: FlexAttention), each eager and compiled, and the fastest setting on each device, which Tables 5 and 17 report.

		GPU		CPU
Model	Params	jets/s	MB	ms/jet
Teachers
ParT	2.1M	38,503	388	7.3
MIParT	0.7M	28,320	1,007	8.5
L-GATr ∗	1.1M	5,641	785	47.6
L-GATr-slim ∗	2.0M	17,485	351	7.6

Appendix

Table 17: Inference cost on jets with exactly 50 particles and random momenta: GPU throughput and peak memory at a batch of 512 jets and CPU latency for one jet on one core, each in the fastest 32-bit setting.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Students
Transformer	188k	0.98512	393 $\pm$ 5	1,612 $\pm$ 121	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98556	413 $\pm$ 16	1,714 $\pm$ 83	537	7,601	1,722	11.0
ParT	205k	0.98565	384 $\pm$ 7	1,603 $\pm$ 76	21.9	26,228	2,166	2.3
MIParT	200k	0.98580	410 $\pm$ 5	1,767 $\pm$ 83	20.4	12,181	1,709	7.8

Appendix

Table 18: Performance and inference cost of the students trained without a teacher.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teacher
ParT	2.1M	0.98571	394 $\pm$ 14	1,586 $\pm$ 19	197	18,198	1,099	7.9
Students
Transformer	188k	0.98617	441 $\pm$ 20	1,826 $\pm$ 143	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98630	446 $\pm$ 25	1,897 $\pm$ 96	537	7,601	1,722	11.0

Appendix

Table 19: Performance and inference cost of the students distilled from ParT and of ParT itself.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teacher
MIParT	0.7M	0.98632	467 $\pm$ 5	1,871 $\pm$ 69	76.0	7,225	5,132	8.9
Students
Transformer	188k	0.98646	482 $\pm$ 9	1,834 $\pm$ 95	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98655	488 $\pm$ 2	1,758 $\pm$ 71	537	7,601	1,722	11.0

Appendix

Table 20: Performance and inference cost of the students distilled from MIParT and of MIParT itself.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teacher
L-GATr ∗	1.1M	0.98680	528 $\pm$ 14	2,189 $\pm$ 176	1,868	4,263	707	50.5
Students
Transformer	188k	0.98680	537 $\pm$ 14	2,294 $\pm$ 166	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98671	527 $\pm$ 9	2,260 $\pm$ 37	537	7,601	1,722	11.0

Appendix

Table 21: Performance and inference cost of the students distilled from L-GATr and of L-GATr itself.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teacher
L-GATr-slim ∗	2.0M	0.98711	541 $\pm$ 5	2,321 $\pm$ 46	309	16,546	339	8.1
Students
Transformer	188k	0.98691	546 $\pm$ 10	2,374 $\pm$ 103	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98685	539 $\pm$ 6	2,261 $\pm$ 38	537	7,601	1,722	11.0

Appendix

Table 22: Performance and inference cost of the students distilled from L-GATr-slim and of L-GATr-slim itself.

			$1/\epsilon_{B}$			GPU		CPU
Model	Params	AUC	$\epsilon_{S}{=}0.5$	$\epsilon_{S}{=}0.3$	MFLOP	jets/s	MB	ms/jet
Teacher
OmniLearned-S †	2.7M	0.98746	546 $\pm$ 11	2,403 $\pm$ 100	427	3,736	13,221	17.0
Students
Transformer	188k	0.98705	565 $\pm$ 6	2,325 $\pm$ 121	21.6	70,175	87	2.7
LorentzNet ∗	203k	0.98713	536 $\pm$ 12	2,483 $\pm$ 120	537	7,601	1,722	11.0

Appendix

Table 23: Performance and inference cost of the students distilled from OmniLearned-S and of OmniLearned-S itself.

Jet Tagging

Jet tagging uses particle-collision data to classify jets by their underlying particles or processes, supporting precision measurements and searches for new physics.

See all 18 papers

Oct 1, 2026

Distillation of Tabular Foundation Models into Efficient Predictors

Tabular foundation models (TFMs) achieve strong predictive performance through in-context learning, yet repeatedly conditioning on labeled data makes inference expensive. Knowledge distillation can reduce this cost by transferring their predictive ability to lightweight, dataset-specific students. However, the dependence of TFM predictions on both a labeled context and a query introduces two design questions: how to construct teacher supervision and whether expanding query coverage improves distillation. We examine these questions across two TFMs and both neural and tree-based students, and derive an effective distillation recipe. The recipe uses the full labeled training set as teacher context and trains students solely on teacher predictions for observed and synthetic queries. On TabArena, the resulting students outperform their supervised trained tuned-and-ensembled counterparts by 57-98 Elo points. Applied unchanged to TALENT, the same recipe improves matched default students on 236-258 of 300 datasets and reduces median primary error by 4.0-6.4%. The distilled students also achieve median inference speedups of 3.0-21.6 times over their teachers, offering a practical trade-off between predictive performance and repeated inference cost. Code is available at https://github.com/nums-ai/TFM_Distillation .

Minho Jeong, Dooho Lee, Jinmo Lee +1

Nums AI

Tabular Foundation Models Efficient Neural Network Inference Tabular ICL Knowledge Distillation

Ask PDF

Apr 4, 2026

Training a Student Expert via Semi-Supervised Foundation Model Distillation

Foundation models deliver strong perception but are often too computationally heavy to deploy, and adapting them typically requires costly annotations. We introduce a semi-supervised knowledge distillation (SSKD) framework that compresses pre-trained vision foundation models (VFMs) into compact experts using limited labeled and abundant unlabeled data, and instantiate it for instance segmentation where per-pixel labels are particularly expensive. The framework unfolds in three stages: (1) domain adaptation of the VFM(s) via self-training with contrastive calibration, (2) knowledge transfer through a unified multi-objective loss, and (3) student refinement to mitigate residual pseudo-label bias. Central to our approach is an instance-aware pixel-wise contrastive loss that fuses mask and class scores to extract informative negatives and enforce clear inter-instance margins. By maintaining this contrastive signal across both adaptation and distillation, we align teacher and student embeddings and more effectively leverage unlabeled images. On Cityscapes and ADE20K, our

\approx 11\times

smaller student improves over its zero-shot VFM teacher(s) by +11.9 and +8.6 AP, surpasses adapted teacher(s) by +3.4 and +1.5 AP, and outperforms state-of-the-art SSKD methods on benchmarks.

Pardis Taghavi, Tian Liu, Renjie Li +2

Texas A&M University

Knowledge Distillation Instance Segmentation Vision Foundation Model Adaptation

Ask PDF

Aug 4, 2026

Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-

K

logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29% faster per iteration, and reaching up to 41% higher throughput on a single H200 GPU. Second, we introduce a fused, chunked KL loss that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.

Bakbergen Ryskulov, Iker García-Ferrero, David Montero +5

Multiverse Computing

Online Knowledge Distillation Efficient Language Model Training Language Model Distillation Knowledge Distillation

Ask PDF

Abstract

Jet Tagging

Explore similar work

Distillation of Tabular Foundation Models into Efficient Predictors

Training a Student Expert via Semi-Supervised Foundation Model Distillation

Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss