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.
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
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
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
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
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)
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)
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
# 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
# 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
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-safePhysics-Informed Neural Network
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
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
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)
breakContinual Learning
EWC is exposed through the legacy ArcV2 API.
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
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
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)