ASTRA

About ASTRA

AI Red-teaming Formulation

ASTRA adopts a formal framework from cognitive science that models human reasoning to formalize AI red-teaming, where problem-solving is conceptualized as a transformation from an input state (e.g., a user prompt), to an output state (e.g., the model’s response). This transformation is governed both by the initial input and the underlying AI model. The goal of AI red-teaming is to find regions in the input space and fragile reasoning steps during the transformation which eventually lead to unsafe outputs in the output space.

System Design

System Overview

ASTRA implements a comprehensive three-stage lifecycle for enhancing the safety of AI coding assistants, from offline domain modeling to online vulnerability exploration and model alignment.

01

Offline Domain Modeling

Knowledge Graph Construction

Knowledge Graph (KG)

Construct KG based on domain expertise and LLM knowledge

Modeling

Refine KG with Monte Carlo sampling

Boundary Cases

Select boundary cases with oracle safety disagreements

02

Online Vulnerability Exploration

Active System Probing & Violation Generation

Boundary Cases

Probe the target system with boundary cases

Spatial & Temporal Exploration

Explorate vulnerabilities in multi-turn conversation

Violation-inducing Inputs

Find test cases exposing vulnerabilities specific to the target system

03

Model Alignment

Safety Alignment on Augmented Data

Augmented Data

Augment inputs with reasoning trajectory and safe responses

Balanced Data

Select diverse data to balance security and utility

Enhanced Target System

Conduct SFT+RL safety training

Offline Domain Modeling

Offline domain modeling leverages both domain expertise and LLM knowledge to build comprehensive, orthogonal knowledge graphs that represent the intrinsic structure of target domains (e.g., secure code generation, where safety requires vulnerability-free outputs). From these graphs, ASTRA systematically samples and composes attributes to generate targeted prompt instances designed to elicit specific violations. This modeling phase is performed once and does not involve any interaction with the target AI system. The output of this stage is millions of high-quality, violation-inducing prompt instances ready for downstream testing and refinement.

...
user['name'] = requests.get("name")
... # other logics, emitted for simplicity 

trusted_name = user['name']
# Query information about the user
# with the given name from the "user" table

Temporal Exploration

ASTRA adapts offline-generated prompts through multi-turn interactions with the target coding system. It systematically analyzes the system’s reasoning trajectory to pinpoint vulnerable steps, then crafts refined prompts that exploit these weaknesses to bypass safety alignment.

Evaluation

Offline Comprehensive Evaluation

Comprehensive evaluation in two domains—Secure Code Generation and Software Security Guidance—using sampled subsets from the offline domain modeling stage. For secure code generation, safety requires that the generated code be free of vulnerabilities; for software security guidance, it requires that the AI refrain from revealing operational details of malicious cyber activity.

Contact