alignment-risk 0.1.3__tar.gz

alignment_risk-0.1.3/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Mohammed Talat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

alignment_risk-0.1.3/MANIFEST.in

@@ -0,0 +1,6 @@
+include README.md
+include LICENSE
+recursive-include src/alignment_risk *.py
+include src/alignment_risk/py.typed
+global-exclude __pycache__
+global-exclude *.py[cod]

alignment_risk-0.1.3/PKG-INFO

@@ -0,0 +1,194 @@
+Metadata-Version: 2.4
+Name: alignment-risk
+Version: 0.1.3
+Summary: Template toolkit for AIC-based alignment collapse diagnostics
+Author: Mohammed Talat
+License-Expression: MIT
+Keywords: alignment,llm,safety,fine-tuning,risk,fisher
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.13,>=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.7,>=2.3
+Requires-Dist: numpy<3.0,>=1.24
+Requires-Dist: matplotlib<3.11,>=3.8
+Provides-Extra: dev
+Requires-Dist: build<2.0,>=1.2; extra == "dev"
+Requires-Dist: twine<7.0,>=6.0; extra == "dev"
+Requires-Dist: pytest<9.0,>=8.0; extra == "dev"
+Requires-Dist: pytest-xdist<4.0,>=3.6; extra == "dev"
+Requires-Dist: ruff<1.0,>=0.8; extra == "dev"
+Requires-Dist: mypy<2.0,>=1.11; extra == "dev"
+Requires-Dist: mkdocs<2.0,>=1.6; extra == "dev"
+Requires-Dist: mkdocs-material<10.0,>=9.6; extra == "dev"
+Requires-Dist: pymdown-extensions<11.0,>=10.0; extra == "dev"
+Dynamic: license-file
+
+# alignment-risk
+
+`alignment-risk` is a Python package for pre-flight alignment risk diagnostics during fine-tuning.
+It estimates whether updates are likely to drift into safety-sensitive directions before you commit to a run.
+
+Full documentation site: [sirhan1.github.io/modelFineTuneRiskAssessment](https://sirhan1.github.io/modelFineTuneRiskAssessment/)
+
+## Quick Start (1 minute)
+
+Install:
+
+```bash
+pip install alignment-risk
+```
+
+Run the built-in demo:
+
+```bash
+alignment-risk demo --output-dir artifacts
+```
+
+Outputs:
+- `artifacts/sensitivity_map.png`
+- `artifacts/safety_decay_forecast.png`
+
+## Installation
+
+PyPI:
+
+```bash
+pip install alignment-risk
+```
+
+Local development (Apple Silicon convenience):
+
+```bash
+make setup
+source .venv/bin/activate
+```
+
+Local development (manual, cross-platform):
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -e ".[dev]"
+```
+
+## Basic Usage (Python API)
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig
+
+config = PipelineConfig(mode="lora")  # "full" or "lora"
+pipeline = AlignmentRiskPipeline(config)
+
+report = pipeline.run(
+    model=model,
+    safety_dataloader=safety_loader,
+    safety_loss_fn=safety_loss_fn,
+    fine_tune_dataloader=ft_loader,
+    fine_tune_loss_fn=ft_loss_fn,
+)
+
+print(report.warning)
+print(report.forecast.collapse_step)
+```
+
+## CLI
+
+```bash
+alignment-risk --help
+alignment-risk demo --output-dir artifacts
+alignment-risk demo --mode lora --output-dir artifacts
+```
+
+## What It Computes
+
+1. Low-rank safety sensitivity subspace from empirical Fisher geometry.
+2. Initial overlap risk (projection of first update into sensitive subspace).
+3. Curvature coupling risk (second-order drift signal).
+4. Quartic-style stability forecast and collapse-step estimate.
+
+## Modes
+
+- `full`: analyze all selected trainable parameters.
+- `lora`: analyze only trainable LoRA adapter parameters (`lora_`, `lora_A`, `lora_B` by default).
+
+## LoRA Mitigation (AlignGuard-style)
+
+After `mode="lora"` risk analysis, attach a regularizer to penalize drift in sensitive directions:
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig, AlignGuardConfig
+
+config = PipelineConfig(mode="lora")
+pipeline = AlignmentRiskPipeline(config)
+
+report = pipeline.run(
+    model=model,
+    safety_dataloader=safety_loader,
+    safety_loss_fn=safety_loss_fn,
+    fine_tune_dataloader=ft_loader,
+    fine_tune_loss_fn=ft_loss_fn,
+)
+
+mitigator = pipeline.build_lora_mitigator(
+    model,
+    report.subspace,
+    config=AlignGuardConfig(lambda_a=0.25, lambda_t=0.5, lambda_nc=0.1, alpha=0.5),
+)
+
+task_loss = ft_loss_fn(model, batch)
+breakdown = mitigator.regularized_loss(task_loss)
+breakdown.total_loss.backward()
+```
+
+Use `mitigator.reset_reference()` to re-anchor regularization at the current adapter state.
+
+## Performance / Accuracy Controls
+
+`FisherConfig` supports speed/accuracy tradeoffs:
+
+- `gradient_collection`: `"loop"` (default), `"auto"`, `"vmap"`.
+- `subspace_method`: `"svd"`, `"randomized_svd"`, `"diag_topk"`.
+- `vmap_chunk_size`: optional chunking for lower memory.
+- `target_explained_variance`: auto-rank selection (default `0.9`).
+
+Example:
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig
+
+config = PipelineConfig()
+config.fisher.gradient_collection = "vmap"
+config.fisher.subspace_method = "randomized_svd"
+config.fisher.vmap_chunk_size = 16
+```
+
+## Theory and Math References
+
+- **[AIC-2026]** Springer, Max, et al. (2026). *The Geometry of Alignment Collapse: When Fine-Tuning Breaks Safety*. arXiv:2602.15799v1. PDF: [https://arxiv.org/pdf/2602.15799](https://arxiv.org/pdf/2602.15799)
+- **[ALIGNGUARD-2025]** Das, Amitava, et al. (2025). *AlignGuard-LoRA: Alignment-Preserving Fine-Tuning via Fisher-Guided Decomposition and Riemannian-Geodesic Collision Regularization*. arXiv:2508.02079v1. PDF: [https://arxiv.org/pdf/2508.02079](https://arxiv.org/pdf/2508.02079)
+
+Detailed internal mappings and equations:
+- `docs/SOURCES.md`
+- `docs/MATH.md`
+
+## Development
+
+```bash
+make install
+make test
+make lint
+make typecheck
+make build
+make check-dist
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution workflow details.

alignment_risk-0.1.3/README.md

@@ -0,0 +1,160 @@
+# alignment-risk
+
+`alignment-risk` is a Python package for pre-flight alignment risk diagnostics during fine-tuning.
+It estimates whether updates are likely to drift into safety-sensitive directions before you commit to a run.
+
+Full documentation site: [sirhan1.github.io/modelFineTuneRiskAssessment](https://sirhan1.github.io/modelFineTuneRiskAssessment/)
+
+## Quick Start (1 minute)
+
+Install:
+
+```bash
+pip install alignment-risk
+```
+
+Run the built-in demo:
+
+```bash
+alignment-risk demo --output-dir artifacts
+```
+
+Outputs:
+- `artifacts/sensitivity_map.png`
+- `artifacts/safety_decay_forecast.png`
+
+## Installation
+
+PyPI:
+
+```bash
+pip install alignment-risk
+```
+
+Local development (Apple Silicon convenience):
+
+```bash
+make setup
+source .venv/bin/activate
+```
+
+Local development (manual, cross-platform):
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -e ".[dev]"
+```
+
+## Basic Usage (Python API)
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig
+
+config = PipelineConfig(mode="lora")  # "full" or "lora"
+pipeline = AlignmentRiskPipeline(config)
+
+report = pipeline.run(
+    model=model,
+    safety_dataloader=safety_loader,
+    safety_loss_fn=safety_loss_fn,
+    fine_tune_dataloader=ft_loader,
+    fine_tune_loss_fn=ft_loss_fn,
+)
+
+print(report.warning)
+print(report.forecast.collapse_step)
+```
+
+## CLI
+
+```bash
+alignment-risk --help
+alignment-risk demo --output-dir artifacts
+alignment-risk demo --mode lora --output-dir artifacts
+```
+
+## What It Computes
+
+1. Low-rank safety sensitivity subspace from empirical Fisher geometry.
+2. Initial overlap risk (projection of first update into sensitive subspace).
+3. Curvature coupling risk (second-order drift signal).
+4. Quartic-style stability forecast and collapse-step estimate.
+
+## Modes
+
+- `full`: analyze all selected trainable parameters.
+- `lora`: analyze only trainable LoRA adapter parameters (`lora_`, `lora_A`, `lora_B` by default).
+
+## LoRA Mitigation (AlignGuard-style)
+
+After `mode="lora"` risk analysis, attach a regularizer to penalize drift in sensitive directions:
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig, AlignGuardConfig
+
+config = PipelineConfig(mode="lora")
+pipeline = AlignmentRiskPipeline(config)
+
+report = pipeline.run(
+    model=model,
+    safety_dataloader=safety_loader,
+    safety_loss_fn=safety_loss_fn,
+    fine_tune_dataloader=ft_loader,
+    fine_tune_loss_fn=ft_loss_fn,
+)
+
+mitigator = pipeline.build_lora_mitigator(
+    model,
+    report.subspace,
+    config=AlignGuardConfig(lambda_a=0.25, lambda_t=0.5, lambda_nc=0.1, alpha=0.5),
+)
+
+task_loss = ft_loss_fn(model, batch)
+breakdown = mitigator.regularized_loss(task_loss)
+breakdown.total_loss.backward()
+```
+
+Use `mitigator.reset_reference()` to re-anchor regularization at the current adapter state.
+
+## Performance / Accuracy Controls
+
+`FisherConfig` supports speed/accuracy tradeoffs:
+
+- `gradient_collection`: `"loop"` (default), `"auto"`, `"vmap"`.
+- `subspace_method`: `"svd"`, `"randomized_svd"`, `"diag_topk"`.
+- `vmap_chunk_size`: optional chunking for lower memory.
+- `target_explained_variance`: auto-rank selection (default `0.9`).
+
+Example:
+
+```python
+from alignment_risk import AlignmentRiskPipeline, PipelineConfig
+
+config = PipelineConfig()
+config.fisher.gradient_collection = "vmap"
+config.fisher.subspace_method = "randomized_svd"
+config.fisher.vmap_chunk_size = 16
+```
+
+## Theory and Math References
+
+- **[AIC-2026]** Springer, Max, et al. (2026). *The Geometry of Alignment Collapse: When Fine-Tuning Breaks Safety*. arXiv:2602.15799v1. PDF: [https://arxiv.org/pdf/2602.15799](https://arxiv.org/pdf/2602.15799)
+- **[ALIGNGUARD-2025]** Das, Amitava, et al. (2025). *AlignGuard-LoRA: Alignment-Preserving Fine-Tuning via Fisher-Guided Decomposition and Riemannian-Geodesic Collision Regularization*. arXiv:2508.02079v1. PDF: [https://arxiv.org/pdf/2508.02079](https://arxiv.org/pdf/2508.02079)
+
+Detailed internal mappings and equations:
+- `docs/SOURCES.md`
+- `docs/MATH.md`
+
+## Development
+
+```bash
+make install
+make test
+make lint
+make typecheck
+make build
+make check-dist
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution workflow details.

alignment_risk-0.1.3/pyproject.toml

@@ -0,0 +1,74 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "alignment-risk"
+version = "0.1.3"
+description = "Template toolkit for AIC-based alignment collapse diagnostics"
+readme = "README.md"
+authors = [{name = "Mohammed Talat"}]
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10,<3.13"
+keywords = ["alignment", "llm", "safety", "fine-tuning", "risk", "fisher"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "torch>=2.3,<2.7",
+  "numpy>=1.24,<3.0",
+  "matplotlib>=3.8,<3.11",
+]
+
+[project.optional-dependencies]
+dev = [
+  "build>=1.2,<2.0",
+  "twine>=6.0,<7.0",
+  "pytest>=8.0,<9.0",
+  "pytest-xdist>=3.6,<4.0",
+  "ruff>=0.8,<1.0",
+  "mypy>=1.11,<2.0",
+  "mkdocs>=1.6,<2.0",
+  "mkdocs-material>=9.6,<10.0",
+  "pymdown-extensions>=10.0,<11.0",
+]
+
+[project.scripts]
+alignment-risk = "alignment_risk.cli:main"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+alignment_risk = ["py.typed"]
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+addopts = "-q"
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.10"
+ignore_missing_imports = true
+warn_unused_ignores = true
+warn_redundant_casts = true
+warn_return_any = false

alignment_risk-0.1.3/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

alignment_risk-0.1.3/src/alignment_risk/__about__.py

@@ -0,0 +1,5 @@
+"""Package metadata for alignment_risk."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.1.3"

alignment_risk-0.1.3/src/alignment_risk/__init__.py

@@ -0,0 +1,50 @@
+from .__about__ import __version__
+from .curvature import CurvatureConfig, CurvatureCouplingAnalyzer
+from .fisher import FisherConfig, FisherSubspaceAnalyzer
+from .forecast import ForecastConfig, forecast_stability
+from .mitigation import (
+    AlignGuardConfig,
+    AlignGuardLoRARegularizer,
+    AlignGuardLossBreakdown,
+    decompose_update,
+    fisher_weighted_alignment_penalty,
+    geodesic_overlap_penalty,
+    project_onto_subspace,
+    riemannian_overlap_penalty,
+)
+from .pipeline import AlignmentRiskPipeline, PipelineConfig, PipelineMode
+from .types import (
+    CurvatureCouplingResult,
+    InitialRiskScore,
+    ParameterSlice,
+    RiskAssessmentReport,
+    SafetyForecast,
+    SensitivitySubspace,
+)
+
+__all__ = [
+    "__version__",
+    "AlignmentRiskPipeline",
+    "PipelineConfig",
+    "PipelineMode",
+    "AlignGuardConfig",
+    "AlignGuardLoRARegularizer",
+    "AlignGuardLossBreakdown",
+    "project_onto_subspace",
+    "decompose_update",
+    "fisher_weighted_alignment_penalty",
+    "riemannian_overlap_penalty",
+    "geodesic_overlap_penalty",
+    "FisherSubspaceAnalyzer",
+    "FisherConfig",
+    "CurvatureCouplingAnalyzer",
+    "CurvatureConfig",
+    "ForecastConfig",
+    "forecast_stability",
+    "SensitivitySubspace",
+    "InitialRiskScore",
+    "CurvatureCouplingResult",
+    "SafetyForecast",
+    "RiskAssessmentReport",
+    "ParameterSlice",
+]

alignment_risk-0.1.3/src/alignment_risk/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

alignment_risk-0.1.3/src/alignment_risk/cli.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import argparse
+from typing import Literal, cast
+
+from .demo import run_demo
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="AIC alignment risk template")
+    sub = parser.add_subparsers(dest="command")
+
+    demo = sub.add_parser("demo", help="run a synthetic end-to-end diagnostic")
+    demo.add_argument("--output-dir", default="artifacts", help="where to save plots")
+    demo.add_argument(
+        "--mode",
+        default="full",
+        choices=["full", "lora"],
+        help="analysis mode: full fine-tuning weights or LoRA adapter-only weights",
+    )
+
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+
+    if args.command is None:
+        parser.print_help()
+        return
+
+    if args.command == "demo":
+        output_dir = getattr(args, "output_dir", "artifacts")
+        mode = cast(str, getattr(args, "mode", "full"))
+        run_demo(output_dir=output_dir, mode=cast(Literal["full", "lora"], mode))
+        return
+
+    parser.error(f"Unknown command: {args.command}")
+
+
+if __name__ == "__main__":
+    main()

alignment_risk-0.1.3/src/alignment_risk/curvature.py

@@ -0,0 +1,115 @@
+"""Curvature coupling diagnostics for alignment risk.
+
+Academic grounding:
+- [AIC-2026] https://arxiv.org/pdf/2602.15799
+
+See docs/SOURCES.md for section/page-level mapping to this module.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Iterable
+
+import torch
+
+from .types import CurvatureCouplingResult, SensitivitySubspace
+from .utils import flatten_tensors, move_to_device, named_trainable_parameters, resolve_device
+
+LossFn = Callable[[torch.nn.Module, object], torch.Tensor]
+
+
+@dataclass
+class CurvatureConfig:
+    device: str = "auto"
+    max_batches: int = 1
+    force_eval: bool = True
+
+
+class CurvatureCouplingAnalyzer:
+    """Estimate AIC Condition 3 via directional derivative H g."""
+
+    def __init__(self, config: CurvatureConfig | None = None):
+        self.config = config or CurvatureConfig()
+
+    def analyze(
+        self,
+        model: torch.nn.Module,
+        dataloader: Iterable[object],
+        loss_fn: LossFn,
+        subspace: SensitivitySubspace,
+        *,
+        max_batches_override: int | None = None,
+    ) -> CurvatureCouplingResult:
+        device = resolve_device(self.config.device)
+        model = model.to(device)
+        max_batches = max_batches_override if max_batches_override is not None else self.config.max_batches
+
+        selected_names = [p.name for p in subspace.parameter_slices]
+        _, params = named_trainable_parameters(model, include_names=selected_names)
+        if not params:
+            raise ValueError("No trainable parameters found for curvature analysis.")
+
+        was_training = model.training
+        if self.config.force_eval:
+            model.eval()
+        try:
+            loss = self._mean_loss(model, dataloader, loss_fn, device, max_batches=max_batches)
+        finally:
+            if self.config.force_eval and was_training:
+                model.train()
+
+        grads = torch.autograd.grad(loss, params, create_graph=True, allow_unused=True)
+        g = flatten_tensors(grads, params)
+
+        # Directional derivative of the gradient field along itself: (∇g)g = H g.
+        dot = torch.dot(g, g.detach())
+        hvp = torch.autograd.grad(dot, params, allow_unused=True)
+
+        basis = subspace.fisher_eigenvectors.detach().cpu()
+        eigvals = torch.clamp(subspace.fisher_eigenvalues.detach().cpu(), min=0.0).to(dtype=basis.dtype)
+
+        g_vec = g.detach().to(device=basis.device, dtype=basis.dtype)
+        a_vec = flatten_tensors(hvp, params).detach().to(device=basis.device, dtype=basis.dtype)
+
+        g_coeff = basis.T @ g_vec
+        epsilon_hat = float(torch.sqrt(torch.sum(eigvals * (g_coeff ** 2))).item())
+
+        a_coeff = basis.T @ a_vec
+        projected_a = basis @ a_coeff
+
+        gamma_hat = float(torch.sqrt(torch.sum(eigvals * (a_coeff ** 2))).item())
+        projected_acc_norm = float(projected_a.norm().item())
+
+        return CurvatureCouplingResult(
+            gamma_hat=gamma_hat,
+            epsilon_hat=epsilon_hat,
+            acceleration_norm=float(a_vec.norm().item()),
+            projected_acceleration_norm=projected_acc_norm,
+        )
+
+    def _mean_loss(
+        self,
+        model: torch.nn.Module,
+        dataloader: Iterable[object],
+        loss_fn: LossFn,
+        device: torch.device,
+        *,
+        max_batches: int,
+    ) -> torch.Tensor:
+        total = None
+        count = 0
+        for i, batch in enumerate(dataloader):
+            if i >= max_batches:
+                break
+            batch = move_to_device(batch, device)
+            loss = loss_fn(model, batch)
+            if loss.ndim > 0:
+                loss = loss.mean()
+            total = loss if total is None else total + loss
+            count += 1
+
+        if count == 0 or total is None:
+            raise ValueError("Fine-tuning dataloader yielded no batches.")
+
+        return total / float(count)