Add Uncertainty Metric to Naive Bayes Classifier Using Conformal Prediction

[prayas7102]

Enhancement: Add Uncertainty Metric to Naive Bayes Classifier

Description

To improve the reliability and interpretability of our Naive Bayes classifier, we should implement an uncertainty metric using conformal prediction. This enhancement will provide a measure of confidence for each classification, allowing for more nuanced decision-making and potentially improving the overall accuracy of our security analysis.

Proposed Implementation

1. Implement conformal prediction algorithm alongside our existing Naive Bayes classifier.
2. Calculate prediction sets and confidence scores for each classification.
3. Integrate the uncertainty metric into the classification output.
4. Update the reporting mechanism to include uncertainty information.

Benefits

• Improved reliability of classification results
• Better handling of edge cases and ambiguous inputs
• Enhanced decision-making capabilities for security analysis
• More informative output for users of NodejsSecurify

Tasks

• Research and select appropriate conformal prediction method
• Implement conformal prediction algorithm
• Integrate uncertainty metric with existing Naive Bayes classifier
• Update classification output format to include uncertainty information
• Modify reporting system to display uncertainty metrics
• Write tests for the new functionality
• Update documentation to reflect the new feature

Additional Notes

This enhancement aligns with our goal of providing comprehensive and accurate security analysis for Node.js applications. The uncertainty metric will add another layer of sophistication to our machine learning-based approach.

[Commit2Cosmos]

Hi @prayas7102, I can try to tackle this.
From what I understand ICP can be used. There are no JS libraries that do this, so I'd have to implement it from scratch. Because there're 4 files (SecurityHeaders, BruteForceAttack, InputValidation and InsecureAuthentication) that use the BayesClassifier it probably makes sence to write a function dealing with CP in a separate file and apply it to each instance.

Here's my pseudo solution:

import { BayesClassifier } from 'natural';

function trainClassifier(tokenizedSamples: string[][], labels: number[]): BayesClassifier {
    const classifier = new BayesClassifier();
    for (let i = 0; i < tokenizedSamples.length; i++) {
        classifier.addDocument(tokenizedSamples[i], labels[i].toString());
    }
    classifier.train();
    return classifier;
}

function calibrate(classifier: BayesClassifier, calibrationData: string[][], labels: number[]) {
    let nonConformityScores: { predicted: string, nonConformity: number }[] = [];
    for (let i = 0; i < calibrationData.length; i++) {
        const label = labels[i];
        const data = calibrationData[i];
        
        // Get predicted label and class probabilities
        const predicted = classifier.classify(data);
        const probabilities = classifier.getClassifications(data);

        //* Ensure probabilities is always ordered as [ { label: '0', value: ... }, { label: '1', value: ... } ]
        probabilities.sort((a, _) => a.label === '0' ? -1 : 1);

        const confidence = probabilities[label].value;
        
        const nonConformity = Math.abs(confidence - 0.5);
        
        nonConformityScores.push({
            predicted,
            nonConformity
        });
    }
    return nonConformityScores;
}

function predictWithConformal(
    classifier: BayesClassifier,
    newExample: string[],
    calibrationScores: { predicted: string, nonConformity: number }[], 
    confidenceLevel: number
) {
    const predictionSet: string[] = [];
    
    //* Calculate the nonconformity score for the new example
    const probabilities = classifier.getClassifications(newExample);
    const predicted = classifier.classify(newExample);
    const confidence = probabilities.find(p => p.label === predicted)?.value || 0;
    const nonConformity = Math.abs(confidence - 0.5);
    
    //* Compute the proportion of calibration scores greater than or equal to the current nonConformity score
    for (let i = 0; i < calibrationScores.length; i++) {
        const score = calibrationScores[i];
        const proportion = (calibrationScores.filter(s => s.nonConformity >= nonConformity).length) / calibrationScores.length;
        
        if (proportion > 1 - confidenceLevel) {
            predictionSet.push(score.predicted);
        }
    }

    //* If no valid prediction set is found, return the default classifier prediction
    return predictionSet.length ? predictionSet : [predicted];
}

const tokenizedSnippet: string[] = ['this', 'is', 'a', 'sample', 'sentence'];
const tokenizedSamples: string[][] = [
    ['this', 'is', 'a', 'positive', 'example'],
    ['this', 'is', 'a', 'negative', 'example']
];
const labels: number[] = [1, 0];  // Binary labels: 1 for positive, 0 for negative

//* Train and calibrate the classifier
//! Should split data into training/calibration sets; What split to use?
const classifier = trainClassifier(tokenizedSamples, labels);
const calibrationScores = calibrate(classifier, tokenizedSamples, labels);

//! What confidenceLevel to use? Should it be the same for each file?
const confidenceLevel = 0.95;
const prediction = predictWithConformal(classifier, tokenizedSnippet, calibrationScores, confidenceLevel);

Please let me know you thoughts on this and if my understanding is correct!

[prayas7102]

thanks @Commit2Cosmos for coming up with the solution, but i would like to tell you about the recent changes in source code, we've already unified the prediction code of the 4 .ts file into one. You can check this PR #16 for more info.

Coming to your proposal you can take a fresh pull, add yr changes in a new branch and raise a PR. Also check that yr changes doesn't compromise with previous log results. For this you can check the NodeJsSecurify.log file which gets generated. And then share your version of log file.

[Commit2Cosmos]

Thanks for pointing this out, was thinking of raising an issue about it myself! I've implemented the changes and it seems to work, but have a few questions:

• How would you like the confidence scores to be displayed? (Thinking e.g. "Code vulnerable to Brute force Attack in this file with _% certainty!!!")
• What proportion of the dataset should I use for calculating the confidence scores? (20%?)
• Would you like them to be displayed for detected vulnerabilities only or for negatives too?
• Is it possible to test how accurate these scores are? (i.e. separate tests or ..?)

[prayas7102]

Thanks for pointing this out, was thinking of raising an issue about it myself! I've implemented the changes and it seems to work, but have a few questions:

• How would you like the confidence scores to be displayed? (Thinking e.g. "Code vulnerable to Brute force Attack in this file with _% certainty!!!")
• What proportion of the dataset should I use for calculating the confidence scores? (20%?)
• Would you like them to be displayed for detected vulnerabilities only or for negatives too?
• Is it possible to test how accurate these scores are? (i.e. separate tests or ..?)

Hey @Commit2Cosmos,

I was reviewing your questions, and it struck me that we could prioritize addressing Issue #6 before moving on to the implementation. By doing so, we can develop a weighted model that will return both vulnerability predictions and an uncertainty metric.

This approach would simplify solving the issue, making the overall process smoother. Also, I assume you're using CSVs from the datasets directory—please confirm.

Let me know what you think!

Coming to you r questions, which are vital for both these issues,

1. Yes, preferred.
2. You can split data into 80% training and 20% testing.
3. Lets see for both cases how does it turns out for now.
4. Not Sure I may have to check it manually or from chatgpt etc.

[Commit2Cosmos]

Thank you! I agree that it makes sence to implement the ability to save and reuse models' weights before introducing new features like this one, as it could make testing the new model a bit more challenging. Hence, I'll have a look if I'm able to tackle that issue first before returning to this one

Also, I assume you're using CSVs from the datasets directory—please confirm.

I'm guessing the datasets in .ts come from those in .csv files? The parseCSV file isn't called anywhere it seems, could you clarify the relation between the two. I was simply going to use a part of the dataset that the classifier is already training on in Vulnerability.ts

[prayas7102]

I'm guessing the datasets in .ts come from those in .csv files? The parseCSV file isn't called anywhere it seems, could you clarify the relation between the two. I was simply going to use a part of the dataset that the classifier is already training on in Vulnerability.ts

Yes, the datasets in .ts and .csv files are the same.