TorchLean Transformer IBP Workflow

Small end-to-end workflow:

TorchLean (MHA + LayerNorm + MSE) → lower to NN.IR.Graph → run:

• IBP (runIBP)
• optional CROWN output bounds (outputBoxCROWN?)
• optional objective-dependent backward CROWN (backwardObjectiveBox?)

The example has one batch, two tokens, two features, and one attention head. On the CLI's rounded backends, both CROWN queries use directed backward propagation and may fall back to IBP bounds. This is a runtime integration check, not a checked model certificate.

Run: lake exe verify -- torchlean-transformer-ibp lake exe verify -- torchlean-transformer-ibp --with-crown lake exe verify -- torchlean-transformer-ibp --arithmetic ieee

def NN.Verification.Builtin.TransformerIBPWorkflow.n : ℕ

Sequence length for the transformer verification example.

def NN.Verification.Builtin.TransformerIBPWorkflow.dModel : ℕ

Model embedding dimension.

def NN.Verification.Builtin.TransformerIBPWorkflow.numHeads : ℕ

Number of attention heads.

def NN.Verification.Builtin.TransformerIBPWorkflow.headDim : ℕ

Per-head embedding dimension.

def NN.Verification.Builtin.TransformerIBPWorkflow.batch : ℕ

Batch size for the transformer verification example.

def NN.Verification.Builtin.TransformerIBPWorkflow.xShape : Spec.Shape

Input shape (batch × n × dModel).

def NN.Verification.Builtin.TransformerIBPWorkflow.wProjShape : Spec.Shape

Projection weight shape for Q/K/V: (dModel × (numHeads*headDim)).

def NN.Verification.Builtin.TransformerIBPWorkflow.wOShape : Spec.Shape

Output projection weight shape: ((numHeads*headDim) × dModel).

def NN.Verification.Builtin.TransformerIBPWorkflow.gammaShape : Spec.Shape

LayerNorm scale parameter shape, matching the feature dimension.

def NN.Verification.Builtin.TransformerIBPWorkflow.betaShape : Spec.Shape

LayerNorm beta shape, matching the feature dimension.

def NN.Verification.Builtin.TransformerIBPWorkflow.targetShape : Spec.Shape

MSE target shape (matches the model output shape).

def NN.Verification.Builtin.TransformerIBPWorkflow.paramShapes : List Spec.Shape

Parameter shapes list for modelLoss (Wq,Wk,Wv,Wo,gamma,beta,target).

def NN.Verification.Builtin.TransformerIBPWorkflow.modelLoss {α : Type} [TorchLean.Storage α] [Context α] : Runtime.Autograd.Model.Program α (paramShapes ++ [xShape]) []

TorchLean program: mha -> layer_norm -> mse_loss, returning a scalar loss.

def NN.Verification.Builtin.TransformerIBPWorkflow.runMain {α : Type} [TorchLean.Storage α] [Context α] [ToString α] [TorchLean.Runtime.FromFloat α] [MLTheory.CROWN.BoundOps α] [MLTheory.CROWN.NonlinearBoundOps α] (withCrown : Bool) : IO Unit

Runtime-selected typed runner used by the CLI entrypoint.

def NN.Verification.Builtin.TransformerIBPWorkflow.main (args : List String) : IO Unit

CLI entry point for the transformer-IBP workflow.

This is wired into lake exe verify -- torchlean-transformer-ibp.

By default this command lowers the small attention/LayerNorm example to the verification IR and runs IBP on its scalar loss. Pass --with-crown to also request CROWN output and objective bounds. The separate torchlean-crown-ops command exercises linear/softmax and linear/MSE examples. Missing bounds or unexpected output dimensions fail the command.