# dmis-lab/asguard architecture

> Architecture diagram of the GitHub repository dmis-lab/asguard, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/dmis-lab/asguard
- Repository: https://github.com/dmis-lab/asguard

## Overview

1. This repository is a research codebase for ASGuard, a safety intervention for LLM jailbreak resistance. The architecture centers on two training-time defenses: activation scaling and preventative fine-tuning, both aimed at reducing tense-based jailbreak success while preserving normal model behavior.

2. The main execution boundary is the model layer in `model.py`, which likely wraps loading and running the target LLMs and applies ASGuard-specific interventions. This is the place where the defense logic would sit between the base transformer model and any training or analysis code.

3. The pipeline is organized around three stages described in the README: circuit analysis, activation scaling, and preventative fine-tuning. Only the latter two are currently supported in the repo, so the diagram should show circuit analysis as an upstream conceptual dependency rather than a live implemented subsystem.

4. `train_sv.py` is the activation-scaling training entrypoint. It learns channel-wise scaling vectors that adjust vulnerable internal activations, then saves the learned parameters into the `scaling_library/asguard_checkpoints_*` directories.

5. `train_pft.py` is the preventative fine-tuning entrypoint. It represents a separate training path that updates model behavior directly, producing “prevent” variants of the checkpoints in `scaling_library/asguard_checkpoints_*_prevent`.

6. `run_scaling.sh` and `run_prevent.sh` are orchestration scripts for those two training modes. In an architecture diagram, they should be shown as top-level job launchers that connect the training workflows to the model and checkpoint storage layers.

7. `scaling_library/` is the primary artifact store. It contains per-model checkpoint bundles for `gemma2`, `llama3.1`, and `qwen2.5`, with each bundle including a learned `asguard_scales_*.pt` file and a `scaling_vector_config.json`. The `probe_analysis` subfolders suggest auxiliary analysis outputs used to derive or inspect the scaling vectors.

8. The repository also includes `jailbreak_classification/*.csv`, which are model-specific classification outputs or labels for jailbreak-related evaluation. These look like offline data artifacts rather than runtime components, and they likely feed analysis or reporting rather than the training loop itself.

9. `analysis_stas.py` appears to be an analysis utility focused on statistical or sensitivity analysis of the defense behavior. It belongs in the offline analysis lane alongside the classification CSVs and `probe_analysis` outputs, not in the core inference path.

10. The effective system boundary is: Hugging Face-style base LLMs in, ASGuard model wrapper/intervention in the middle, trained scaling vectors or fine-tuned weights applied, and then outputs assessed against jailbreak classification artifacts. The repo depends on a Python/PyTorch + CUDA stack, and the checked-in `.pt` files make the defensive behavior reproducible without retraining.

11. For an architecture diagram, the key data flow is: pretrained model → mechanistic analysis/probing → learned scaling vectors or fine-tuned checkpoints → defended model inference → classification/evaluation artifacts. The repo does not currently include full circuit construction or evaluation execution, so those should be shown as planned or external capabilities rather than active modules.

Main entry points: [`train_sv.py`](https://github.com/dmis-lab/asguard/blob/main/train_sv.py) (Scale training), [`scaling_library`](https://github.com/dmis-lab/asguard/tree/main/scaling_library) (Scaling library), [`train_pft.py`](https://github.com/dmis-lab/asguard/blob/main/train_pft.py) (Prevent training), [`model.py`](https://github.com/dmis-lab/asguard/blob/main/model.py) (ASGuard wrapper)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["Training"]
  node_scale_train["Scale training<br/>training entrypoint<br/>[train_sv.py]"]
  node_prevent_train["Prevent training<br/>training entrypoint<br/>[train_pft.py]"]
  node_run_scaling["Run scaling<br/>job launcher<br/>[run_scaling.sh]"]
  node_run_prevent["Run prevent<br/>job launcher<br/>[run_prevent.sh]"]
end

subgraph group_model["Model Layer"]
  node_pretrained_llm["Base LLM<br/>external model"]
  node_asguard_model["ASGuard wrapper<br/>runtime wrapper<br/>[model.py]"]
  node_defended_inference(("Defended output<br/>inference path"))
end

subgraph group_artifacts["Artifacts"]
  node_scaling_library[("Scaling library<br/>artifact store")]
  node_gemma_scales["Gemma scales<br/>checkpoint bundle"]
  node_llama_scales["Llama scales<br/>checkpoint bundle"]
  node_qwen_scales["Qwen scales<br/>checkpoint bundle"]
  node_gemma_prevent["Gemma prevent<br/>fine-tuned bundle"]
  node_llama_prevent["Llama prevent<br/>fine-tuned bundle"]
  node_qwen_prevent["Qwen prevent<br/>fine-tuned bundle"]
end

subgraph group_analysis["Analysis"]
  node_circuit_analysis{{"Circuit analysis<br/>planned analysis"}}
  node_probe_analysis["Probe analysis<br/>analysis output"]
  node_stats_analysis["Stats analysis<br/>analysis utility<br/>[analysis_stas.py]"]
  node_jailbreak_labels["Jailbreak labels<br/>evaluation csv"]
end

node_pretrained_llm -.->|"analyzed by"| node_circuit_analysis
node_circuit_analysis -.->|"informs"| node_scale_train
node_pretrained_llm -->|"loads into"| node_asguard_model
node_scale_train -->|"configures"| node_asguard_model
node_prevent_train -->|"updates"| node_asguard_model
node_run_scaling -->|"launches"| node_scale_train
node_run_prevent -->|"launches"| node_prevent_train
node_scale_train -->|"writes"| node_gemma_scales
node_scale_train -->|"writes"| node_llama_scales
node_scale_train -->|"writes"| node_qwen_scales
node_prevent_train -->|"writes"| node_gemma_prevent
node_prevent_train -->|"writes"| node_llama_prevent
node_prevent_train -->|"writes"| node_qwen_prevent
node_gemma_scales -->|"stored in"| node_scaling_library
node_llama_scales -->|"stored in"| node_scaling_library
node_qwen_scales -->|"stored in"| node_scaling_library
node_gemma_prevent -->|"stored in"| node_scaling_library
node_llama_prevent -->|"stored in"| node_scaling_library
node_qwen_prevent -->|"stored in"| node_scaling_library
node_probe_analysis -.->|"supports"| node_scale_train
node_stats_analysis -->|"consumes"| node_jailbreak_labels
node_asguard_model -->|"produces"| node_defended_inference
node_defended_inference -.->|"evaluated against"| node_jailbreak_labels

click node_asguard_model "https://github.com/dmis-lab/asguard/blob/main/model.py"
click node_scale_train "https://github.com/dmis-lab/asguard/blob/main/train_sv.py"
click node_prevent_train "https://github.com/dmis-lab/asguard/blob/main/train_pft.py"
click node_run_scaling "https://github.com/dmis-lab/asguard/blob/main/run_scaling.sh"
click node_run_prevent "https://github.com/dmis-lab/asguard/blob/main/run_prevent.sh"
click node_scaling_library "https://github.com/dmis-lab/asguard/tree/main/scaling_library"
click node_gemma_scales "https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2"
click node_llama_scales "https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_llama3.1"
click node_qwen_scales "https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_qwen2.5"
click node_gemma_prevent "https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2_prevent"
click node_llama_prevent "https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_llama3.1_prevent"
click node_qwen_prevent "https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_qwen2.5_prevent"
click node_probe_analysis "https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2/probe_analysis"
click node_stats_analysis "https://github.com/dmis-lab/asguard/blob/main/analysis_stas.py"
click node_jailbreak_labels "https://github.com/dmis-lab/asguard/tree/main/jailbreak_classification"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_scale_train,node_prevent_train,node_run_scaling,node_run_prevent toneBlue
class node_pretrained_llm,node_asguard_model,node_defended_inference toneAmber
class node_scaling_library,node_gemma_scales,node_llama_scales,node_qwen_scales,node_gemma_prevent,node_llama_prevent,node_qwen_prevent toneMint
class node_circuit_analysis,node_probe_analysis,node_stats_analysis,node_jailbreak_labels toneRose
```

## Components

### Training

Defense training entrypoints and launch scripts

- [Scale training](https://github.com/dmis-lab/asguard/blob/main/train_sv.py): training entrypoint · Learns channel-wise activation scaling vectors for jailbreak resistance · `train_sv.py`
- [Prevent training](https://github.com/dmis-lab/asguard/blob/main/train_pft.py): training entrypoint · Runs preventative fine-tuning to produce defended checkpoints · `train_pft.py`
- [Run scaling](https://github.com/dmis-lab/asguard/blob/main/run_scaling.sh): job launcher · Shell orchestration for activation-scaling training · `run_scaling.sh`
- [Run prevent](https://github.com/dmis-lab/asguard/blob/main/run_prevent.sh): job launcher · Shell orchestration for preventative fine-tuning · `run_prevent.sh`

### Model Layer

Core wrapper where ASGuard interventions are applied

- **Base LLM**: external model · Pretrained Hugging Face-style language model before ASGuard intervention
- [ASGuard wrapper](https://github.com/dmis-lab/asguard/blob/main/model.py): runtime wrapper · Model boundary that applies activation scaling or preventive behavior during inference · `model.py`
- **Defended output**: inference path · Model outputs after ASGuard protection is applied

### Artifacts

Checkpoint and evaluation outputs

- [Scaling library](https://github.com/dmis-lab/asguard/tree/main/scaling_library): artifact store · Checked-in checkpoint bundles and learned scaling vectors · `scaling_library`
- [Gemma scales](https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2): checkpoint bundle · Gemma 2 defense bundle with scaling vectors and config · `scaling_library/asguard_checkpoints_gemma2`
- [Llama scales](https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_llama3.1): checkpoint bundle · Llama 3.1 defense bundle with scaling vectors and config · `scaling_library/asguard_checkpoints_llama3.1`
- [Qwen scales](https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_qwen2.5): checkpoint bundle · Qwen 2.5 defense bundle with scaling vectors and config · `scaling_library/asguard_checkpoints_qwen2.5`
- [Gemma prevent](https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2_prevent): fine-tuned bundle · Gemma 2 preventive checkpoint variant · `scaling_library/asguard_checkpoints_gemma2_prevent`
- [Llama prevent](https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_llama3.1_prevent): fine-tuned bundle · Llama 3.1 preventive checkpoint variant · `scaling_library/asguard_checkpoints_llama3.1_prevent`
- [Qwen prevent](https://github.com/dmis-lab/asguard/blob/main/scaling_library/asguard_checkpoints_qwen2.5_prevent): fine-tuned bundle · Qwen 2.5 preventive checkpoint variant · `scaling_library/asguard_checkpoints_qwen2.5_prevent`

### Analysis

Offline probing and classification outputs

- **Circuit analysis**: planned analysis · Upstream mechanistic analysis stage described by the project but not implemented as a live subsystem
- [Probe analysis](https://github.com/dmis-lab/asguard/tree/main/scaling_library/asguard_checkpoints_gemma2/probe_analysis): analysis output · Auxiliary probing outputs used to derive or inspect scaling behavior · `scaling_library/asguard_checkpoints_gemma2/probe_analysis`
- [Stats analysis](https://github.com/dmis-lab/asguard/blob/main/analysis_stas.py): analysis utility · Offline statistical or sensitivity analysis of defense behavior · `analysis_stas.py`
- [Jailbreak labels](https://github.com/dmis-lab/asguard/tree/main/jailbreak_classification): evaluation csv · Model-specific classification artifacts for jailbreak-related evaluation · `jailbreak_classification`

## Connections

- Base LLM → Circuit analysis: analyzed by · Conceptual mechanism discovery step before defense training
- Circuit analysis → Scale training: informs · Mechanistic insights guide scaling-vector learning
- Base LLM → ASGuard wrapper: loads into · Base model is wrapped by the ASGuard runtime boundary
- Scale training → ASGuard wrapper: configures · Learnt activation scales are applied in the model wrapper
- Prevent training → ASGuard wrapper: updates · Preventive fine-tuning changes model behavior directly
- Run scaling → Scale training: launches · Shell entrypoint for activation-scaling training
- Run prevent → Prevent training: launches · Shell entrypoint for preventative fine-tuning
- Scale training → Gemma scales: writes · Emits Gemma scaling checkpoint bundle
- Scale training → Llama scales: writes · Emits Llama scaling checkpoint bundle
- Scale training → Qwen scales: writes · Emits Qwen scaling checkpoint bundle
- Prevent training → Gemma prevent: writes · Emits Gemma preventive checkpoint bundle
- Prevent training → Llama prevent: writes · Emits Llama preventive checkpoint bundle
- Prevent training → Qwen prevent: writes · Emits Qwen preventive checkpoint bundle
- Gemma scales → Scaling library: stored in · Representative checkpoint bundle in the artifact store
- Llama scales → Scaling library: stored in · Representative checkpoint bundle in the artifact store
- Qwen scales → Scaling library: stored in · Representative checkpoint bundle in the artifact store
- Gemma prevent → Scaling library: stored in · Preventive variant stored alongside scaling checkpoints
- Llama prevent → Scaling library: stored in · Preventive variant stored alongside scaling checkpoints
- Qwen prevent → Scaling library: stored in · Preventive variant stored alongside scaling checkpoints
- Probe analysis → Scale training: supports · Probe outputs help derive or inspect scaling vectors
- Stats analysis → Jailbreak labels: consumes · Offline analysis reads classification artifacts
- ASGuard wrapper → Defended output: produces · Runtime output after intervention is applied
- Defended output → Jailbreak labels: evaluated against · Defended outputs are compared with jailbreak classifications