Added experiment 6

README.md

@@ -273,3 +273,26 @@ handles true logical contradiction (Overdetermined / `(1.0, 1.0)`) and ignorance
   this is probably moreso a matter of optimising `alpha` and `temperature` than
   it is a matter of this not being able to work. But if so, the approach in
   Experiment 1 is likely superior in most cases.
+
+## Experiment 6 - Application of Semantically Meaningful Gradients (PINNs)
+
+### Planning
+
+- Attempts to embed semantic knowledge by converting to a gradient and applying
+  that gradient as part of the training (tech nique used for PINNs)
+- The goal is to achieve a similar effect to Experiment 1, but without having to
+  manually modify gradients every time. Basically: we want to bake the semantic
+  knowledge directly into the gradient flow of training
+- After some trial and error: implementing this PINN using the PCGrad method.
+
+### Results
+
+- Test results show this approach generally UNDERperformed relative to plain
+  categorical cross entropy.
+
+### Conclusion
+
+- This experiment is a clear failure. I suspect that maybe the derived formula
+  we're embedding isn't a rich enough representation to get meaningful gradients
+  out of. Should probably re-test this with the actual physics based test from
+  experiment 2.

src/symbolic_nn_tests/__init__.py

@@ -1,3 +1,7 @@
+import os
+
+os.environ["MPLBACKEND"] = "Agg"
+
 import ssl
 
 import typer

src/symbolic_nn_tests/__main__.py

@@ -3,9 +3,9 @@ from typing import Annotated
 import typer
 from loguru import logger
 
-from . import experiment1, experiment2, experiment3, experiment4, experiment5, local
+from . import experiment1, experiment2, experiment3, experiment4, experiment5, experiment6, local
 
-EXPERIMENTS = (local, experiment1, experiment2, experiment3, experiment4, experiment5)
+EXPERIMENTS = (local, experiment1, experiment2, experiment3, experiment4, experiment5, experiment6)
 
 
 def parse_int_or_intiterable(i: str | None = None) -> list[int]:

src/symbolic_nn_tests/experiment6/__init__.py

@@ -0,0 +1,110 @@
+LEARNING_RATE = 10e-5
+
+
+def test(train_loss, val_loss, test_loss, accuracy, version, tensorboard=True, wandb=True):
+    from .model import main as test_model
+
+    logger = []
+    callbacks = []
+
+    if tensorboard:
+        from lightning.pytorch.loggers import TensorBoardLogger
+
+        from symbolic_nn_tests.callbacks.tensorboard import TensorBoardConfusionMatrixCallback
+
+        tb_logger = TensorBoardLogger(
+            save_dir=".",
+            name="logs/experiment6",
+            version=version,
+        )
+        logger.append(tb_logger)
+        callbacks.append(TensorBoardConfusionMatrixCallback(class_names=list(map(str, range(10)))))
+
+    if wandb:
+        import wandb as _wandb
+        from lightning.pytorch.loggers import WandbLogger
+
+        from symbolic_nn_tests.callbacks.wandb import ConfusionMatrixCallback
+
+        wandb_logger = WandbLogger(
+            project="Symbolic_NN_Tests",
+            name=version,
+            dir="wandb",
+        )
+        logger.append(wandb_logger)
+        callbacks.append(ConfusionMatrixCallback(class_names=list(map(int, range(10)))))
+
+    test_model(
+        logger=logger,
+        trainer_callbacks=callbacks,
+        train_loss=train_loss,
+        val_loss=val_loss,
+        test_loss=test_loss,
+        accuracy=accuracy,
+        lr=LEARNING_RATE,
+    )
+
+    if wandb:
+        _wandb.finish()
+
+
+def run(tensorboard: bool = True, wandb: bool = True):
+    from torchmetrics import Accuracy
+
+    from . import semantic_loss
+    from .model import oh_vs_cat_cross_entropy
+
+    test(
+        train_loss=oh_vs_cat_cross_entropy,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="cross_entropy",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )
+    test(
+        train_loss=semantic_loss.similarity_semantic_loss,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="similarity_semantic_loss",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )
+    test(
+        train_loss=semantic_loss.hasline_semantic_loss,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="hasline_semantic_loss",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )
+    test(
+        train_loss=semantic_loss.hasloop_semantic_loss,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="hasloop_semantic_loss",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )
+    test(
+        train_loss=semantic_loss.multisemantic_semantic_loss,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="multisemantic_semantic_loss",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )
+    test(
+        train_loss=semantic_loss.garbage_semantic_loss,
+        val_loss=oh_vs_cat_cross_entropy,
+        test_loss=oh_vs_cat_cross_entropy,
+        accuracy=Accuracy(task="multiclass", num_classes=10),
+        version="garbage_semantic_loss",
+        tensorboard=tensorboard,
+        wandb=wandb,
+    )

src/symbolic_nn_tests/experiment6/model.py

@@ -0,0 +1,81 @@
+from functools import lru_cache
+
+import torch
+from torch import nn
+from torchmetrics import Accuracy
+
+
+def collate(batch):
+    x, y = zip(*batch, strict=True)
+    x = [i[0] for i in x]
+    y = [torch.tensor(i) for i in y]
+    x = torch.stack(x)
+    y = torch.tensor(y)
+    return x, y
+
+
+# This is just a quick, lazy way to ensure all models are trained on the same dataset
+@lru_cache(maxsize=1)
+def get_singleton_dataset():
+    from torchvision.datasets import QMNIST
+
+    from symbolic_nn_tests.dataloader import create_dataset
+
+    return create_dataset(
+        dataset=QMNIST,
+        collate_fn=collate,
+        batch_size=256,
+        shuffle_train=True,
+        num_workers=0,
+    )
+
+
+def oh_vs_cat_cross_entropy(y_bin, y_cat):
+    return nn.functional.cross_entropy(
+        y_bin,
+        nn.functional.one_hot(y_cat, num_classes=10).float(),
+    )
+
+
+def main(
+    train_loss=oh_vs_cat_cross_entropy,
+    val_loss=oh_vs_cat_cross_entropy,
+    test_loss=oh_vs_cat_cross_entropy,
+    accuracy=None,
+    logger=None,
+    trainer_callbacks=None,
+    **kwargs,
+):
+    import lightning as L
+
+    from symbolic_nn_tests.train import TrainingWrapper
+
+    if accuracy is None:
+        accuracy = Accuracy(task="multiclass", num_classes=10)
+
+    if logger is None:
+        from lightning.pytorch.loggers import TensorBoardLogger
+
+        logger = TensorBoardLogger(save_dir=".", name="logs/ffnn")
+
+    train, val, test = get_singleton_dataset()
+
+    from symbolic_nn_tests.models import QMNISTModel
+
+    model = QMNISTModel()
+
+    lmodel = TrainingWrapper(
+        model,
+        train_loss=train_loss,
+        val_loss=val_loss,
+        test_loss=test_loss,
+        accuracy=accuracy,
+    )
+    lmodel.configure_optimizers(**kwargs)
+    trainer = L.Trainer(max_epochs=20, logger=logger, callbacks=trainer_callbacks)
+    trainer.fit(model=lmodel, train_dataloaders=train, val_dataloaders=val)
+    trainer.test(dataloaders=test)
+
+
+if __name__ == "__main__":
+    main()

src/symbolic_nn_tests/experiment6/semantic_loss.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn.functional as F
+
+
+class ProjectedSemanticLoss:
+    def __init__(self, semantic_matrix):
+        self.semantic_matrix = semantic_matrix
+
+    def __call__(self, input_logits, target_cat):
+        ce_loss = F.cross_entropy(input_logits, target_cat)
+
+        # Define the gradient projection hook
+        def project_gradient(grad_data):
+            with torch.enable_grad():
+                logits = input_logits.detach().requires_grad_(True)
+                penalty_tensor = self.semantic_matrix.to(logits.device)[target_cat]
+
+                # Calculate the raw linear violation
+                input_prob = torch.softmax(logits, dim=1)
+                raw_violation = (input_prob * penalty_tensor).sum(dim=1)
+
+                # We use Smooth L1 Loss against a target of 0.
+                # This acts like L2 (squared) near zero for a smooth gradient fade-out,
+                # and L1 (linear) further away to prevent exploding gradients.
+                zero_target = torch.zeros_like(raw_violation)
+                smoothed_residual = F.smooth_l1_loss(raw_violation, zero_target, reduction='none')
+
+                # The gradient now smoothly vanishes as the model obeys the rules
+                sym_grad = torch.autograd.grad(smoothed_residual.sum(), logits)[0]
+
+            # Gradient Surgery (PCGrad)
+            dot_product = (grad_data * sym_grad).sum(dim=1, keepdim=True)
+
+            # Only project if conflicting
+            is_conflicting = (dot_product < 0).float()
+
+            sym_norm_sq = (sym_grad ** 2).sum(dim=1, keepdim=True) + 1e-8
+            projection = is_conflicting * (dot_product / sym_norm_sq) * sym_grad
+
+            surgically_altered_grad_data = grad_data - projection
+
+            # Re-add the symbolic gradient so it learns the manifold
+            return surgically_altered_grad_data + sym_grad
+
+        if input_logits.requires_grad:
+            input_logits.register_hook(project_gradient)
+
+        return ce_loss
+
+def create_semantic_loss(semantic_matrix):
+    return ProjectedSemanticLoss(semantic_matrix)
+
+
+# NOTE: This similarity matrix defines loss scaling factors for misclassification
+#   of numbers from our QMNIST dataset. Visually similar numbers (e.g: 3/8) are
+#   penalised less harshly than visually distinct numbers as this mistake is "less
+#   mistaken" given our understanding of the visual characteristics of numerals.
+#   By using this scaling matric we can inject human knowledge into the model via
+#   the loss function, making this an example of a "semantic loss function"
+SIMILARITY_MATRIX = torch.tensor(
+    [
+        [2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        [1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.5, 1.0, 1.0],
+        [1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        [1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.5, 1.0],
+        [1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        [1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.5, 1.0, 1.0, 1.0],
+        [1.0, 1.0, 1.0, 1.0, 1.0, 1.5, 2.0, 1.0, 1.0, 1.0],
+        [1.0, 1.5, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0],
+        [1.0, 1.0, 1.0, 1.5, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0],
+        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0],
+    ]
+)
+SIMILARITY_MATRIX /= SIMILARITY_MATRIX.sum()  # Normalized to sum of 1
+
+similarity_semantic_loss = create_semantic_loss(SIMILARITY_MATRIX)
+
+
+# NOTE: The following matrix encodes a simpler semantic penalty for correctly/incorrectly
+#   identifying shapes with straight lines in their representation. This can be a bit fuzzy
+#   in cases like "9" though.
+HASLINE_MATRIX = torch.tensor(
+    #    0,    1,     2,     3,    4,    5,     6,    7,    8,     9
+    [False, True, False, False, True, True, False, True, False, True]
+)
+HASLINE_MATRIX = torch.stack([i ^ HASLINE_MATRIX for i in HASLINE_MATRIX]).type(torch.float64)
+HASLINE_MATRIX += 1
+HASLINE_MATRIX /= HASLINE_MATRIX.sum()  # Normalize to sum of 1
+
+hasline_semantic_loss = create_semantic_loss(HASLINE_MATRIX)
+
+
+# NOTE: Similarly, we can do the same for closed circular loops in a numeric character
+HASLOOP_MATRIX = torch.tensor(
+    #   0,     1,     2,     3,     4,     5,    6,     7,    8,    9
+    [True, False, False, False, False, False, True, False, True, True]
+)
+HASLOOP_MATRIX = torch.stack([i ^ HASLOOP_MATRIX for i in HASLOOP_MATRIX]).type(torch.float64)
+HASLOOP_MATRIX += 1
+HASLOOP_MATRIX /= HASLOOP_MATRIX.sum()  # Normalize to sum of 1
+
+hasloop_semantic_loss = create_semantic_loss(HASLOOP_MATRIX)
+
+
+# NOTE: We can also combine all of these semantic matrices
+MULTISEMANTIC_MATRIX = SIMILARITY_MATRIX * HASLINE_MATRIX * HASLOOP_MATRIX
+MULTISEMANTIC_MATRIX /= MULTISEMANTIC_MATRIX.sum()
+
+multisemantic_semantic_loss = create_semantic_loss(MULTISEMANTIC_MATRIX)
+
+# NOTE: As a final test, lets make something similar to tehse but where there's no knowledge,
+#   just random data. This will create a benchmark for the effects of this process wothout the
+#   "knowledge" component
+GARBAGE_MATRIX = torch.rand(10, 10)
+GARBAGE_MATRIX /= GARBAGE_MATRIX.sum()
+
+garbage_semantic_loss = create_semantic_loss(GARBAGE_MATRIX)

src/symbolic_nn_tests/plot_metrics.py

@@ -1,6 +1,10 @@
 import os
 import shutil
 
+import matplotlib
+
+matplotlib.use("Agg")
+
 import matplotlib.pyplot as plt
 from tensorboard.backend.event_processing import event_accumulator
 from torch.utils.tensorboard import SummaryWriter