diff --git a/QTracker_training/.gitignore b/QTracker_training/.gitignore
index 2586f11..75252a4 100644
--- a/QTracker_training/.gitignore
+++ b/QTracker_training/.gitignore
@@ -3,7 +3,12 @@
 checkpoints/
 Slurm_Files/
 local/
+docs/
 
 run.sh
 
-__pycache__/
\ No newline at end of file
+__pycache__/
+
+.cursor/settings.json
+.claude/settings.local.json
+CLAUDE.local.md
\ No newline at end of file
diff --git a/QTracker_training/CLAUDE.md b/QTracker_training/CLAUDE.md
index 10105b3..d1ae070 100644
--- a/QTracker_training/CLAUDE.md
+++ b/QTracker_training/CLAUDE.md
@@ -264,3 +264,5 @@ Pre-commit hook configured with ruff (v0.14.4):
 6. **File size considerations**: Training datasets can be very large (500K-19M events). Use `data/skim.py` or `data/skim_flat.py` to create manageable subsets for development.
 
 7. **Mixed precision**: Enabled by default in TrackFinder for performance. Final output layers cast to FP32 for numerical stability.
+
+8. **Commit and PR Rules**: Do NOT mention Claude Code. Preface commit msg with tag like `feat` or `fix`.
diff --git a/QTracker_training/data/multi_track/gen_training_random.py b/QTracker_training/data/multi_track/gen_training_random.py
index dca8b9f..a7afd55 100644
--- a/QTracker_training/data/multi_track/gen_training_random.py
+++ b/QTracker_training/data/multi_track/gen_training_random.py
@@ -23,9 +23,14 @@ def _get_value(x):
 
 
 def combine_files(
-    file1, file2, output_file, pairsmup, pairsmum, use_random, at_least_one_pair
+    file1,
+    file2,
+    output_file,
+    pairsmup,
+    pairsmum,
+    use_random,
+    at_least_one_pair,
 ):
-
     if os.path.exists(output_file):
         os.remove(output_file)
 
@@ -108,7 +113,6 @@ def combine_files(
     ev = 0
 
     while idx_mup < n1 and idx_mum < n2:
-
         if use_random:
             lower_bound = 1 if at_least_one_pair else 0
             max_possible_pairs = min(pairsmup, pairsmum)
@@ -156,7 +160,6 @@ def combine_files(
 
         # -------- Fill paired tracks --------
         for k in range(current_pairs):
-
             tree1.GetEntry(idx_mup)
             tree2.GetEntry(idx_mum)
 
@@ -203,6 +206,15 @@ def combine_files(
             idx_mup += 1
             idx_mum += 1
 
+        # --- Canonical ordering: sort filled pair slots by mu+ element ID sum ---
+        if current_pairs > 1:
+            pair_indices = list(range(current_pairs))
+            pair_indices.sort(
+                key=lambda k: int(np.sum(hitarray_mup[k][hitarray_mup[k] > 0]))
+            )
+            hitarray_mup[:current_pairs] = hitarray_mup[pair_indices]
+            hitarray_mum[:current_pairs] = hitarray_mum[pair_indices]
+
         # -------- Extra mu+ --------
         for _ in range(extra_mup):
             tree1.GetEntry(idx_mup)
@@ -364,6 +376,18 @@ def add_hit_array(input_file, output_file):
         action="store_true",
         help="Ensure at least one pair of mu+ and mu- in each event",
     )
+    parser.add_argument(
+        "--output_mom1",
+        type=str,
+        default="momentum_training-1.root",
+        help="Output path for mu+ momentum training file",
+    )
+    parser.add_argument(
+        "--output_mom2",
+        type=str,
+        default="momentum_training-2.root",
+        help="Output path for mu- momentum training file",
+    )
 
     args = parser.parse_args()
 
@@ -383,8 +407,8 @@ def add_hit_array(input_file, output_file):
     if pairsmup <= 0 or pairsmum <= 0:
         raise ValueError("pairsmup and pairsmum must be >= 1")
 
-    file1_array_output = "momentum_training-1.root"
-    file2_array_output = "momentum_training-2.root"
+    file1_array_output = args.output_mom1
+    file2_array_output = args.output_mom2
     for file_name in [args.output, file1_array_output, file2_array_output]:
         if os.path.exists(file_name):
             os.remove(file_name)
diff --git a/QTracker_training/eval_multi_track.py b/QTracker_training/eval_multi_track.py
index 381a736..07ab71d 100644
--- a/QTracker_training/eval_multi_track.py
+++ b/QTracker_training/eval_multi_track.py
@@ -7,6 +7,7 @@
 absl.logging.set_verbosity("error")
 
 import argparse
+import itertools
 import ROOT  # noqa: F401
 import numpy as np
 import tensorflow as tf
@@ -56,6 +57,62 @@ def chi_squared(y_true, y_pred):
     return chi2_mean
 
 
+def match_predictions(y_test, y_pred_argmax, mask):
+    """Reorder predicted pair slots per event to best match the active GT pairs.
+
+    For each event we restrict to the active GT pairs (slots with any nonzero
+    hit; GT is assumed already canonically sorted). We build a residual-based
+    cost between every predicted slot ``p`` and every active GT slot ``g``:
+
+        cost(p, g) = sum_{unmasked det} |pred_p - gt_g| for mu+
+                   + sum_{unmasked det} |pred_p - gt_g| for mu-
+
+    We enumerate permutations of predicted slots and pick the one minimizing
+    the total cost over the active GT slots. The chosen permutation is applied
+    so the predicted slot matched to GT rank ``k`` lands at index ``k``.
+    GT and model pair-slot counts may differ (e.g. 3-pair GT vs 5-pair model).
+    """
+    matched = y_pred_argmax.copy()
+    det_idx = np.where(mask)[0]
+
+    for ev in range(len(y_test)):
+        gt = y_test[ev].astype(np.int32)  # (n_gt, 2, 62)
+        pred = y_pred_argmax[ev]  # (n_pred, 2, 62)
+        n_gt = gt.shape[0]
+        n_pred = pred.shape[0]
+
+        n_active = max(
+            (k + 1 for k in range(n_gt) if np.any(gt[k] != 0)),
+            default=0,
+        )
+        if n_active <= 1:
+            continue
+
+        n_match = min(n_active, n_pred)
+
+        cost = np.zeros((n_pred, n_match))
+        for p in range(n_pred):
+            for g in range(n_match):
+                c_plus = np.sum(np.abs(pred[p, 0, det_idx] - gt[g, 0, det_idx]))
+                c_minus = np.sum(np.abs(pred[p, 1, det_idx] - gt[g, 1, det_idx]))
+                cost[p, g] = c_plus + c_minus
+
+        best_perm = None
+        best_cost = None
+        for perm in itertools.permutations(range(n_pred), n_match):
+            total = sum(cost[perm[k], k] for k in range(n_match))
+            if best_cost is None or total < best_cost:
+                best_cost = total
+                best_perm = perm
+
+        reordered = pred.copy()
+        for k in range(n_match):
+            reordered[k] = pred[best_perm[k]]
+        matched[ev] = reordered
+
+    return matched
+
+
 def evaluate_model(args):
     # Load data - existing loader handles both formats
     load_result = data_loader.load_data(
@@ -76,6 +133,25 @@ def evaluate_model(args):
         print(f"{'=' * 70}\n")
         y_test = np.stack([y_muPlus_test, y_muMinus_test], axis=2)
         # Shape: (num_events, max_pairs, 2, 62)
+
+        # Apply canonical ordering to GT pairs so evaluation matches training ordering
+        for ev_idx in range(len(y_test)):
+            n_active = max(
+                (
+                    k + 1
+                    for k in range(y_test.shape[1])
+                    if np.any(y_test[ev_idx, k] != 0)
+                ),
+                default=0,
+            )
+            if n_active > 1:
+                active_indices = np.arange(n_active)
+                sort_keys = [
+                    np.sum(y_test[ev_idx, k, 0, :][y_test[ev_idx, k, 0, :] > 0])
+                    for k in active_indices
+                ]
+                sorted_order = active_indices[np.argsort(sort_keys)]
+                y_test[ev_idx, :n_active] = y_test[ev_idx, sorted_order]
     else:
         print("\n" + "=" * 70)
         print("Single-track format detected")
@@ -118,60 +194,28 @@ def evaluate_model(args):
 
     # Extract argmax predictions
     y_pred_argmax = np.argmax(y_pred, axis=-1).astype(np.int32)
-    # Shape: (num_events, max_pairs, 2, 62)
+    # Shape: (num_events, pred_max_pairs, 2, 62)
+
+    gt_max_pairs = y_test.shape[1]
+    pred_max_pairs = y_pred_argmax.shape[1]
+    if gt_max_pairs != pred_max_pairs:
+        print(
+            f"WARNING: GT has {gt_max_pairs} pair slots but model outputs "
+            f"{pred_max_pairs}. Evaluating the first "
+            f"{min(gt_max_pairs, pred_max_pairs)} pair(s)."
+        )
+    eval_max_pairs = min(gt_max_pairs, pred_max_pairs)
+
+    # Reorder predicted pair slots per event to best match the active GT pairs
+    y_pred_argmax = match_predictions(y_test, y_pred_argmax, mask)
 
     # Evaluate each pair
-    for pair_idx in range(max_pairs):
+    for pair_idx in range(eval_max_pairs):
         print(f"\n{'=' * 70}")
         print(f"Evaluating Pair {pair_idx}")
         print(f"{'=' * 70}")
 
-        # ============================================================
-        # Pair Existence Evaluation (captures FP and FN)
-        # ============================================================
-
-        print("\n--- Pair Existence Metrics ---")
-
-        # Ground truth existence
-        gt_exists = np.any(y_test[:, pair_idx, :, :] != 0, axis=(1, 2))
-
-        # Prediction existence (after argmax)
-        pred_exists = np.any(y_pred_argmax[:, pair_idx, :, :] != 0, axis=(1, 2))
-
-        TP = np.sum(gt_exists & pred_exists)
-        TN = np.sum(~gt_exists & ~pred_exists)
-        FP = np.sum(~gt_exists & pred_exists)
-        FN = np.sum(gt_exists & ~pred_exists)
-
-        total = len(gt_exists)
-
-        accuracy_exist = (TP + TN) / total if total > 0 else 0.0
-        precision = TP / (TP + FP) if (TP + FP) > 0 else 0.0
-        recall = TP / (TP + FN) if (TP + FN) > 0 else 0.0
-        specificity = TN / (TN + FP) if (TN + FP) > 0 else 0.0
-        f1 = (
-            2 * precision * recall / (precision + recall)
-            if (precision + recall) > 0
-            else 0.0
-        )
-
-        print(f"Total events: {total}")
-        print(f"True Positives : {TP}")
-        print(f"True Negatives : {TN}")
-        print(f"False Positives: {FP}")
-        print(f"False Negatives: {FN}")
-
-        print(f"\nExistence Accuracy : {accuracy_exist:.4f}")
-        print(f"Precision          : {precision:.4f}")
-        print(f"Recall             : {recall:.4f}")
-        print(f"Specificity        : {specificity:.4f}")
-        print(f"F1 Score           : {f1:.4f}")
-
-        if np.sum(~gt_exists) > 0:
-            fp_rate_empty = FP / np.sum(~gt_exists)
-            print(f"\nFalse Positive Rate on Empty Pairs: {fp_rate_empty:.4f}")
-
-        # Check for non-zero ground truth to determine valid events
+        # Restrict to events where this GT rank is active (any nonzero hit)
         valid_mask = np.any(y_test[:, pair_idx, :, :] != 0, axis=(1, 2))
 
         num_valid = np.sum(valid_mask)
@@ -181,7 +225,7 @@ def evaluate_model(args):
 
         print(f"Valid events: {num_valid}/{len(y_test)}")
 
-        # Extract predictions and ground truth for this pair
+        # Extract matched predictions and ground truth for this pair
         y_p_raw = y_pred_argmax[valid_mask, pair_idx, 0, :]  # (valid_events, 62)
         y_m_raw = y_pred_argmax[valid_mask, pair_idx, 1, :]
 
diff --git a/QTracker_training/models/MultiTrackFinder.py b/QTracker_training/models/MultiTrackFinder.py
index 693b93b..d48d54d 100644
--- a/QTracker_training/models/MultiTrackFinder.py
+++ b/QTracker_training/models/MultiTrackFinder.py
@@ -19,15 +19,12 @@
 
 from backbones import unetpp_backbone
 from data_loader import load_data_denoise
-from losses import multi_track_loss, weighted_bce
+from losses import min_perm_multi_track_loss, weighted_bce
 
 # Set seeds
 tf.random.set_seed(42)
 np.random.seed(42)
 
-# Ensure the checkpoints directory exists
-os.makedirs("checkpoints", exist_ok=True)
-
 # Set mixed precision policy for better performance
 mixed_precision.set_global_policy("mixed_float16")
 
@@ -185,9 +182,12 @@ def train_model(args: argparse.Namespace) -> None:
             optimizer=optimizer,
             loss={
                 "denoise": weighted_bce(pos_weight=args.pos_weight),
-                "segment": multi_track_loss(
+                "segment": min_perm_multi_track_loss(
+                    max_pairs=args.max_pairs,
                     lambda_presence=args.lambda_presence,
                     pos_weight_presence=args.pos_weight_presence,
+                    focal_gamma=args.focal_gamma,
+                    lambda_diversity=args.lambda_diversity,
                 ),
             },
             loss_weights={
@@ -336,6 +336,9 @@ def train_model(args: argparse.Namespace) -> None:
             verbose=2,
         )
 
+    output_dir = os.path.dirname(args.output_model)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
     model.save(args.output_model)
     print(f"Model saved to {args.output_model}")
 
@@ -500,7 +503,7 @@ def train_model(args: argparse.Namespace) -> None:
     parser.add_argument(
         "--lambda_presence",
         type=float,
-        default=0.2,
+        default=1.0,
         help="Weight for presence term in multi-track loss.",
     )
     parser.add_argument(
@@ -509,6 +512,18 @@ def train_model(args: argparse.Namespace) -> None:
         default=5.0,
         help="Positive class weight for presence term in multi-track loss.",
     )
+    parser.add_argument(
+        "--focal_gamma",
+        type=float,
+        default=2.0,
+        help="Gamma for focal loss in presence term (0 = standard BCE).",
+    )
+    parser.add_argument(
+        "--lambda_diversity",
+        type=float,
+        default=0.05,
+        help="Weight for inter-pair diversity penalty.",
+    )
     args = parser.parse_args()
 
     train_model(args)
diff --git a/QTracker_training/models/losses.py b/QTracker_training/models/losses.py
index eab9b05..ef8ca90 100644
--- a/QTracker_training/models/losses.py
+++ b/QTracker_training/models/losses.py
@@ -1,3 +1,5 @@
+import itertools
+
 import tensorflow as tf
 from typing import Callable
 
@@ -5,6 +7,20 @@
 DISTANCE_LAMBDA = 5e-4
 EPSILON = 1e-7
 
+MAX_PERMS_P = (
+    5  # P (max_pairs) should not exceed this; beyond that, P! becomes impractical.
+)
+
+
+def _build_perm_table(p: int) -> tf.Tensor:
+    """Build permutation index table for a given number of slots."""
+    if p > MAX_PERMS_P:
+        raise ValueError(
+            f"max_pairs={p} exceeds limit of {MAX_PERMS_P} "
+            f"(would require {p}! = {len(list(itertools.permutations(range(p))))} permutations)"
+        )
+    return tf.constant(list(itertools.permutations(range(p))), dtype=tf.int32)
+
 
 def custom_loss(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
     """
@@ -41,6 +57,110 @@ def custom_loss(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
     return tf.reduce_mean(loss_mup + loss_mum + OVERLAP_LAMBDA * overlap_penalty)
 
 
+def min_perm_multi_track_loss(
+    max_pairs: int = 5,
+    lambda_presence: float = 1.0,
+    pos_weight_presence: float = 5.0,
+    focal_gamma: float = 2.0,
+    lambda_diversity: float = 0.05,
+) -> Callable:
+    """
+    Multi-track segmentation loss with min-over-permutations matching,
+    focal presence, and diversity penalty.
+
+    Uses brute-force enumeration of all P! permutations to find the optimal
+    assignment between predicted slots and GT slots. For P<=5 (at most 120
+    perms), this is negligible compute vs the model forward pass and avoids
+    the tf.py_function required by Hungarian matching (which is brittle under
+    MirroredStrategy multi-GPU training).
+
+    Components:
+      1) Min-over-permutations sparse categorical CE on nonzero targets.
+      2) Focal BCE for hit presence (inside the permutation min).
+      3) Inter-pair diversity penalty (outside the min — depends only on y_pred).
+
+    Args:
+        max_pairs: number of output pair slots (must not exceed MAX_PERMS_P=5)
+        lambda_presence: weight for presence/focal term
+        pos_weight_presence: weight multiplier for positive presence labels
+        focal_gamma: gamma parameter for focal loss (0 = standard BCE)
+        lambda_diversity: weight for inter-pair diversity penalty
+
+    Returns:
+        Loss function with signature (y_true, y_pred) -> scalar.
+    """
+    all_perms = _build_perm_table(max_pairs)  # (P!, P)
+
+    def loss(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        """
+        y_true: (B, P, 2, 62)     — sparse element IDs (0 = no hit)
+        y_pred: (B, P, 2, 62, C)  — softmax probabilities
+        """
+        y_true = tf.cast(y_true, tf.int32)
+        y_pred = tf.cast(y_pred, tf.float32)
+
+        # --- Step 1: Evaluate loss for all P! permutations ---
+        # Gather produces (B, P!, P, 2, 62); transpose to (P!, B, P, 2, 62)
+        y_true_all = tf.transpose(
+            tf.gather(y_true, all_perms, axis=1), perm=[1, 0, 2, 3, 4]
+        )
+
+        # Tile predictions across all P! permutations: (P!, B, P, 2, 62, C)
+        num_perms = tf.shape(y_true_all)[0]
+        y_pred_exp = tf.tile(tf.expand_dims(y_pred, 0), [num_perms, 1, 1, 1, 1, 1])
+
+        # Hit mask for all perms
+        mask_all = tf.cast(tf.not_equal(y_true_all, 0), tf.float32)
+
+        # Masked sparse CE: (120, B, P, 2, 62)
+        ce_all = tf.keras.losses.sparse_categorical_crossentropy(y_true_all, y_pred_exp)
+        ce_masked = ce_all * mask_all
+        num_hits = tf.reduce_sum(mask_all, axis=[2, 3, 4])  # (120, B)
+        cls_per_perm = tf.reduce_sum(ce_masked, axis=[2, 3, 4]) / (num_hits + EPSILON)
+
+        # Focal presence loss per perm
+        y_hit = mask_all
+        p0 = tf.clip_by_value(y_pred_exp[..., 0], EPSILON, 1.0 - EPSILON)
+        p_hit = 1.0 - p0
+        p_t = y_hit * p_hit + (1.0 - y_hit) * (1.0 - p_hit)
+        focal_weight = tf.pow(1.0 - p_t, focal_gamma)
+        class_weight = 1.0 + (pos_weight_presence - 1.0) * y_hit
+        bce = tf.keras.backend.binary_crossentropy(y_hit, p_hit)
+        presence_per_perm = tf.reduce_mean(
+            bce * focal_weight * class_weight, axis=[2, 3, 4]
+        )  # (120, B)
+
+        # --- Step 2: Min over permutations ---
+        total_per_perm = cls_per_perm + lambda_presence * presence_per_perm
+        best_loss = tf.reduce_min(total_per_perm, axis=0)  # (B,)
+        matched_loss = tf.reduce_mean(best_loss)
+
+        # --- Step 3: Diversity penalty (y_pred only, outside min) ---
+        C = tf.shape(y_pred)[-1]
+        elem_indices = tf.cast(tf.range(C), tf.float32)
+        soft_ids = tf.reduce_sum(y_pred * elem_indices, axis=-1)  # (B, P, 2, 62)
+
+        soft_i = tf.expand_dims(soft_ids, 2)
+        soft_j = tf.expand_dims(soft_ids, 1)
+        pairwise_dist = tf.reduce_sum(
+            tf.square(soft_i - soft_j), axis=[-2, -1]
+        )  # (B, P, P)
+
+        P = tf.shape(soft_ids)[1]
+        diag_mask = 1.0 - tf.eye(P, dtype=tf.float32)
+
+        diversity_penalty = tf.reduce_sum(
+            tf.exp(-pairwise_dist / 1000.0) * diag_mask
+        ) / (
+            tf.cast(P * (P - 1), tf.float32) * tf.cast(tf.shape(y_pred)[0], tf.float32)
+            + EPSILON
+        )
+
+        return matched_loss + lambda_diversity * diversity_penalty
+
+    return loss
+
+
 def multi_track_loss(
     lambda_presence: float = 0.2,
     pos_weight_presence: float = 5.0,
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_raw_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_raw_residuals.png
new file mode 100644
index 0000000..3e5df4d
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_raw_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_refined_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_refined_residuals.png
new file mode 100644
index 0000000..2a331e8
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair0_refined_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_raw_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_raw_residuals.png
new file mode 100644
index 0000000..23fd60b
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_raw_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_refined_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_refined_residuals.png
new file mode 100644
index 0000000..5bc07ac
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair1_refined_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_raw_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_raw_residuals.png
new file mode 100644
index 0000000..e9fa06a
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_raw_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_refined_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_refined_residuals.png
new file mode 100644
index 0000000..0aa0400
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair2_refined_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_raw_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_raw_residuals.png
new file mode 100644
index 0000000..edd5f18
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_raw_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_refined_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_refined_residuals.png
new file mode 100644
index 0000000..7d8f6b0
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair3_refined_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_raw_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_raw_residuals.png
new file mode 100644
index 0000000..64a0040
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_raw_residuals.png differ
diff --git a/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_refined_residuals.png b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_refined_residuals.png
new file mode 100644
index 0000000..1a24a0b
Binary files /dev/null and b/QTracker_training/plots/multi_track/multi_track_finder_v2_pair4_refined_residuals.png differ
diff --git a/QTracker_training/scripts/eval_multi.slurm b/QTracker_training/scripts/eval_multi.slurm
new file mode 100644
index 0000000..72629c6
--- /dev/null
+++ b/QTracker_training/scripts/eval_multi.slurm
@@ -0,0 +1,27 @@
+#!/bin/bash
+#SBATCH -A spinquest_standard
+#SBATCH -p gpu
+#SBATCH --gres=gpu:1
+#SBATCH -c 4
+#SBATCH -t 12:00:00
+#SBATCH -J multi-track-eval
+#SBATCH -o Slurm_Files/multi-track-eval.out
+#SBATCH -e Slurm_Files/multi-track-eval.err
+#SBATCH --mem=256000
+
+module purge
+
+# For apptainer exec
+APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer
+IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif
+CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/
+
+### Main Body ###
+MAX_PAIRS=3
+
+# Evaluate
+${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
+    python3 /mnt/code/eval_multi_track.py \
+     /mnt/code/data/multi_track/processed_files/mc_events_val.root \
+     /mnt/code/checkpoints/multi_track_finder_v2.keras \
+     --max_pairs $MAX_PAIRS
diff --git a/QTracker_training/scripts/preprocess_multitrack.slurm b/QTracker_training/scripts/preprocess_multi.slurm
similarity index 68%
rename from QTracker_training/scripts/preprocess_multitrack.slurm
rename to QTracker_training/scripts/preprocess_multi.slurm
index bb3a522..8ff021f 100644
--- a/QTracker_training/scripts/preprocess_multitrack.slurm
+++ b/QTracker_training/scripts/preprocess_multi.slurm
@@ -11,7 +11,7 @@
 module purge
 
 # For apptainer exec
-APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer
+APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer
 IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif
 CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/
 
@@ -90,64 +90,62 @@ CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/
 # --- 4. Generate training data by combining μ⁺ and μ⁻ signal tracks ---
 echo "Generate training data for train set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/gen_training_random.py \
-        data/raw_files/Dimuons_Train_500K_track1.root \
-        data/raw_files/Dimuons_Train_500K_track2.root \
-        --output data/multi_track/processed_files/finder_training_train.root \
-        --pairsmup 5 \
-        --pairsmum 5 \
-        --random
-
-mv momentum_training-1.root data/multi_track/processed_files/momentum_training-1_train.root 
-mv momentum_training-2.root data/multi_track/processed_files/momentum_training-2_train.root
+    python3 /mnt/code/data/multi_track/gen_training_random.py \
+        /mnt/code/data/raw_files/Dimuons_Train_500K_track1.root \
+        /mnt/code/data/raw_files/Dimuons_Train_500K_track2.root \
+        --output /mnt/code/data/multi_track/processed_files/finder_training_train.root \
+        --pairsmup 3 \
+        --pairsmum 3 \
+        --random \
+        --output_mom1 /mnt/code/data/multi_track/processed_files/momentum_training-1_train.root \
+        --output_mom2 /mnt/code/data/multi_track/processed_files/momentum_training-2_train.root
 
 echo "Generate training data for val set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/gen_training_random.py \
-        data/raw_files/Dimuons_Val_200K_track1.root \
-        data/raw_files/Dimuons_Val_200K_track2.root \
-        --output data/multi_track/processed_files/finder_training_val.root \
-        --pairsmup 5 \
-        --pairsmum 5 \
-        --random
-
-mv momentum_training-1.root data/multi_track/processed_files/momentum_training-1_val.root 
-mv momentum_training-2.root data/multi_track/processed_files/momentum_training-2_val.root
+    python3 /mnt/code/data/multi_track/gen_training_random.py \
+        /mnt/code/data/raw_files/Dimuons_Val_200K_track1.root \
+        /mnt/code/data/raw_files/Dimuons_Val_200K_track2.root \
+        --output /mnt/code/data/multi_track/processed_files/finder_training_val.root \
+        --pairsmup 3 \
+        --pairsmum 3 \
+        --random \
+        --output_mom1 /mnt/code/data/multi_track/processed_files/momentum_training-1_val.root \
+        --output_mom2 /mnt/code/data/multi_track/processed_files/momentum_training-2_val.root
 
 # --- 5. Inject background muon tracks to signal file from 4 ---
 echo "Inject low-level background tracks into train set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/messy_gen.py \
-        data/multi_track/processed_files/finder_training_train.root \
-        data/multi_track/processed_files/single_muons_train.root \
-        --output data/multi_track/processed_files/mc_events_train_low.root \
+    python3 /mnt/code/data/multi_track/messy_gen.py \
+        /mnt/code/data/multi_track/processed_files/finder_training_train.root \
+        /mnt/code/data/processed_files/single_muons_train.root \
+        --output /mnt/code/data/multi_track/processed_files/mc_events_train_low.root \
         --uniform_tracks 0 \
         --lower_bound 0 \
         --num_tracks 16
 
 echo "Inject mid-level background tracks into train set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/messy_gen.py \
-        data/multi_track/processed_files/finder_training_train.root \
-        data/multi_track/processed_files/single_muons_train.root \
-        --output data/multi_track/processed_files/mc_events_train_med.root \
+    python3 /mnt/code/data/multi_track/messy_gen.py \
+        /mnt/code/data/multi_track/processed_files/finder_training_train.root \
+        /mnt/code/data/processed_files/single_muons_train.root \
+        --output /mnt/code/data/multi_track/processed_files/mc_events_train_med.root \
         --uniform_tracks 0 \
         --lower_bound 17 \
         --num_tracks 33
         
 echo "Inject high-level background tracks into train set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/messy_gen.py \
-        data/multi_track/processed_files/finder_training_train.root \
-        data/multi_track/processed_files/single_muons_train.root \
-        --output data/multi_track/processed_files/mc_events_train_high.root \
+    python3 /mnt/code/data/multi_track/messy_gen.py \
+        /mnt/code/data/multi_track/processed_files/finder_training_train.root \
+        /mnt/code/data/processed_files/single_muons_train.root \
+        --output /mnt/code/data/multi_track/processed_files/mc_events_train_high.root \
         --uniform_tracks 0 \
         --lower_bound 34 \
         --num_tracks 50
 
 echo "Inject background tracks into val set"
 ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/multi_track/messy_gen.py \
-        data/multi_track/processed_files/finder_training_val.root \
-        data/multi_track/processed_files/single_muons_val.root \
-        --output data/multi_track/processed_files/mc_events_val.root
+    python3 /mnt/code/data/multi_track/messy_gen.py \
+        /mnt/code/data/multi_track/processed_files/finder_training_val.root \
+        /mnt/code/data/processed_files/single_muons_val.root \
+        --output /mnt/code/data/multi_track/processed_files/mc_events_val.root
diff --git a/QTracker_training/scripts/preprocess_old.slurm b/QTracker_training/scripts/preprocess_old.slurm
deleted file mode 100644
index e61b70e..0000000
--- a/QTracker_training/scripts/preprocess_old.slurm
+++ /dev/null
@@ -1,79 +0,0 @@
-#!/bin/bash
-#SBATCH -A spinquest_standard
-#SBATCH -p standard
-#SBATCH -c 4
-#SBATCH -t 72:00:00
-#SBATCH -J preprocess_old
-#SBATCH -o Slurm_Files/preprocess_old.out
-#SBATCH -e Slurm_Files/preprocess_old.err
-#SBATCH --mem=128000
-
-module purge
-
-# For apptainer exec
-APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer
-IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif
-CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/
-
-# # --- 0. Create train/val/test sets ---
-# echo "Split dimuon file into train, val, and test sets"
-# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-#     python3 data/split.py \
-#         /project/ptgroup/spinquest/QTracker-Data/JPsi-Target_10M.rus.root \
-#         --train_output data/old/raw_files/jpsi_train.root \
-#         --val_output data/old/raw_files/jpsi_val.root \
-#         --test_output data/old/raw_files/jpsi_test.root
-
-# # --- 1. Skim train and val files to reduce sample size (for downstream ML) ---
-# echo "Sample 20,000 events from raw val file (simple random)"
-# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-#     python3 data/skim.py \
-#         data/old/raw_files/jpsi_val.root \
-#         --output_file data/old/raw_files/jpsi_val_20k.root \
-#         --max_events 20000 \
-#         --random 1
-
-# # --- 2. Split signal ROOT file into μ⁺ and μ⁻ tracks ---
-# echo "Separate dimuons val set"
-# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-#     python3 data/separate.py data/old/raw_files/jpsi_val_20k.root
-
-# # --- 4. Generate training data by combining μ⁺ and μ⁻ signal tracks ---
-# echo "Generate training data for val set"
-# ${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-#     python3 data/gen_training.py \
-#         data/old/raw_files/jpsi_val_20k_track1.root \
-#         data/old/raw_files/jpsi_val_20k_track2.root \
-#         --output data/old/processed_files/finder_training_val.root
-
-# mv momentum_training-1.root data/old/processed_files/momentum_training-1_val.root 
-# mv momentum_training-2.root data/old/processed_files/momentum_training-2_val.root
-
-# --- 5. Inject background muon tracks to signal file from 4 ---
-echo "Inject background tracks into val set"
-${APPTAINER} exec --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 data/messy_gen.py \
-        data/old/processed_files/finder_training_val.root \
-        data/processed_files/single_muons_val.root \
-        --output data/old/processed_files/mc_events_val.root
-
-# Run QTracker
-${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 /mnt/code/QTracker.py \
-     /mnt/code/data/old/processed_files/mc_events_val.root \
-     --output_file /mnt/code/data/old/processed_files/qtracker_reco_old.root
-
-${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 /mnt/code/Util/imass_plot.py \
-     /mnt/code/data/old/processed_files/qtracker_reco_old.root \
-     --output_plot /mnt/code/plots/invariant_mass_old.png
-
-${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 /mnt/code/QTracker.py \
-     /project/ptgroup/spinquest/Donghwa/Qtracker_basic_old/QTracker_prod/data/processed_files/mc_events_val.root \
-     --output_file /mnt/code/data/processed_files/qtracker_reco_jpsi.root
-
-${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
-    python3 /mnt/code/Util/imass_plot.py \
-     /mnt/code/data/processed_files/qtracker_reco_jpsi.root \
-     --output_plot /mnt/code/plots/invariant_mass_jpsi.png
\ No newline at end of file
diff --git a/QTracker_training/scripts/train_multi.slurm b/QTracker_training/scripts/train_multi.slurm
index b7d755d..8c7b0d7 100644
--- a/QTracker_training/scripts/train_multi.slurm
+++ b/QTracker_training/scripts/train_multi.slurm
@@ -6,38 +6,36 @@
 #SBATCH -C gpupod
 #SBATCH -c 4
 #SBATCH -t 72:00:00
-#SBATCH -J multi-track-finder
-#SBATCH -o Slurm_Files/multi-track-finder.out
-#SBATCH -e Slurm_Files/multi-track-finder.err
+#SBATCH -J train_multi
+#SBATCH -o Slurm_Files/train_multi.out
+#SBATCH -e Slurm_Files/train_multi.err
 #SBATCH --mem=256000
 
 module purge
 
 # For apptainer exec
-APPTAINER=/apps/software/standard/core/apptainer/1.3.4/bin/apptainer
+APPTAINER=/apps/software/standard/core/apptainer/1.4.5/bin/apptainer
 IMAGE=/project/ptgroup/spinquest/David/TfRootBuild.sif
 CODEDIR=/project/ptgroup/spinquest/Donghwa/Qtracker_basic/QTracker_training/
 
-
 ### Main Body ###
-# --- Training and Evaluating Track Finder ---
 BATCH_NORM=1
 USE_ATTENTION=1
 USE_ATTENTION_FFN=0
-MAX_PAIRS=5
+MAX_PAIRS=3
 
 # GPU validation
 ${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
     echo $CUDA_VISIBLE_DEVICES
 
-# Denoise-base 16 and 32 are not too different in performance
+# Training with improved loss (min-perm matching + focal + diversity)
 ${APPTAINER} exec --nv --env CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
     python3 /mnt/code/models/MultiTrackFinder.py \
      /mnt/code/data/multi_track/processed_files/mc_events_train_low.root \
      /mnt/code/data/multi_track/processed_files/mc_events_val.root \
      --train_root_file_med /mnt/code/data/multi_track/processed_files/mc_events_train_med.root \
      --train_root_file_high /mnt/code/data/multi_track/processed_files/mc_events_train_high.root \
-     --output_model /mnt/code/checkpoints/multi_track_finder.keras \
+     --output_model /mnt/code/checkpoints/multi_track_finder_v2.keras \
      --lr_low 0.0003 \
      --lr_med 0.0001 \
      --lr_high 0.00003 \
@@ -56,11 +54,14 @@ ${APPTAINER} exec --nv --env CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES --bind $
      --factor 0.3 \
      --pos_weight 20.0 \
      --max_pairs $MAX_PAIRS \
-     --lambda_presence 0.2 \
-     --pos_weight_presence 5.0
+     --lambda_presence 1.0 \
+     --pos_weight_presence 5.0 \
+     --focal_gamma 2.0 \
+     --lambda_diversity 0.05
 
+# Evaluate
 ${APPTAINER} exec --nv --bind ${CODEDIR}:/mnt/code "${IMAGE}" \
     python3 /mnt/code/eval_multi_track.py \
      /mnt/code/data/multi_track/processed_files/mc_events_val.root \
-     /mnt/code/checkpoints/multi_track_finder.keras \
+     /mnt/code/checkpoints/multi_track_finder_v2.keras \
      --max_pairs $MAX_PAIRS