cassandra-anon/cassandra-usenix-asl-tram2

CASSANDRA — ASL configuration on TRAM2

Six independently trained CTI-BERT checkpoints that extract MITRE ATT&CK techniques from cyber threat intelligence (CTI) reports. This repository holds the recommended TRAM2 deployment system from the CASSANDRA paper, together with the per-sentence predictions needed to reproduce the TRAM2 ASL operating points below without a GPU. It does not reproduce the whole paper; the companion cassandra-usenix-evaluation-artifacts repository does that.

An extractor here is a complete configuration, not a checkpoint: encoder, classification head, loss, decision threshold, and how many checkpoints are averaged at inference. All three are stated for every number below, because checkpoint count changes both the score and the cost.

Results

Per-report F1 on the TRAM2 test split of Büchel et al., averaged over the 30 test reports that have at least one in-vocabulary technique (31 reports total; one has no in-vocabulary ground truth and is excluded, as in the source protocol).

Intervals are percentile bootstrap over the scored reports, 2,000 resamples, seeded with numpy.random.default_rng(0). The bundled verify.py ships that exact implementation and reproduces every point estimate and every interval endpoint in this table; it asserts them and fails on any mismatch. The τ = 0.63 threshold was selected on a held-out development split, never on test.

Per-checkpoint scores at τ = 0.50, with the weight variant selected on development macro-F1:

Three checkpoints, not six, is the recommendation. Across all 63 non-empty subsets of these six checkpoints, mean F1 rises 72.28 → 76.34 from one to three checkpoints (+4.07, 95% CI [2.70, 5.58], p < 0.001) but only 76.34 → 76.91 from three to six (+0.57, 95% CI [−0.51, 1.63], p = 0.317). Three checkpoints capture 87.7% of the available aggregation gain for half the storage and half the encoder passes.

The gain from averaging is driven primarily by precision. Precision rises 65.4% → 72.9% from one checkpoint to six, while recall varies little and without trend (86.7%–87.5%). A report's prediction is the union over its sentence-level decisions, so a single spurious firing anywhere in a report inserts a technique that no other sentence can retract. Averaging independently trained checkpoints attenuates firings that are specific to one training run while leaving evidence that several runs agree on, which is consistent with the effect being variance reduction — and with its returns diminishing as checkpoints are added.

Reference point

The published Llama-3.1-8B result on this benchmark is 72.50 report-F1. The three-checkpoint configuration here evaluates 373.4M parameters per sentence against that model's 8B, roughly 21× fewer. This is a published reference rather than a controlled experiment: training and inference implementations differ, so the comparison should not be read as isolating model scale or architecture.

Deployment cost (measured)

AMD Ryzen 7 7700X, 16 threads, batch size 16, sequential checkpoint inference, seven timed repeats after warm-up, on real TRAM2 sentences.

The full 3,820-sentence TRAM2 test set labels in under six minutes on this CPU at the recommended operating point. Latency and memory are implementation- and machine-specific; between-session timing varied by 25–30%.

Architecture and training

• —Encoder: ibm-research/CTI-BERT. The shared base encoder is 124,441,344 parameters; with the label-aware head each released TRAM2 checkpoint totals 124,480,513.
• —Head: label-aware attention. Each of the 50 techniques owns a learned query vector that attends over the token sequence, so different techniques can read different evidence from the same sentence, rather than sharing one [CLS] summary. The head is custom, so loading requires the bundled modeling.py.
• —Loss: asymmetric loss (ASL) with γ₊ = 0, γ₋ = 4, margin m = 0.05.
• —Optimization: AdamW, base LR 2×10⁻⁵, cosine schedule, 0.10 warm-up ratio, weight decay 0.01 (excluding bias and LayerNorm), layer-wise learning-rate decay λ = 0.85, per-device batch 16 on two visible GPUs, effective total batch 32, max sequence length 512. Mixed precision was requested, but because autocast is thread-local and does not reach DataParallel replica threads, the classification head trained in float32.
• —Budget: 30 epochs, no early stopping.
• —Weight selection: three weight variants are maintained per seed — the final ("base") parameters, an exponential moving average (decay 0.999), and a stochastic weight average over the last 25% of epochs. For each seed, the variant with the highest development macro-F1 is the one released and scored; the choice per seed is listed in the table above, and three of the six seeds selected SWA. Selection uses development data only and never touches the test split. Readers comparing against the paper should note that the released checkpoints reflect this three-way selection.
• —Split: report-level train/test split of Büchel et al. (120 train / 31 test reports). Development is an 80:20 sentence-level split of the training reports at random_state=0, identical across seeds.

Scope and limitations

• —Evaluated on the 50 TRAM2 techniques, not the full ~691-technique ATT&CK matrix.
• —The test split has 30 scored reports, so confidence intervals are wide. Directional differences smaller than a few F1 points are not resolvable at this sample size.
• —ASL produces the higher observed ensemble score on TRAM2, although its advantage over binary cross-entropy is not statistically established (+2.89 F1, 95% CI [−1.31, 7.41], p = 0.188). On the sparser AnnoCTR benchmark ASL is significantly worse than BCE (−4.13 F1, 95% CI [−6.70, −1.72], p < 0.001), where BCE is preferred. We tested the natural explanation — that AnnoCTR simply offers fewer positive examples per technique — in a preregistered 300-run intervention, and the preregistered decision rule was not met. We therefore report the benchmark dependence as an observation and deliberately do not offer label density as a loss-selection rule. Validate the loss on your own development data.
• —Techniques absent from training are not recovered. This system has no demonstrated zero-shot capability.
• —Intended as decision support for analysts, not an autonomous labeling authority. Rare-technique predictions in particular warrant review.

Usage

pip install -r requirements.txt
python inference_example.py --threshold 0.5     # predetermined 3-checkpoint system (76.70)
python inference_example.py --all-seeds         # all six checkpoints (76.91)

inference_example.py defaults to the predetermined three checkpoints (seeds 42, 123, 456) — the recommended deployment system, and the configuration reported as 76.70. Pass --all-seeds for the six-checkpoint system at roughly twice the storage and encoder passes per sentence.

from modeling import load_seed, predict_ensemble

seeds = [load_seed(f"seeds/seed-{s}") for s in (42, 123, 456)]   # recommended: 3 checkpoints
preds = predict_ensemble(
    seeds,
    ["The malware created a scheduled task to maintain persistence."],
    threshold=0.5,
)

Use τ = 0.50 with three checkpoints for the recommended operating point, or τ = 0.63 with all six to reproduce the maximum observed score.

Verifying the reported numbers

Every seed directory contains seed_probs.npz with the per-sentence sigmoid probabilities, the test gold labels, and the report boundaries. This is sufficient to recompute every F1, confidence interval, threshold sweep, and subset score in this card without running the model:

python verify.py          # requires only numpy

It recomputes 76.70, 76.91, and 77.43, the per-seed scores, the subset curve, and the confidence intervals from the released predictions, and fails loudly on any mismatch. It verifies the TRAM2 ASL operating points only; the companion artifacts repository verifies the paper's other results.

Licensing

The code in this repository (modeling.py, inference_example.py, verify.py) is released under Apache-2.0; see LICENSE.

The weights are a derivative of ibm-research/CTI-BERT. That model card currently declares no license metadata, so we cannot assert that Apache-2.0 extends to the derived weights. Users intending to redistribute or deploy commercially should confirm the upstream terms with the base model's authors. We make no license claim over the fine-tuned weights beyond the code.

Related repositories

• —cassandra-usenix-evaluation-artifacts — all preserved prediction outputs, ablation and matched controlled comparisons, subset analysis, measured deployment costs, and analysis scripts.

Citation

Paper under anonymous review; citation will be added on acceptance.