File size: 3,076 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 | """Non-neural baselines.
These exist to make the fine-tuned encoder falsifiable. A large share of
published cyber-ML results do not clear a TF-IDF + logistic regression floor,
and a model that fails to beat these has no reason to be downloaded.
* ``frequency`` — always predict the k most common training techniques.
Establishes what a model that has learned nothing about the input scores.
* ``keyword`` — substring match on ATT&CK technique names. Zero training.
* ``tfidf_lr`` — TF-IDF word+char n-grams, one-vs-rest logistic regression.
The real floor. Usually much stronger than people expect.
"""
from __future__ import annotations
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import FeatureUnion
from . import attack_meta
def frequency_scores(train_Y: np.ndarray, n_eval: int) -> np.ndarray:
"""Score every technique by its training prevalence, identically for all rows."""
prior = train_Y.mean(axis=0)
return np.tile(prior, (n_eval, 1))
def keyword_scores(texts: list[str], labels: list[str]) -> np.ndarray:
"""1.0 where the technique's name appears verbatim in the sentence, else 0."""
kw = attack_meta.technique_keywords(labels)
scores = np.zeros((len(texts), len(labels)), dtype=np.float32)
lowered = [t.lower() for t in texts]
for j, tid in enumerate(labels):
forms = kw.get(tid, [])
if not forms:
continue
for i, text in enumerate(lowered):
if any(f in text for f in forms):
scores[i, j] = 1.0
return scores
def _vectorizer() -> FeatureUnion:
return FeatureUnion([
("word", TfidfVectorizer(
ngram_range=(1, 2), min_df=2, sublinear_tf=True,
strip_accents="unicode", lowercase=True)),
("char", TfidfVectorizer(
analyzer="char_wb", ngram_range=(3, 5), min_df=3, sublinear_tf=True,
lowercase=True)),
])
def tfidf_lr_scores(
train_texts: list[str], train_Y: np.ndarray, eval_sets: dict[str, list[str]]
) -> dict[str, np.ndarray]:
"""One-vs-rest logistic regression over TF-IDF word + character n-grams.
Classes with a single training polarity (all-negative after splitting) are
handled explicitly rather than allowed to raise, and score a constant 0.
"""
vec = _vectorizer()
Xtr = vec.fit_transform(train_texts)
Xev = {name: vec.transform(texts) for name, texts in eval_sets.items()}
out = {name: np.zeros((X.shape[0], train_Y.shape[1]), dtype=np.float32)
for name, X in Xev.items()}
for j in range(train_Y.shape[1]):
y = train_Y[:, j]
if y.sum() == 0 or y.sum() == len(y):
continue # degenerate class: leave scores at 0
clf = LogisticRegression(
C=4.0, max_iter=2000, class_weight="balanced", solver="liblinear")
clf.fit(Xtr, y)
for name, X in Xev.items():
out[name][:, j] = clf.predict_proba(X)[:, 1]
return out
|