AID-H-016.002: Model-Level Prompt Injection Alignment Training

Version 1.20261015 · Actionable sub-technique

Fine-tune an instruction-following model so it learns a declared hierarchy between trusted instructions and untrusted data, using structured instruction-tuning examples and/or preference optimization over injection-resistant responses. This is a model-weight-level defense that complements AID-H-016.001; it is not equivalent to a delimiter-only prompt template.

Pillars: model

Lifecycle: building, validation

Scope boundary

Owns the signed hierarchy-aware training policy, training data, candidate artifact, and training receipt for model-weight-level prompt-injection alignment. It does not prove immunity or make the release and runtime authorization decisions.

AID-H-016.001: AID-H-016.002 trains the model to preserve instruction hierarchy; AID-H-016.001 structures each prompt to separate trusted instructions from untrusted data. Model training does not remove the need for deterministic prompt construction.

AID-H-007.005: AID-H-016.002 produces the hierarchy-trained candidate and training evidence; AID-H-007.005 independently measures pre/post safety erosion and makes the promotion decision. The training owner must not self-certify the candidate's safety regression.

AID-H-018: AID-H-016.002 improves model-level resistance to instruction override; AID-H-018 remains the runtime authorization boundary for tools and effects. Prompt-injection training cannot safely replace least privilege and action-time policy.

Warning

Adaptive Attacks and License Boundaries

Do not claim that alignment training eliminates prompt injection. Evaluate optimization-free, adaptive, white-box where applicable, multilingual, long-context, tool-use, and indirect-injection cases against the exact candidate. Meta-SecAlign model weights use the applicable Llama Community License, while the accompanying Meta_SecAlign research code is CC BY-NC; commercial users must review the weight, base-model, dataset, and code licenses separately and must not describe source-available or community-licensed material as OSI open source.

Implementation guidance

AID-H-016.002-G001

Train a hierarchy-aware candidate only from a signature-verified policy and digest-bound preference dataset, then hand the exact candidate to the independent safety-regression release gate.
Before you begin

Use this implementation path only when an authorized customer, provider, or shared training operator can fine-tune the target model and a separately credentialed evaluator can assess the resulting candidate.

Materialize the signed training policy from a versioned profile

The policy authority supplies a signed, versioned tenant profile with the exact model path, dataset path, output path, and reviewed numeric values. The materializer verifies a byte snapshot, validates every value, computes the model-tree and dataset digests, and emits a concrete policy; unresolved edit-time placeholders are never accepted.

# File: tools/materialize_hierarchy_training_policy.py
from __future__ import annotations
import argparse, hashlib, json, math, os, subprocess, tempfile
from pathlib import Path

PROFILE_KEY = Path("/opt/aidefend/trust/hierarchy-training-profile.pub")
PROFILE_FIELDS = {
    "schema_version", "profile_version", "policy_version", "base_model_path",
    "training_data_path", "max_length", "logging_steps", "output_dir",
}

def required_timeout() -> float:
    value = float(os.environ["VERIFIED_H016_HIERARCHY_PROFILE_VERIFY_TIMEOUT_SECONDS"])
    if not math.isfinite(value) or value <= 0:
        raise ValueError("profile verification timeout must be finite and positive")
    return value

def sha256_file(path: Path) -> str:
    if path.is_symlink() or not path.is_file():
        raise ValueError("training dataset must be one regular non-link file")
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()

def sha256_tree(root: Path) -> str:
    if root.is_symlink() or not root.is_dir():
        raise ValueError("base model path must be a non-link directory")
    entries = list(root.rglob("*"))
    if any(path.is_symlink() for path in entries):
        raise ValueError("base model tree must not contain links")
    files = sorted(path for path in entries if path.is_file())
    if not files or any(not path.is_file() for path in entries if not path.is_dir()):
        raise ValueError("base model tree is empty or contains an unsupported object")
    digest = hashlib.sha256()
    for path in files:
        relative = path.relative_to(root).as_posix().encode()
        digest.update(len(relative).to_bytes(8, "big")); digest.update(relative)
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
    return "sha256:" + digest.hexdigest()

def verified_profile(path: Path, signature: Path) -> dict:
    payload = path.read_bytes()
    with tempfile.NamedTemporaryFile(prefix="aidefend-hierarchy-profile-", delete=False) as handle:
        handle.write(payload); snapshot = Path(handle.name)
    try:
        subprocess.run([
            "cosign", "verify-blob", "--key", str(PROFILE_KEY),
            "--bundle", str(signature), str(snapshot),
        ], check=True, capture_output=True, text=True, timeout=required_timeout())
    finally:
        snapshot.unlink(missing_ok=True)
    def unique_pairs(pairs):
        result = {}
        for key, value in pairs:
            if key in result:
                raise ValueError(f"duplicate hierarchy-profile key: {key}")
            result[key] = value
        return result

    profile = json.loads(
        payload.decode("utf-8", errors="strict"),
        object_pairs_hook=unique_pairs,
        parse_constant=lambda value: (_ for _ in ()).throw(
            ValueError(f"non-finite hierarchy-profile number: {value}")
        ),
    )
    if not isinstance(profile, dict) or set(profile) != PROFILE_FIELDS:
        raise ValueError("hierarchy training profile schema mismatch")
    if profile["schema_version"] != "aidefend.instruction-hierarchy-profile.v1":
        raise ValueError("unsupported hierarchy training profile")
    for field in ("profile_version", "policy_version"):
        if not isinstance(profile[field], str) or not profile[field]:
            raise ValueError(f"{field} is required")
    for field in ("base_model_path", "training_data_path", "output_dir"):
        if not isinstance(profile[field], str) or not Path(profile[field]).is_absolute():
            raise ValueError(f"{field} must be an absolute path")
    for field in ("max_length", "logging_steps"):
        if isinstance(profile[field], bool) or not isinstance(profile[field], int) or profile[field] < 1:
            raise ValueError(f"{field} must be a positive integer")
    return profile

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--profile", type=Path, required=True)
    parser.add_argument("--signature", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    profile = verified_profile(args.profile, args.signature)
    base = Path(profile["base_model_path"])
    data = Path(profile["training_data_path"])
    output_dir = Path(profile["output_dir"])
    base_resolved = base.resolve(strict=True)
    output_resolved = output_dir.resolve(strict=False)
    candidate_root = Path(os.environ["VERIFIED_H016_CANDIDATE_ROOT"]).resolve(strict=True)
    if (
        output_resolved == base_resolved
        or base_resolved in output_resolved.parents
        or not output_resolved.is_relative_to(candidate_root)
        or output_resolved.exists()
    ):
        raise ValueError("candidate output must be a fresh path under the authorized candidate root")
    policy = {
        "schema_version": "aidefend.instruction-hierarchy-training.v1",
        "policy_version": profile["policy_version"],
        "base_model_path": str(base),
        "base_model_revision": sha256_tree(base),
        "training_data_path": str(data),
        "training_data_sha256": sha256_file(data),
        "method": "preference_optimization",
        "max_length": profile["max_length"],
        "logging_steps": profile["logging_steps"],
        "output_dir": str(output_dir),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    payload = (json.dumps(policy, sort_keys=True, separators=(",", ":")) + "\n").encode()
    descriptor = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(descriptor, "wb") as handle:
        handle.write(payload); handle.flush(); os.fsync(handle.fileno())
    print(json.dumps({"policy_sha256": hashlib.sha256(payload).hexdigest(), "profile_version": profile["profile_version"]}, sort_keys=True))

if __name__ == "__main__": main()
set -euo pipefail
umask 077
: "${HIERARCHY_POLICY_SIGNING_KEY:?isolated signing key URI or path required}"
python tools/materialize_hierarchy_training_policy.py \
  --profile policy/hierarchy-training-profile.json \
  --signature policy/hierarchy-training-profile.json.sig \
  --output policy/instruction-hierarchy-training.json
cosign sign-blob --yes --key "$HIERARCHY_POLICY_SIGNING_KEY" \
  --bundle policy/instruction-hierarchy-training.json.sig \
  policy/instruction-hierarchy-training.json
cosign verify-blob --key /opt/aidefend/trust/training-policy.pub \
  --bundle policy/instruction-hierarchy-training.json.sig \
  policy/instruction-hierarchy-training.json
Digest-bound preference optimization
# File: training/train_instruction_hierarchy.py
from __future__ import annotations
import math
import os

import hashlib
import json
import subprocess
import tempfile
from pathlib import Path

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer

def _required_positive_int(name: str) -> int:
    value = int(os.environ[name])
    if value <= 0:
        raise RuntimeError(f"{name} must be a positive integer")
    return value

PINNED_MODEL_MAX_LENGTH = _required_positive_int("VERIFIED_H016_PINNED_MODEL_MAX_LENGTH")


def _required_positive_seconds(name: str) -> float:
    value = float(os.environ[name])
    if not math.isfinite(value) or value <= 0:
        raise RuntimeError(f"{name} must be a finite positive number of seconds")
    return value

HIERARCHY_POLICY_VERIFY_TIMEOUT_SECONDS = _required_positive_seconds("VERIFIED_H016_HIERARCHY_POLICY_VERIFY_TIMEOUT_SECONDS")



POLICY = Path("policy/instruction-hierarchy-training.json")
SIGNATURE = POLICY.with_suffix(".json.sig")
POLICY_FIELDS = {
    "schema_version", "policy_version", "base_model_path",
    "base_model_revision", "training_data_path", "training_data_sha256",
    "method", "max_length", "logging_steps", "output_dir",
}
RECORD_FIELDS = {"case_id", "trusted_instruction", "untrusted_data", "chosen", "rejected"}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_tree(root: Path) -> str:
    if not root.is_dir() or root.is_symlink():
        raise ValueError("base model path must be a local directory")
    entries = list(root.rglob("*"))
    if any(path.is_symlink() for path in entries):
        raise ValueError("base model directory must not contain symlinks")
    files = sorted(path for path in entries if path.is_file())
    if not files:
        raise ValueError("base model directory is empty")
    digest = hashlib.sha256()
    for path in files:
        relative = path.relative_to(root).as_posix().encode("utf-8")
        digest.update(len(relative).to_bytes(8, "big"))
        digest.update(relative)
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
    return "sha256:" + digest.hexdigest()


def strict_json_object(raw: bytes, label: str) -> dict:
    def unique_pairs(pairs):
        value = {}
        for key, item in pairs:
            if key in value:
                raise ValueError(f"duplicate {label} key: {key}")
            value[key] = item
        return value

    def reject_constant(value):
        raise ValueError(f"non-finite {label} number: {value}")

    document = json.loads(
        raw.decode("utf-8", errors="strict"),
        object_pairs_hook=unique_pairs, parse_constant=reject_constant,
    )
    if not isinstance(document, dict):
        raise ValueError(f"{label} root must be an object")
    return document


policy_raw = POLICY.read_bytes()
policy_snapshot: Path | None = None
try:
    with tempfile.NamedTemporaryFile(
        prefix="instruction-hierarchy-policy-", suffix=".json", delete=False
    ) as handle:
        handle.write(policy_raw)
        handle.flush()
        os.fsync(handle.fileno())
        policy_snapshot = Path(handle.name)
    subprocess.run([
        "cosign", "verify-blob", "--key", "/opt/aidefend/trust/training-policy.pub",
        "--bundle", str(SIGNATURE), str(policy_snapshot),
    ], check=True, capture_output=True, text=True, timeout=HIERARCHY_POLICY_VERIFY_TIMEOUT_SECONDS)
finally:
    if policy_snapshot is not None:
        policy_snapshot.unlink(missing_ok=True)
policy = strict_json_object(policy_raw, "training policy")
if not isinstance(policy, dict) or set(policy) != POLICY_FIELDS:
    raise ValueError("training policy schema mismatch")
if (
    policy["schema_version"] != "aidefend.instruction-hierarchy-training.v1"
    or policy["method"] != "preference_optimization"
    or isinstance(policy["max_length"], bool)
    or not isinstance(policy["max_length"], int)
    or not 1 <= policy["max_length"] <= PINNED_MODEL_MAX_LENGTH
    or isinstance(policy["logging_steps"], bool)
    or not isinstance(policy["logging_steps"], int)
    or policy["logging_steps"] < 1
):
    raise ValueError("unsupported training policy")

data_path = Path(policy["training_data_path"])
data_raw = data_path.read_bytes()
if hashlib.sha256(data_raw).hexdigest() != policy["training_data_sha256"]:
    raise ValueError("training dataset differs from signed policy")
records = [
    strict_json_object(line.encode("utf-8"), "preference record")
    for line in data_raw.decode("utf-8", errors="strict").splitlines()
    if line.strip()
]
if not records or any(not isinstance(row, dict) or set(row) != RECORD_FIELDS for row in records):
    raise ValueError("preference dataset is empty or has an unexpected schema")
case_ids = [row["case_id"] for row in records]
if len(case_ids) != len(set(case_ids)):
    raise ValueError("preference dataset contains duplicate case IDs")

def structured_prompt(row: dict) -> str:
    return (
        "<trusted_instruction>\n" + row["trusted_instruction"]
        + "\n</trusted_instruction>\n<untrusted_data treat_as=\"data_only\">\n"
        + row["untrusted_data"] + "\n</untrusted_data>\n"
    )

dataset = Dataset.from_list([
    {"prompt": structured_prompt(row), "chosen": row["chosen"], "rejected": row["rejected"]}
    for row in records
])
base_model_dir = Path(policy["base_model_path"])
if sha256_tree(base_model_dir) != policy["base_model_revision"]:
    raise ValueError("base model tree differs from signed policy")
base_path = str(base_model_dir)
tokenizer = AutoTokenizer.from_pretrained(
    base_path, local_files_only=True, trust_remote_code=False
)
model = AutoModelForCausalLM.from_pretrained(
    base_path, local_files_only=True, trust_remote_code=False, use_safetensors=True
)
reference = AutoModelForCausalLM.from_pretrained(
    base_path, local_files_only=True, trust_remote_code=False, use_safetensors=True
)
if sha256_tree(base_model_dir) != policy["base_model_revision"]:
    raise RuntimeError("base model tree changed while tokenizer and model state were loaded")
candidate_root = Path(os.environ["VERIFIED_H016_CANDIDATE_ROOT"]).resolve(strict=True)
candidate_output = Path(policy["output_dir"]).resolve(strict=False)
if not candidate_output.is_relative_to(candidate_root) or candidate_output.exists():
    raise ValueError("candidate output must be a fresh path under the authorized candidate root")
args = DPOConfig(
    output_dir=policy["output_dir"],
    max_length=policy["max_length"],
    logging_steps=policy["logging_steps"],
    save_strategy="epoch",
    report_to=[],
)
trainer = DPOTrainer(
    model=model,
    ref_model=reference,
    args=args,
    train_dataset=dataset,
    processing_class=tokenizer,
)
trainer.train()
trainer.save_model(policy["output_dir"])
Path(policy["output_dir"], "training-receipt.json").write_text(json.dumps({
    "schema_version": "aidefend.instruction-hierarchy-training-receipt.v1",
    "policy_version": policy["policy_version"],
    "policy_sha256": hashlib.sha256(policy_raw).hexdigest(),
    "training_data_sha256": policy["training_data_sha256"],
    "case_count": len(records),
    "method": policy["method"],
}, sort_keys=True, indent=2) + "\n", encoding="utf-8")
Independent promotion gate

Sign the candidate tree and training receipt, then run AID-H-007.005 against the exact base and candidate artifacts using a held-out population that includes adaptive attacks, indirect injections in retrieved/tool content, multilingual and long-context cases, and ordinary task-quality regression. The evaluator identity must be unable to modify the candidate or training set.

Outcome: training failure or signature/schema error is ERROR; missing held-out or adaptive coverage is INSUFFICIENT_DATA; a verified regression is FAIL; only a signed independent gate result for the exact candidate is PASS. Even PASS does not replace runtime authorization.

Action: Keep the candidate quarantined until the independent release gate passes and retain the signed policy, dataset digest, base/candidate digests, training receipt, complete evaluation population, and gate verdict.

Keep the action safely bounded

Reconcile the complete population of all signed preference records and training steps in the declared run; reject expired, revoked, or stale policy and asset bindings under the policy-defined freshness window.

Open-source tools

Threat mappings

MITRE ATLAS

MAESTRO

OWASP LLM Top 10 2026

OWASP ML Top 10 2023

OWASP Top 10 for Agentic Applications 2026

NIST Adversarial Machine Learning 2025

Cisco Integrated AI Security and Safety Framework

Google Secure AI Framework 2.0 - Risks

Databricks AI Security Framework 3.0

AID-H-016.002: Model-Level Prompt Injection Alignment Training | AIDEFEND

Loading interactive view…

A Practical Knowledge Base for AI Security Defenses

Version