How dmis-lab/asguard works
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 Scale training · scaling_library Scaling library · train_pft.py Prevent training · model.py ASGuard wrapper
Components 18
Training
Defense training entrypoints and launch scripts
- Scale training · Learns channel-wise activation scaling vectors for jailbreak resistance
train_sv.py - Prevent training · Runs preventative fine-tuning to produce defended checkpoints
train_pft.py - Run scaling · Shell orchestration for activation-scaling training
run_scaling.sh - Run prevent · Shell orchestration for preventative fine-tuning
run_prevent.sh
Model Layer
Core wrapper where ASGuard interventions are applied
- Base LLM · Pretrained Hugging Face-style language model before ASGuard intervention
- ASGuard wrapper · Model boundary that applies activation scaling or preventive behavior during inference
model.py - Defended output · Model outputs after ASGuard protection is applied
Artifacts
Checkpoint and evaluation outputs
- Scaling library · Checked-in checkpoint bundles and learned scaling vectors
scaling_library - Gemma scales · Gemma 2 defense bundle with scaling vectors and config
scaling_library/asguard_checkpoints_gemma2 - Llama scales · Llama 3.1 defense bundle with scaling vectors and config
scaling_library/asguard_checkpoints_llama3.1 - Qwen scales · Qwen 2.5 defense bundle with scaling vectors and config
scaling_library/asguard_checkpoints_qwen2.5 - Gemma prevent · Gemma 2 preventive checkpoint variant
scaling_library/asguard_checkpoints_gemma2_prevent - Llama prevent · Llama 3.1 preventive checkpoint variant
scaling_library/asguard_checkpoints_llama3.1_prevent - Qwen prevent · Qwen 2.5 preventive checkpoint variant
scaling_library/asguard_checkpoints_qwen2.5_prevent
Analysis
Offline probing and classification outputs
- Circuit analysis · Upstream mechanistic analysis stage described by the project but not implemented as a live subsystem
- Probe analysis · Auxiliary probing outputs used to derive or inspect scaling behavior
scaling_library/asguard_checkpoints_gemma2/probe_analysis - Stats analysis · Offline statistical or sensitivity analysis of defense behavior
analysis_stas.py - Jailbreak labels · Model-specific classification artifacts for jailbreak-related evaluation
jailbreak_classification
Connections 23
- 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