File size: 5,079 Bytes
468c4c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """Metrics and threshold selection for multi-label ATT&CK classification.
Design notes that matter for reading the numbers:
* **macro-F1 is the headline, not micro-F1.** 79% of sentences carry no label
and the technique distribution is heavily long-tailed (T1027 has 678 examples,
T1072 has ~20). Micro-F1 is dominated by a handful of frequent techniques and
flatters every model. Macro-F1 is what tells you whether the tail works.
* **Two threshold regimes are reported.** A single global threshold, and
per-class thresholds tuned on dev. Per-class tuning fits 49 free parameters on
a dev set with very few positives per class, so it can overfit; reporting both
makes the size of that effect visible instead of hiding it.
"""
from __future__ import annotations
import json
from dataclasses import dataclass
import numpy as np
from . import config
def to_matrix(records: list[dict], labels: list[str]) -> np.ndarray:
index = {l: i for i, l in enumerate(labels)}
Y = np.zeros((len(records), len(labels)), dtype=np.int8)
for i, rec in enumerate(records):
for l in rec["labels"]:
j = index.get(l)
if j is not None:
Y[i, j] = 1
return Y
def _prf(tp: float, fp: float, fn: float) -> tuple[float, float, float]:
p = tp / (tp + fp) if tp + fp else 0.0
r = tp / (tp + fn) if tp + fn else 0.0
f = 2 * p * r / (p + r) if p + r else 0.0
return p, r, f
@dataclass
class Report:
micro_f1: float
macro_f1: float
micro_precision: float
micro_recall: float
macro_precision: float
macro_recall: float
per_class: dict
n_examples: int
n_positive_labels: int
def as_dict(self) -> dict:
return {
"micro_f1": round(self.micro_f1, 4),
"macro_f1": round(self.macro_f1, 4),
"micro_precision": round(self.micro_precision, 4),
"micro_recall": round(self.micro_recall, 4),
"macro_precision": round(self.macro_precision, 4),
"macro_recall": round(self.macro_recall, 4),
"n_examples": self.n_examples,
"n_positive_labels": self.n_positive_labels,
"per_class": self.per_class,
}
def evaluate(Y_true: np.ndarray, Y_pred: np.ndarray, labels: list[str]) -> Report:
tp = (Y_true & Y_pred).sum(axis=0).astype(float)
fp = ((1 - Y_true) & Y_pred).sum(axis=0).astype(float)
fn = (Y_true & (1 - Y_pred)).sum(axis=0).astype(float)
mip, mir, mif = _prf(tp.sum(), fp.sum(), fn.sum())
per_class, ps, rs, fs = {}, [], [], []
for j, name in enumerate(labels):
p, r, f = _prf(tp[j], fp[j], fn[j])
per_class[name] = {
"precision": round(p, 4),
"recall": round(r, 4),
"f1": round(f, 4),
"support": int(Y_true[:, j].sum()),
"predicted": int(Y_pred[:, j].sum()),
}
ps.append(p)
rs.append(r)
fs.append(f)
return Report(
micro_f1=mif,
macro_f1=float(np.mean(fs)),
micro_precision=mip,
micro_recall=mir,
macro_precision=float(np.mean(ps)),
macro_recall=float(np.mean(rs)),
per_class=per_class,
n_examples=int(Y_true.shape[0]),
n_positive_labels=int(Y_true.sum()),
)
def tune_global_threshold(
Y_true: np.ndarray, scores: np.ndarray, grid=None
) -> tuple[float, float]:
"""Pick the single threshold maximising macro-F1 on the given (dev) set."""
grid = grid or config.THRESHOLD_GRID
best_t, best_f = 0.5, -1.0
for t in grid:
Y_pred = (scores >= t).astype(np.int8)
tp = (Y_true & Y_pred).sum(axis=0).astype(float)
fp = ((1 - Y_true) & Y_pred).sum(axis=0).astype(float)
fn = (Y_true & (1 - Y_pred)).sum(axis=0).astype(float)
f = float(np.mean([_prf(tp[j], fp[j], fn[j])[2] for j in range(Y_true.shape[1])]))
if f > best_f:
best_t, best_f = t, f
return best_t, best_f
def tune_per_class_thresholds(
Y_true: np.ndarray, scores: np.ndarray, grid=None
) -> np.ndarray:
"""One threshold per technique, each maximising that technique's dev F1."""
grid = grid or config.THRESHOLD_GRID
out = np.full(Y_true.shape[1], 0.5)
for j in range(Y_true.shape[1]):
best_t, best_f = 0.5, -1.0
yt = Y_true[:, j]
for t in grid:
yp = (scores[:, j] >= t).astype(np.int8)
_, _, f = _prf(
float((yt & yp).sum()),
float(((1 - yt) & yp).sum()),
float((yt & (1 - yp)).sum()),
)
if f > best_f:
best_t, best_f = t, f
out[j] = best_t
return out
def apply_thresholds(scores: np.ndarray, thresholds) -> np.ndarray:
return (scores >= np.asarray(thresholds)).astype(np.int8)
def save_report(name: str, scheme: str, reports: dict) -> None:
path = config.RESULTS_DIR / f"{name}__{scheme}.json"
path.write_text(json.dumps(reports, indent=2), encoding="utf-8")
print(f" -> {path.relative_to(config.REPO_ROOT)}")
|