Skip to content

Latest commit

 

History

4 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Distributed CyberRAG

Utility-aware peer selection and distributed retrieval-augmented generation for cybersecurity knowledge bases.

This project studies a practical question in distributed RAG:

Given a cybersecurity question and multiple heterogeneous knowledge peers, which peer should the system query?

Rather than assuming that the most semantically relevant peer is the most useful one, Distributed CyberRAG measures downstream answer utility directly. The current evaluation focuses on Root Cause Mapping (RCM), where vulnerability descriptions are mapped to CWE weakness identifiers using evidence retrieved from NVD, CWE, MITRE ATT&CK, CAPEC, and CISA KEV.

The main result is simple but important: semantic relevance is a poor proxy for actual peer utility. On the full 290-question RCM benchmark, NVD achieves 80.7% accuracy, while an oracle that selects the best singleton peer reaches 85.2%. CWE uniquely rescues 13 questions that NVD misses, but lightweight learned routers recover almost none of this oracle headroom.


Overview

A conventional RAG pipeline retrieves from a single centralized corpus:

Question
   |
   v
Retriever
   |
   v
Reranker
   |
   v
Generator
   |
   v
Answer

Distributed CyberRAG instead assumes that evidence is distributed across specialized peers:

                         +------------------+
                         |       NVD        |
                         +------------------+
                                  |
                         +------------------+
                         |       CWE        |
                         +------------------+
                                  |
Question ---> Router --->+------------------+
                         | MITRE ATT&CK     |
                         +------------------+
                                  |
                         +------------------+
                         |      CAPEC       |
                         +------------------+
                                  |
                         +------------------+
                         |    CISA KEV      |
                         +------------------+
                                  |
                                  v
                           Local Retrieval
                                  |
                                  v
                              Reranker
                                  |
                                  v
                              Generator
                                  |
                                  v
                                Answer

The central problem becomes utility-aware routing:

Which peer is most likely to improve the final answer?

This is different from:

Which peer description is most semantically similar to the query?

The experiments in this repository quantify that distinction.


Knowledge Peers

The system currently evaluates five cybersecurity knowledge sources:

Peer Role
NVD CVE descriptions, weakness mappings, severity, and vulnerability metadata
CWE Weakness definitions and taxonomy
MITRE ATT&CK Adversary techniques and behaviors
CAPEC Attack patterns
CISA KEV Known exploited vulnerabilities

Each peer maintains an independent local vector index.

A query can therefore be evaluated against individual peers or combinations of peers without merging all source data into one centralized index.


Task: Root Cause Mapping

The primary evaluation task is RCM, or Root Cause Mapping.

Given a natural-language vulnerability description, the model must predict the corresponding CWE identifier.

Example:

Question:
A web application's file-management feature lets remote users supply a
filename or path parameter when deleting files. Due to insufficient
validation, an attacker can use path traversal sequences to delete files
outside the intended directory.

Target:
CWE-73

Another example:

Question:
A storage operating system has a flaw where arithmetic on numeric values
can exceed the maximum representable integer and wrap around.

Target:
CWE-190

The full RCM benchmark contains:

290 questions

Pipeline

For a query (q) and peer (p), the evaluation pipeline is:

Query
  |
  v
Local Vector Retrieval
  |
  v
Top-k Candidate Evidence
  |
  v
Cross-Encoder Reranking
  |
  v
Top Evidence Passages
  |
  v
Candidate CWE Extraction
  |
  v
Constrained Identifier Selection
  |
  v
Predicted CWE
  |
  v
Exact-Match Utility

Retrieval

Local peer indexes use dense embeddings.

Default embedding model:

sentence-transformers/all-MiniLM-L6-v2

The NVD peer uses a larger retrieval depth because of its larger and more heterogeneous corpus.

Typical configuration:

NVD:        local_k = 50
other peers: local_k <= 12

Reranking

Retrieved evidence is reranked using:

BAAI/bge-reranker-v2-m3

The highest-ranked evidence passages are passed downstream.

Typical configuration:

evidence_k = 8

Generation

Experiments use Qwen causal language models, including:

Qwen/Qwen3-1.7B
Qwen/Qwen3-8B

For RCM, unconstrained free-form generation was found to be unreliable because the model could produce explanations instead of the required identifier.

The RCM pipeline therefore extracts candidate CWE identifiers from retrieved evidence and performs constrained identifier selection over those candidates.

Conceptually:

Evidence
   |
   v
Extract identifiers
   |
   v
[CWE-335, CWE-341, CWE-384, ...]
   |
   v
Score/select candidate
   |
   v
CWE-384

This ensures that RCM predictions are valid identifiers supported by retrieved evidence.


Utility Definition

For RCM, answer utility is exact-match correctness:

utility = 1    predicted identifier is a valid gold identifier
utility = 0    otherwise

The marginal utility of adding a peer is:

gain = utility(new peer set) - utility(base peer set)

Therefore:

gain = +1    peer addition fixes an incorrect answer
gain =  0    answer correctness does not change
gain = -1    peer addition breaks a previously correct answer

This allows peer usefulness to be measured directly from downstream behavior rather than semantic similarity.


Experimental Results

Singleton Peer Accuracy

Each peer was evaluated independently across all 290 RCM questions.

Peer Correct Accuracy
NVD 234 / 290 80.7%
CWE 60 / 290 20.7%
MITRE ATT&CK 0 / 290 0.0%
CAPEC 0 / 290 0.0%
CISA KEV 0 / 290 0.0%

NVD is overwhelmingly the strongest singleton peer for RCM.

However, NVD is not sufficient for every question.


NVD vs. CWE

The two useful RCM peers exhibit complementary behavior.

Outcome Questions Fraction
NVD only correct 187 64.5%
CWE only correct 13 4.5%
Both correct 47 16.2%
Both wrong 43 14.8%

The key result is the 13 CWE-only questions.

These represent cases where selecting CWE instead of the much stronger NVD peer improves the answer.

Examples include weakness classes involving:

CWE-73   External Control of File Name or Path
CWE-502  Deserialization of Untrusted Data
CWE-190  Integer Overflow or Wraparound
CWE-77   Command Injection
CWE-89   SQL Injection
CWE-297  Improper Validation of Certificate with Host Mismatch
CWE-681  Incorrect Conversion between Numeric Types
CWE-284  Improper Access Control
CWE-427  Uncontrolled Search Path Element
CWE-129  Improper Validation of Array Index

This establishes meaningful routing headroom even though NVD dominates globally.


Oracle Singleton Router

An oracle router chooses a peer that produces a correct answer whenever any singleton peer can answer the question correctly.

Results:

Router Correct Accuracy
Always NVD 234 / 290 80.7%
Oracle singleton 247 / 290 85.2%

Oracle headroom:

+13 questions
+4.5 percentage points

Therefore, peer routing can theoretically improve RCM performance over simply querying NVD.

The challenge is predicting those 13 switches without damaging the much larger set of questions where NVD is already correct.


Why Similarity Routing Fails

A natural baseline is to route to whichever peer profile is most semantically similar to the query.

This performs extremely poorly.

Router Correct Accuracy
Always NVD 234 / 290 80.7%
Similarity router 65 / 290 22.4%
Oracle singleton 247 / 290 85.2%

Routing distribution:

CWE       279 / 290 = 96.2%
NVD        10 / 290 =  3.4%
CISA KEV    1 / 290 =  0.3%

The similarity router almost always chooses CWE.

This happens because CWE's peer description is semantically close to root-cause questions, even though NVD evidence is much more useful for actually answering them.

Measured correlation between profile similarity and downstream utility is weak:

Pearson correlation:  +0.1536
Spearman correlation: +0.2073

The result demonstrates that:

Semantic relevance and answer utility are not equivalent.

A peer can sound highly relevant to the task while still being a poor evidence source for producing the correct answer.


Learned Utility Routers

The repository evaluates lightweight learned routers that attempt to predict when the system should switch from NVD to CWE.

The target is:

switch_to_cwe = 1

only when:

NVD is wrong
AND
CWE is correct

This is highly imbalanced:

positive examples: 13 / 290
positive rate: 4.5%

Retrieval-Feature Router

Features include:

NVD top-1 reranker score
NVD top-2 reranker score
NVD reranker margin
NVD mean reranker score
NVD reranker-score standard deviation
NVD profile similarity

CWE top-1 reranker score
CWE top-2 reranker score
CWE reranker margin
CWE mean reranker score
CWE reranker-score standard deviation
CWE profile similarity

CWE - NVD feature deltas

Models evaluated include:

Logistic Regression
Histogram Gradient Boosting

Five-fold out-of-fold evaluation is used.

Best observed routing accuracy:

235 / 290 = 81.0%

This is only one question better than always selecting NVD.


Semantic Router

A second router incorporates query embeddings.

Queries are embedded using:

sentence-transformers/all-MiniLM-L6-v2

PCA representations were evaluated at:

8 dimensions
16 dimensions
32 dimensions

These semantic features are combined with routing features and evaluated using cross-validation.

Best observed result:

235 / 290 = 81.0%

Again, this recovers almost none of the oracle headroom.


Final Routing Comparison

Strategy Correct Accuracy
Always NVD 234 / 290 80.7%
Always CWE 60 / 290 20.7%
Similarity router 65 / 290 22.4%
Retrieval-feature router 235 / 290 81.0%
Semantic router 235 / 290 81.0%
Oracle singleton 247 / 290 85.2%

The learned routers demonstrate that identifying the small set of useful CWE switches is substantially harder than measuring aggregate peer relevance.


Multi-Peer Utility

A pilot experiment also evaluates whether combining peers improves performance.

For the first 50 RCM questions, singleton and pairwise marginal utilities were measured.

Pairwise observations:

positive:     0   (0.0%)
negative:     8   (2.8%)
neutral:    282  (97.2%)

No hard singleton failure was rescued by pairwise aggregation in the pilot.

Some combinations actively degraded correct predictions.

Examples:

NVD + CWE
NVD + CISA KEV
CWE + NVD

This suggests that naively increasing the amount of retrieved evidence does not necessarily improve generation.

Additional peers can introduce distracting candidates or evidence that changes the final identifier selection.


Failure Analysis

The oracle singleton still fails on:

43 / 290 questions

These residual failures were decomposed into pipeline stages.

Final validated decomposition:

Failure Stage Count Fraction
Selection 31 72.1%
Reranking 6 14.0%
Retrieval 6 14.0%

Retrieval Failure

The correct CWE does not appear in the retrieved candidate pool.

Query
  |
  v
Vector Retrieval
  |
  +---- correct evidence absent

These failures require better retrieval, indexing, query transformation, or broader candidate generation.


Reranking Failure

The relevant evidence exists in the retrieved pool but is removed before the final evidence set.

Retrieved candidates
  |
  | correct evidence exists
  v
Reranker
  |
  +---- correct evidence falls outside top-k

Potential improvements include larger evidence windows, stronger rerankers, or diversity-aware reranking.


Selection Failure

The correct CWE is present in the available evidence, but the final model selects another identifier.

Evidence
  |
  | contains correct CWE
  v
Candidate extraction
  |
  v
Candidate selection
  |
  +---- wrong CWE chosen

Selection accounts for:

31 / 43 = 72.1%

of residual oracle failures.

This is the dominant remaining bottleneck.


Stability Validation

Failure decomposition was validated against live retrieval and reranking rather than relying solely on cached evidence.

Three cases whose live top-8 evidence differed from cached results were rerun three times:

rcm-043
rcm-044
rcm-048

The reranked results were deterministic across all three repetitions for each peer.

Example:

rcm-043

NVD:
majority_hit = False (0/3)

CWE:
majority_hit = True (3/3)
rcm-044

NVD:
majority_hit = True (3/3)

CWE:
majority_hit = False (0/3)
rcm-048

NVD:
majority_hit = False (0/3)

CWE:
majority_hit = True (3/3)

This indicates that the observed differences were not caused by stochastic reranking instability.


Key Findings

The RCM experiments support six main conclusions.

1. Semantic relevance is a poor proxy for peer utility

The similarity router achieves only:

22.4%

compared with:

80.7%

for simply querying NVD.

Peer descriptions capture topical relevance but not whether retrieved evidence will produce the correct downstream answer.

2. NVD dominates RCM

NVD alone answers:

234 / 290 = 80.7%

of the benchmark correctly.

For this task, it should therefore be treated as the default peer.

3. CWE provides real but sparse complementary utility

CWE uniquely rescues:

13

NVD failures.

This raises oracle singleton performance to:

247 / 290 = 85.2%

4. The routing problem is highly asymmetric

A useful router must identify only 13 beneficial switches while avoiding a much larger number of harmful switches.

Therefore, routing should be treated as a high-precision intervention problem, not ordinary balanced classification.

5. Naive multi-peer aggregation is not automatically beneficial

The 50-question pairwise pilot produced:

0 positive marginal interactions
8 negative interactions
282 neutral interactions

More evidence can introduce noise and degrade identifier selection.

6. Selection is now the dominant bottleneck

After retrieval and reranking analysis:

72.1% of residual failures are selection failures.

Future improvements should therefore focus heavily on evidence-grounded candidate selection rather than only increasing retrieval depth.


Repository Structure

distributed-cyberrag/
├── data/
│   └── cticonnect/
│       └── entity_linking/
│           └── rcm.jsonl
│
├── src/
│   └── cyberrag/
│       ├── index.py
│       ├── models.py
│       └── types.py
│
├── scripts/
│   ├── compute_oracle_utilities.py
│   ├── check_gold_candidate_recall.py
│   ├── build_router_features.py
│   ├── evaluate_router_models.py
│   ├── evaluate_semantic_router.py
│   ├── analyze_rcm_failures.py
│   ├── analyze_rcm_failure_stages.py
│   ├── validate_rcm_failure_decomposition.py
│   ├── check_unstable_failure_cases.py
│   ├── build_final_rcm_summary.py
│   └── plot_final_rcm_results.py
│
├── results/
│   ├── final_rcm/
│   │   ├── figures/
│   │   │   ├── failure_decomposition.png
│   │   │   ├── nvd_cwe_outcomes.png
│   │   │   ├── peer_accuracy.png
│   │   │   └── routing_accuracy.png
│   │   │
│   │   ├── summary/
│   │   │   ├── failure_decomposition.csv
│   │   │   ├── final_results.txt
│   │   │   ├── nvd_cwe_outcomes.csv
│   │   │   ├── peer_accuracy.csv
│   │   │   └── routing_accuracy.csv
│   │   │
│   │   ├── oracle_rcm_singleton.jsonl
│   │   ├── rcm_failure_analysis.jsonl
│   │   ├── rcm_failure_stage_analysis.jsonl
│   │   ├── router_features_rcm.jsonl
│   │   ├── router_oof_hist_gb.jsonl
│   │   ├── router_oof_logistic.jsonl
│   │   ├── router_semantic_pca8.jsonl
│   │   ├── router_semantic_pca16.jsonl
│   │   └── router_semantic_pca32.jsonl
│   │
│   └── pilot_rcm/
│       └── oracle_rcm_pairwise_50q.jsonl
│
├── pyproject.toml
└── README.md

Installation

Clone the repository:

git clone https://github.com/namtran1812/distributed-cyberrag.git
cd distributed-cyberrag

Create a virtual environment:

python -m venv .venv
source .venv/bin/activate

Install the package:

pip install -e .

For development:

pip install -e ".[dev]"

The project requires Python 3.10 or newer.

Core dependencies include:

numpy
requests
pyyaml
tqdm
sentence-transformers
transformers
accelerate
torch
faiss-cpu
matplotlib

Running the RCM Evaluation

Singleton Evaluation

Example:

python scripts/compute_oracle_utilities.py \
  --task rcm \
  --stage singleton \
  --embedding-model sentence-transformers/all-MiniLM-L6-v2 \
  --reranker BAAI/bge-reranker-v2-m3 \
  --generator Qwen/Qwen3-1.7B \
  --device cuda \
  --local-k 50 \
  --evidence-k 8

This evaluates each knowledge peer independently.


Candidate Recall Analysis

python scripts/check_gold_candidate_recall.py

This checks whether the gold CWE identifier is recoverable from each peer's retrieved and reranked evidence.


Build Router Features

python scripts/build_router_features.py

The generated feature dataset is written to:

results/router_features_rcm.jsonl

Evaluate Retrieval-Based Routers

python scripts/evaluate_router_models.py

This evaluates lightweight learned routing models using cross-validation.

Outputs include:

results/router_oof_logistic.jsonl
results/router_oof_hist_gb.jsonl

Evaluate Semantic Routers

python scripts/evaluate_semantic_router.py

This evaluates query-embedding routers using PCA-reduced semantic representations.

Outputs include:

results/router_semantic_pca8.jsonl
results/router_semantic_pca16.jsonl
results/router_semantic_pca32.jsonl

Analyze Failures

python scripts/analyze_rcm_failures.py

Then:

python scripts/analyze_rcm_failure_stages.py

These scripts distinguish failures caused by candidate availability, retrieval, reranking, and final identifier selection.


Validate Failure Decomposition

python scripts/validate_rcm_failure_decomposition.py

For unstable or changed top-k cases:

python scripts/check_unstable_failure_cases.py

Build Final Summary

python scripts/build_final_rcm_summary.py

Summary artifacts are written under:

results/final_rcm/summary/

Generate Figures

python scripts/plot_final_rcm_results.py

Generated figures:

results/final_rcm/figures/peer_accuracy.png
results/final_rcm/figures/routing_accuracy.png
results/final_rcm/figures/nvd_cwe_outcomes.png
results/final_rcm/figures/failure_decomposition.png

Testing

Run:

pytest -q

Current test suite:

3 passed

Python syntax can also be checked with:

python -m py_compile \
  scripts/compute_oracle_utilities.py \
  scripts/analyze_rcm_failure_stages.py \
  scripts/analyze_rcm_failures.py \
  scripts/build_final_rcm_summary.py \
  scripts/build_router_features.py \
  scripts/check_gold_candidate_recall.py \
  scripts/check_unstable_failure_cases.py \
  scripts/evaluate_router_models.py \
  scripts/evaluate_semantic_router.py \
  scripts/plot_final_rcm_results.py \
  scripts/validate_rcm_failure_decomposition.py \
  src/cyberrag/models.py

Reproducibility

The final RCM evaluation artifacts are preserved under:

results/final_rcm/

The earlier 50-question pairwise experiment is intentionally separated as:

results/pilot_rcm/oracle_rcm_pairwise_50q.jsonl

This distinction is important because the singleton results cover the complete 290-question benchmark, while pairwise marginal utility was evaluated only as a pilot on 50 questions.


Current Limitations

The current evaluation has several limitations.

First, RCM strongly favors NVD. Results should therefore not be interpreted as evidence that NVD is universally the best cybersecurity peer. Other tasks may exhibit substantially different peer-utility distributions.

Second, only 13 of 290 questions benefit uniquely from switching from NVD to CWE. This makes learned routing extremely imbalanced and limits the amount of supervision available to lightweight models.

Third, the pairwise aggregation experiment currently covers only a 50-question pilot rather than the complete RCM benchmark.

Fourth, exact-match RCM evaluation measures identifier correctness but does not capture answer quality for open-ended cybersecurity tasks.

Finally, the current routing experiments primarily evaluate singleton selection. More sophisticated conditional retrieval policies could decide not only which peer to query but also when additional peers are worth their computational and noise costs.


Future Work

Several directions follow naturally from the current results.

Confidence-Aware Routing

Instead of predicting the best peer directly, estimate whether the default NVD answer is likely to fail:

if confidence(NVD) is high:
    stop
else:
    query CWE

This reframes routing as selective escalation.

Candidate-Level Utility Modeling

Rather than predicting peer utility from global retrieval statistics, model individual candidate CWE identifiers and their supporting evidence.

Evidence Diversity

Pairwise aggregation may benefit from selecting complementary evidence rather than simply concatenating top-ranked passages from multiple peers.

Query-Adaptive Retrieval Depth

Different vulnerability classes may require different retrieval depths:

simple mapping      -> small k
ambiguous mapping   -> larger k

Hierarchical CWE Reasoning

CWE's taxonomy provides parent/child relationships that could help distinguish semantically related weakness categories during final selection.

Additional CTI Tasks

The routing hypothesis should be tested on tasks where ATT&CK, CAPEC, and CISA KEV are expected to provide direct utility rather than evaluating them only through RCM.


Research Takeaway

The central lesson from the current experiments is:

The most relevant-looking knowledge source is not necessarily the source that produces the best answer.

For RCM, CWE appears semantically aligned with nearly every query, causing a similarity router to choose it for 96.2% of questions. Yet NVD answers 80.7% correctly while CWE answers only 20.7%.

At the same time, simply hard-coding NVD leaves measurable performance on the table: CWE uniquely solves 13 questions, giving an oracle router 85.2% accuracy.

Distributed RAG therefore requires a notion of utility, not merely relevance.

The problem is not just:

Where is relevant information?

but:

Which source is likely to change the downstream answer for the better?

The current results show that answering the second question is substantially harder—and potentially more important—than conventional semantic routing.


Final RCM Results

Questions: 290

Peer accuracy
-------------
NVD:       234/290 = 80.7%
CWE:        60/290 = 20.7%
ATT&CK:      0/290 = 0.0%
CAPEC:       0/290 = 0.0%
CISA KEV:    0/290 = 0.0%

Routing
-------
Always NVD:         234/290 = 80.7%
Similarity router:   65/290 = 22.4%
Retrieval router:   235/290 = 81.0%
Semantic router:    235/290 = 81.0%
Oracle singleton:   247/290 = 85.2%

Oracle headroom over NVD
------------------------
13 questions
+4.5 percentage points

NVD/CWE outcomes
----------------
NVD only:       187
CWE only:        13
Both correct:    47
Both wrong:      43

Residual oracle failures
------------------------
Selection: 31/43 = 72.1%
Reranking:  6/43 = 14.0%
Retrieval:  6/43 = 14.0%

License

This repository is intended for research and experimentation in distributed retrieval-augmented generation and cybersecurity knowledge retrieval.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages