Arithmetic Consistency on MNIST¶

In this example, a network receives two MNIST images and predicts three scalars: the first digit, the second digit and their sum. We train these outputs independently and add a property requiring the predicted sum to agree with the two predicted digits within a tolerance of 0.5. This regression formulation keeps the output property affine.

First, we import the required functions.

In [ ]:

import torch
from torch.utils.data import DataLoader, TensorDataset

from CTRAIN.data_loaders import load_mnist
from CTRAIN.model_definitions import MNISTAdditionRegressor
from CTRAIN.model_wrappers import PropertyModelWrapper
from CTRAIN.properties import Outputs, And, Box, PropertySpec

import torch from torch.utils.data import DataLoader, TensorDataset from CTRAIN.data_loaders import load_mnist from CTRAIN.model_definitions import MNISTAdditionRegressor from CTRAIN.model_wrappers import PropertyModelWrapper from CTRAIN.properties import Outputs, And, Box, PropertySpec

Now, we load one batch of MNIST images and arrange consecutive images into pairs. For this introductory example, the resulting pairs form the task dataset. We disable normalisation so that the declared property boxes use pixel coordinates directly.

In [ ]:

mnist_loader, _ = load_mnist(batch_size=128, normalise=False, val_split=False, data_root="./data")
images, labels = next(iter(mnist_loader))
inputs = images.reshape(-1, 2, 28, 28)
operands = labels.reshape(-1, 2).float()
targets = torch.cat((operands, operands.sum(dim=1, keepdim=True)), dim=1)
train_loader = DataLoader(TensorDataset(inputs, targets), batch_size=16, shuffle=True)
model = MNISTAdditionRegressor()

mnist_loader, _ = load_mnist(batch_size=128, normalise=False, val_split=False, data_root="./data") images, labels = next(iter(mnist_loader)) inputs = images.reshape(-1, 2, 28, 28) operands = labels.reshape(-1, 2).float() targets = torch.cat((operands, operands.sum(dim=1, keepdim=True)), dim=1) train_loader = DataLoader(TensorDataset(inputs, targets), batch_size=16, shuffle=True) model = MNISTAdditionRegressor()

We define the arithmetic property and two small input boxes around the first two pairs. The sum is a learned output, rather than a fixed addition layer, so the constraint is something the network must learn.

In [ ]:

y = Outputs()
residual = y[2] - y[0] - y[1]
eps = 0.01
property_spec = PropertySpec(
    "arithmetic_consistency",
    And(residual >= -0.5, residual <= 0.5),
    [Box((pair - eps).clamp(0, 1), (pair + eps).clamp(0, 1)) for pair in inputs[:2]],
)

y = Outputs() residual = y[2] - y[0] - y[1] eps = 0.01 property_spec = PropertySpec( "arithmetic_consistency", And(residual >= -0.5, residual <= 0.5), [Box((pair - eps).clamp(0, 1), (pair + eps).clamp(0, 1)) for pair in inputs[:2]], )

Now, we combine the supervised regression task with an IBP property objective and start training.

In [ ]:

wrapped_model = PropertyModelWrapper(
    model=model,
    input_shape=(2, 28, 28),
    properties=[property_spec],
    task_loss=torch.nn.MSELoss(),
    num_epochs=30,
    training_method="IBP",
    property_weight=1.0,
    device="cpu",
)
wrapped_model.train_model(train_loader)

wrapped_model = PropertyModelWrapper( model=model, input_shape=(2, 28, 28), properties=[property_spec], task_loss=torch.nn.MSELoss(), num_epochs=30, training_method="IBP", property_weight=1.0, device="cpu", ) wrapped_model.train_model(train_loader)

Finally, we check the arithmetic property over the declared boxes.

In [ ]:

report = wrapped_model.certify_properties(method="IBP")
print("Declared boxes certified:", report["certified"])
print("Cell outcomes:", [cell["outcome"] for cell in report["cells"]])

report = wrapped_model.certify_properties(method="IBP") print("Declared boxes certified:", report["certified"]) print("Cell outcomes:", [cell["outcome"] for cell in report["cells"]])

This result concerns arithmetic consistency over the two declared boxes. Consistent predictions can still be incorrect, and certification here does not cover the remaining MNIST data. The classification model with separate digit and sum heads uses a different argmax property; it is not the regression formulation demonstrated here.