Examples

Examples of integrating ARC with different architectures. The minimal example runs as-is after pip install arc-training. The others show the integration pattern, so swap in your own data loaders and loss functions.

Pattern

Every example follows one pattern. Create a Controller, train as normal, call controller.step(loss) after backward(), and skip the optimizer step when action.should_skip is true.

Minimal Example

minimal.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import arc

model = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters())
controller = arc.Controller(model, optimizer)

for step in range(1000):
    x = torch.randn(32, 10)
    y = torch.randint(0, 5, (32,))
    optimizer.zero_grad()
    loss = F.cross_entropy(model(x), y)
    loss.backward()

    action = controller.step(loss)
    if action.should_skip:
        continue
    optimizer.step()

controller.summary()

Image Classification

cnn_example.py

import arc

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        return self.fc(x.view(x.size(0), -1))

model = SimpleCNN()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
controller = arc.Controller(model, optimizer, scheduler=scheduler)

for epoch in range(50):
    for images, labels in train_loader:
        optimizer.zero_grad()
        loss = F.cross_entropy(model(images), labels)
        loss.backward()
        if controller.step(loss).should_skip:
            continue
        optimizer.step()
    scheduler.step()

Object Detection

yolo_example.py

class SimpleYOLO(nn.Module):
    def __init__(self, num_classes=80):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.head = nn.Conv2d(64, num_classes + 5, 1)

    def forward(self, x):
        return self.head(self.backbone(x))

model = SimpleYOLO()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
controller = arc.Controller(model, optimizer)

for epoch in range(300):
    for images, targets in train_loader:
        optimizer.zero_grad()
        loss = compute_yolo_loss(model(images), targets)
        loss.backward()
        if controller.step(loss).should_skip:
            continue
        optimizer.step()

Transformer / NLP

transformer_example.py

class TextClassifier(nn.Module):
    def __init__(self, vocab=10000, d_model=256, nhead=4):
        super().__init__()
        self.emb = nn.Embedding(vocab, d_model)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead, batch_first=True),
            num_layers=2
        )
        self.fc = nn.Linear(d_model, 2)

    def forward(self, x):
        return self.fc(self.encoder(self.emb(x)).mean(dim=1))

model = TextClassifier()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
controller = arc.Controller(model, optimizer)

for epoch in range(10):
    for tokens, labels in dataloader:
        optimizer.zero_grad()
        loss = F.cross_entropy(model(tokens), labels)
        loss.backward()
        if controller.step(loss).should_skip:
            continue
        optimizer.step()

LLM Fine-Tuning (Scheduler + Gradient Accumulation)

llm_finetune.py

from arc import Controller, ControllerConfig

ACCUM = 8
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0,
                                              end_factor=0.0, total_iters=10_000)
controller = Controller(model, optimizer, scheduler=scheduler,
                        config=ControllerConfig(gradient_accumulation_steps=ACCUM))

for i, batch in enumerate(loader):
    loss = model(**batch).loss / ACCUM
    loss.backward()
    action = controller.step(loss)       # full checks run on every ACCUM-th call

    if (i + 1) % ACCUM:
        continue                          # still accumulating
    if action.unrecoverable:
        raise RuntimeError(action.message)
    if not action.should_skip:
        optimizer.step()
        scheduler.step()
    optimizer.zero_grad()

When ARC cuts the learning rate after a rollback, it also scales the scheduler's base_lrs, so the schedule continues from the lower rate instead of overriding it. After 500 healthy steps (lr_recovery_steps) it undoes the cut.

Mixed Precision (fp16)

amp_example.py

scaler = torch.amp.GradScaler("cuda")
controller = arc.Controller(model, optimizer, scaler=scaler)

for x, y in loader:
    optimizer.zero_grad()
    with torch.autocast("cuda", dtype=torch.float16):
        loss = F.cross_entropy(model(x), y)
    scaler.scale(loss).backward()

    action = controller.step(loss)          # ARC unscales; don't call scaler.unscale_()
    if not action.should_skip:
        scaler.step(optimizer)
        scaler.update()

ARC checks the gradients at true scale. The scaler's routine overflow steps come back as a WARNING, not a rollback.

Hugging Face Trainer

hf_example.py

# pip install "arc-training[hf]"
from transformers import Trainer, TrainingArguments
from arc.api.integrations import HuggingFaceCallback

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="out", fp16=True),
    train_dataset=train_ds,
    callbacks=[HuggingFaceCallback()],
)
trainer.train()

The callback runs ARC just before each optimizer step and skips a poisoned update. With fp16=True it leaves overflow handling to the trainer's own scaler.

Lightning

lightning_example.py

# pip install "arc-training[lightning]"
import lightning as L
from arc.api.integrations import LightningCallback

trainer = L.Trainer(max_epochs=10, callbacks=[LightningCallback()])
trainer.fit(lit_module, train_loader)

Save and Resume

resume_example.py

controller = arc.Controller(model, optimizer, scheduler=scheduler)
start = 0
if os.path.exists("run/latest.pt"):
    start = controller.load_checkpoint("run/latest.pt")   # model, optimizer, scheduler, RNG, ARC state

for step in range(start, total_steps):
    ...                                                     # forward, backward, controller.step, optimizer.step
    if step % 1000 == 0:
        controller.save_checkpoint("run/latest.pt")            # atomic, weights_only-safe

Physics-Informed Neural Network

pinn_example.py

from arc import PINNStabilizer

class HeatPINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64), nn.Tanh(),
            nn.Linear(64, 64), nn.Tanh(),
            nn.Linear(64, 1)
        )

    def forward(self, x):
        return self.net(x)

model = HeatPINN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
stabilizer = PINNStabilizer(model, n_loss_terms=3)

for epoch in range(5000):
    inputs = torch.rand(100, 2, requires_grad=True)
    u = model(inputs)
    pde_loss = compute_pde_residual(u, inputs)
    bc_loss  = compute_boundary_loss(model)
    ic_loss  = compute_initial_loss(model)
    total = stabilizer.get_stabilized_loss([pde_loss, bc_loss, ic_loss])
    optimizer.zero_grad()
    total.backward()
    stabilizer.stabilize_step()
    optimizer.step()
    stabilizer.update(epoch, total)

Recovery Ledger and ROI Report

reliability_report.py

from arc import Controller, ControllerConfig
from arc.api.enterprise_report import EnterpriseReport

config = ControllerConfig(
    gpu_count=8, gpu_hourly_cost=2.50,           # $ per GPU-hour
    gpu_power_watts=700, carbon_intensity_kg_per_kwh=0.4,
    baseline_checkpoint_interval=1000,           # you save durable checkpoints every 1000 steps
    team="vision", project="resnet-sweep",
)
controller = Controller(model, optimizer, config=config)

# ... training loop as above ...

controller.export_recovery_ledger("ledger.json")
assert controller.verify_ledger_integrity()
print(controller.reliability_roi_report())
EnterpriseReport(controller).save(json_path="report.json", html_path="report.html")

Savings are measured against restarting from your last durable checkpoint, so baseline_checkpoint_interval must match how often you really save. Without it the report shows work lost and leaves savings as unknown. Open report.html in any browser. It is a single self-contained file, so it can be attached to a ticket or emailed as-is.

RL Fine-Tuning

rl_monitoring.py

from arc.signals.rl_stability import RLStabilityMonitor
from arc.signals.rollout_health import RolloutHealthMonitor

rl_monitor = RLStabilityMonitor(kl_threshold=0.1)
rollout_monitor = RolloutHealthMonitor(ngram_n=3)

for iteration in range(num_iterations):
    samples = generate(policy, prompts)                    # list of token-id lists
    stats = ppo_update(policy, ref_policy, samples)

    report = rl_monitor.update(kl=stats["kl"], entropy=stats["entropy"],
                               reward=stats["reward_mean"],
                               true_objective=held_out_eval(policy),
                               advantage=stats["advantage_mean"])
    health = rollout_monitor.check_batch(samples)

    if report.reward_hacking or not health.is_healthy:
        print(report.messages + health.messages)
        break

Continual Learning

EWC is exposed through the legacy ArcV2 API.

continual_learning.py

from arc import ArcV2

arc = ArcV2.auto(model, optimizer)

arc.begin_task("mnist")
for epoch in range(10):
    for x, y in mnist_loader:
        optimizer.zero_grad()
        loss = F.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()
        arc.step(loss)
    arc.end_epoch(epoch)

arc.consolidate_task(mnist_loader)

arc.begin_task("fashion_mnist")
for epoch in range(10):
    for x, y in fashion_loader:
        optimizer.zero_grad()
        loss = F.cross_entropy(model(x), y) + arc.get_ewc_loss()
        loss.backward()
        optimizer.step()
        arc.step(loss)
    arc.end_epoch(epoch)

Uncertainty Quantification

uncertainty.py

from arc import ConformalPredictor

cp = ConformalPredictor(model, alpha=0.1)
cp.calibrate(calibration_loader)

result = cp.predict(test_input)
print(result.prediction, result.confidence, result.set_members)

Adversarial Defense

adversarial.py

from arc import AdversarialDetector

detector = AdversarialDetector(model)
detector.fit(clean_data_loader)

clean_alert = detector.detect(normal_input)
bad_alert = detector.detect(adversarial_input)

print(clean_alert.is_adversarial)
print(bad_alert.is_adversarial)