Release AS2 framework with Hydra-selected instantiations

69928d6 verified 9 days ago

8.66 kB

"""Objectives and diagnostics owned by the Sudoku instantiation."""

from typing import Any

import torch
import torch.nn.functional as F
from torch import nn


class UniquenessLoss(nn.Module):
    """Soft uniqueness regulariser: penalises high-entropy blank-cell distributions.
    A Sudoku puzzle has a *unique* solution when every blank cell's value is
    fully determined by the declared constraints and the given clues.  A model
    that spreads probability across multiple equally-valid digit assignments
    for blank cells implicitly encodes multiple valid solutions.
    We proxy the uniqueness constraint by minimising the mean normalised
    entropy of post-reasoning distributions over blank cells, driving the
    model to commit to specific symbol assignments rather than remaining
    ambiguous.  This is a fully-differentiable continuous analog of the
    iterative nogood-based uniqueness enforcement in Mulamba et al. (2024,
    *Constraints*, doi:10.1007/s10601-024-09372-9).
    Loss::
        L_uniq = mean_{i ∈ blank} H(p_i) / log(|Σ|)
    where H(p_i) = −Σ_s p_{i,s} log p_{i,s} is per-cell entropy and |Σ| is
    the symbol-domain size.  Normalising by log(|Σ|) maps the range to [0, 1]:
    0 when every blank cell is one-hot (fully committed), 1 when every blank
    cell is uniform (maximally ambiguous).
    """

    def forward(
        self,
        probs: torch.Tensor,
        blank_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            probs:      (B, 81, |Σ|) softmax probabilities — post-reasoning.
            blank_mask: (B, 81) bool tensor, True for blank (non-clue) cells.
        Returns:
            Scalar normalised entropy loss in [0, 1].
        """
        flat_blank = blank_mask.reshape(-1)  # (B*81,)
        flat_probs = probs.reshape(-1, probs.shape[-1])  # (B*81, |Σ|)
        blank_probs = flat_probs[flat_blank]  # (n_blank, |Σ|)
        if blank_probs.shape[0] == 0:
            return probs.sum() * 0.0  # differentiable zero
        entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1)  # (n_blank,)
        log_sigma = torch.log(torch.tensor(float(probs.shape[-1]), device=probs.device))
        return entropy.mean() / log_sigma


def compute_sudoku_objective(
    pre_logits: torch.Tensor,
    post_logits: torch.Tensor,
    labels: torch.Tensor,
    puzzle_mask: torch.Tensor,
    *,
    current_epoch: int,
    lambda_ic: float,
    lambda_blank: float,
    blank_decay_epochs: int,
    ic_warmup_epochs: int,
    lambda_uniq: float,
    constraint_loss: nn.Module,
    uniqueness_loss: nn.Module,
) -> dict[str, Any]:
    """Sudoku supervision, curriculum and diagnostics; separate from AS2 reasoning."""
    # ── 6. Loss computation ───────────────────────────────────────────
    # loss_cnn   — direct CE on the CNN concept head (pre_logits) for clue
    #              cells.  PRIMARY gradient signal for the CNN.
    #
    # loss_sup   — CE on the Transformer output (post_logits) for clue
    #              cells.  Keeps the Transformer grounded at clue positions.
    #
    # loss_blank — CE on blank-cell Transformer outputs against GT labels.
    #              Provides direct supervision for blank cells (the old
    #              model had NONE).  Weighted by a curriculum schedule
    #              that decays from full to 0.1× over training so the
    #              model gradually shifts to relying on constraint reasoning.
    #
    # loss_ic    — T_P fixed-point constraint energy on ALL 81 cells.
    #              Penalises deviation from the T_P operator's fixed point.
    #              Ramps up from 0 over ic_warmup_epochs.
    clue_flat = puzzle_mask.reshape(-1).bool()  # (B*81,)
    blank_flat = ~clue_flat
    labels_flat = labels.reshape(-1)

    # Direct CNN supervision (most important for learning digits)
    loss_cnn = F.cross_entropy(
        pre_logits.reshape(-1, 9)[clue_flat],
        labels_flat[clue_flat],
    )

    # Transformer output supervision at clue positions
    loss_sup = F.cross_entropy(
        post_logits.reshape(-1, 9)[clue_flat],
        labels_flat[clue_flat],
    )

    # Blank-cell CE supervision with curriculum decay
    if blank_flat.any():
        loss_blank = F.cross_entropy(
            post_logits.reshape(-1, 9)[blank_flat],
            labels_flat[blank_flat],
        )
    else:
        loss_blank = torch.tensor(0.0, device=post_logits.device)

    # Curriculum weight: linear decay from 1.0 to 0.1 over blank_decay_epochs
    epoch = current_epoch
    blank_scale = max(0.1, 1.0 - 0.9 * epoch / max(1, blank_decay_epochs))

    # T_P constraint loss with warmup ramp
    soft_model = F.softmax(post_logits, dim=-1)  # (B, 81, 9)
    loss_ic = constraint_loss(soft_model)
    ic_scale = min(1.0, epoch / max(1, ic_warmup_epochs))

    # Uniqueness regulariser: penalise high blank-cell entropy post-reasoning
    blank_mask_2d = ~puzzle_mask.bool()  # (B, 81) bool
    loss_uniq = uniqueness_loss(soft_model, blank_mask_2d)

    loss = (
        loss_cnn
        + loss_sup
        + lambda_blank * blank_scale * loss_blank
        + lambda_ic * ic_scale * loss_ic
        + lambda_uniq * ic_scale * loss_uniq
    )

    # ── 7. Diagnostics ────────────────────────────────────────────────
    # Blank-cell entropy: high → model is uncertain, low → confident
    if blank_flat.any():
        blank_probs = soft_model.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
        blank_entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1).mean()
    else:
        blank_entropy = torch.tensor(0.0, device=post_logits.device)

    # ── 8. Predictions ────────────────────────────────────────────────
    predictions = post_logits.argmax(dim=-1)  # (B, 81)

    # Flat clue-cell tensors:
    # Use pre_logits (CNN output) so that clue_accuracy reflects CNN quality
    # independent of the Transformer.
    clue_logits = pre_logits.reshape(-1, 9)[clue_flat]  # (n_clue, 9)
    clue_labels = labels_flat[clue_flat]  # (n_clue,)

    # ── 9. Perception-vs-reasoning diagnostic tensors ─────────────────
    # Pre-reasoning CNN logits for blank cells (measures perception quality
    # on the cells that actually require reasoning to fill).
    blank_pre_logits = pre_logits.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
    blank_labels = labels_flat[blank_flat]  # (n_blank,)

    # Post-reasoning logits split by clue / blank (measures how much
    # the Transformer improves on CNN quality, and whether it corrupts
    # the confident clue-cell predictions).
    blank_post_logits = post_logits.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
    clue_post_logits = post_logits.reshape(-1, 9)[clue_flat]  # (n_clue, 9)

    return {
        "loss": loss,
        "logits": post_logits,  # (B, 81, 9) — for board-level metrics
        "predictions": predictions,  # (B, 81)
        # ── CNN (pre-reasoning) split ─────────────────────────────────
        "clue_logits": clue_logits,  # (n_clue, 9)  CNN on clue cells
        "clue_labels": clue_labels,  # (n_clue,)
        "blank_pre_logits": blank_pre_logits,  # (n_blank, 9) CNN on blank cells
        "blank_labels": blank_labels,  # (n_blank,)
        # ── Transformer (post-reasoning) split ───────────────────────
        "blank_post_logits": blank_post_logits,  # (n_blank, 9)
        "clue_post_logits": clue_post_logits,  # (n_clue, 9)
        # ── Scalar losses ─────────────────────────────────────────────
        "loss_cnn": loss_cnn.detach(),
        "loss_sup": loss_sup.detach(),
        "loss_blank": loss_blank.detach(),
        "loss_ic": loss_ic.detach(),
        "loss_uniq": loss_uniq.detach(),
        # Weighted contributions (what actually drives the total loss)
        "loss_blank_w": (lambda_blank * blank_scale * loss_blank).detach(),
        "loss_ic_w": (lambda_ic * ic_scale * loss_ic).detach(),
        "loss_uniq_w": (lambda_uniq * ic_scale * loss_uniq).detach(),
        # Curriculum scales for diagnostics
        "blank_scale": blank_scale,
        "ic_scale": ic_scale,
        "blank_entropy": blank_entropy.detach(),
    }

	"""Objectives and diagnostics owned by the Sudoku instantiation."""

	from typing import Any

	import torch
	import torch.nn.functional as F
	from torch import nn


	class UniquenessLoss(nn.Module):
	"""Soft uniqueness regulariser: penalises high-entropy blank-cell distributions.

	A Sudoku puzzle has a unique solution when every blank cell's value is
	fully determined by the declared constraints and the given clues. A model
	that spreads probability across multiple equally-valid digit assignments
	for blank cells implicitly encodes multiple valid solutions.

	We proxy the uniqueness constraint by minimising the mean normalised
	entropy of post-reasoning distributions over blank cells, driving the
	model to commit to specific symbol assignments rather than remaining
	ambiguous. This is a fully-differentiable continuous analog of the
	iterative nogood-based uniqueness enforcement in Mulamba et al. (2024,
	Constraints, doi:10.1007/s10601-024-09372-9).

	Loss::

	L_uniq = mean_{i ∈ blank} H(p_i) / log(\|Σ\|)

	where H(p_i) = −Σ_s p_{i,s} log p_{i,s} is per-cell entropy and \|Σ\| is
	the symbol-domain size. Normalising by log(\|Σ\|) maps the range to [0, 1]:
	0 when every blank cell is one-hot (fully committed), 1 when every blank
	cell is uniform (maximally ambiguous).
	"""

	def forward(
	self,
	probs: torch.Tensor,
	blank_mask: torch.Tensor,
	) -> torch.Tensor:
	"""
	Args:
	probs: (B, 81, \|Σ\|) softmax probabilities — post-reasoning.
	blank_mask: (B, 81) bool tensor, True for blank (non-clue) cells.

	Returns:
	Scalar normalised entropy loss in [0, 1].
	"""
	flat_blank = blank_mask.reshape(-1) # (B*81,)
	flat_probs = probs.reshape(-1, probs.shape[-1]) # (B*81, \|Σ\|)
	blank_probs = flat_probs[flat_blank] # (n_blank, \|Σ\|)
	if blank_probs.shape[0] == 0:
	return probs.sum() * 0.0 # differentiable zero
	entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1) # (n_blank,)
	log_sigma = torch.log(torch.tensor(float(probs.shape[-1]), device=probs.device))
	return entropy.mean() / log_sigma


	def compute_sudoku_objective(
	pre_logits: torch.Tensor,
	post_logits: torch.Tensor,
	labels: torch.Tensor,
	puzzle_mask: torch.Tensor,
	*,
	current_epoch: int,
	lambda_ic: float,
	lambda_blank: float,
	blank_decay_epochs: int,
	ic_warmup_epochs: int,
	lambda_uniq: float,
	constraint_loss: nn.Module,
	uniqueness_loss: nn.Module,
	) -> dict[str, Any]:
	"""Sudoku supervision, curriculum and diagnostics; separate from AS2 reasoning."""
	# ── 6. Loss computation ───────────────────────────────────────────
	# loss_cnn — direct CE on the CNN concept head (pre_logits) for clue
	# cells. PRIMARY gradient signal for the CNN.
	#
	# loss_sup — CE on the Transformer output (post_logits) for clue
	# cells. Keeps the Transformer grounded at clue positions.
	#
	# loss_blank — CE on blank-cell Transformer outputs against GT labels.
	# Provides direct supervision for blank cells (the old
	# model had NONE). Weighted by a curriculum schedule
	# that decays from full to 0.1× over training so the
	# model gradually shifts to relying on constraint reasoning.
	#
	# loss_ic — T_P fixed-point constraint energy on ALL 81 cells.
	# Penalises deviation from the T_P operator's fixed point.
	# Ramps up from 0 over ic_warmup_epochs.
	clue_flat = puzzle_mask.reshape(-1).bool() # (B*81,)
	blank_flat = ~clue_flat
	labels_flat = labels.reshape(-1)

	# Direct CNN supervision (most important for learning digits)
	loss_cnn = F.cross_entropy(
	pre_logits.reshape(-1, 9)[clue_flat],
	labels_flat[clue_flat],
	)

	# Transformer output supervision at clue positions
	loss_sup = F.cross_entropy(
	post_logits.reshape(-1, 9)[clue_flat],
	labels_flat[clue_flat],
	)

	# Blank-cell CE supervision with curriculum decay
	if blank_flat.any():
	loss_blank = F.cross_entropy(
	post_logits.reshape(-1, 9)[blank_flat],
	labels_flat[blank_flat],
	)
	else:
	loss_blank = torch.tensor(0.0, device=post_logits.device)

	# Curriculum weight: linear decay from 1.0 to 0.1 over blank_decay_epochs
	epoch = current_epoch
	blank_scale = max(0.1, 1.0 - 0.9 * epoch / max(1, blank_decay_epochs))

	# T_P constraint loss with warmup ramp
	soft_model = F.softmax(post_logits, dim=-1) # (B, 81, 9)
	loss_ic = constraint_loss(soft_model)
	ic_scale = min(1.0, epoch / max(1, ic_warmup_epochs))

	# Uniqueness regulariser: penalise high blank-cell entropy post-reasoning
	blank_mask_2d = ~puzzle_mask.bool() # (B, 81) bool
	loss_uniq = uniqueness_loss(soft_model, blank_mask_2d)

	loss = (
	loss_cnn
	+ loss_sup
	+ lambda_blank * blank_scale * loss_blank
	+ lambda_ic * ic_scale * loss_ic
	+ lambda_uniq * ic_scale * loss_uniq
	)

	# ── 7. Diagnostics ────────────────────────────────────────────────
	# Blank-cell entropy: high → model is uncertain, low → confident
	if blank_flat.any():
	blank_probs = soft_model.reshape(-1, 9)[blank_flat] # (n_blank, 9)
	blank_entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1).mean()
	else:
	blank_entropy = torch.tensor(0.0, device=post_logits.device)

	# ── 8. Predictions ────────────────────────────────────────────────
	predictions = post_logits.argmax(dim=-1) # (B, 81)

	# Flat clue-cell tensors:
	# Use pre_logits (CNN output) so that clue_accuracy reflects CNN quality
	# independent of the Transformer.
	clue_logits = pre_logits.reshape(-1, 9)[clue_flat] # (n_clue, 9)
	clue_labels = labels_flat[clue_flat] # (n_clue,)

	# ── 9. Perception-vs-reasoning diagnostic tensors ─────────────────
	# Pre-reasoning CNN logits for blank cells (measures perception quality
	# on the cells that actually require reasoning to fill).
	blank_pre_logits = pre_logits.reshape(-1, 9)[blank_flat] # (n_blank, 9)
	blank_labels = labels_flat[blank_flat] # (n_blank,)

	# Post-reasoning logits split by clue / blank (measures how much
	# the Transformer improves on CNN quality, and whether it corrupts
	# the confident clue-cell predictions).
	blank_post_logits = post_logits.reshape(-1, 9)[blank_flat] # (n_blank, 9)
	clue_post_logits = post_logits.reshape(-1, 9)[clue_flat] # (n_clue, 9)

	return {
	"loss": loss,
	"logits": post_logits, # (B, 81, 9) — for board-level metrics
	"predictions": predictions, # (B, 81)
	# ── CNN (pre-reasoning) split ─────────────────────────────────
	"clue_logits": clue_logits, # (n_clue, 9) CNN on clue cells
	"clue_labels": clue_labels, # (n_clue,)
	"blank_pre_logits": blank_pre_logits, # (n_blank, 9) CNN on blank cells
	"blank_labels": blank_labels, # (n_blank,)
	# ── Transformer (post-reasoning) split ───────────────────────
	"blank_post_logits": blank_post_logits, # (n_blank, 9)
	"clue_post_logits": clue_post_logits, # (n_clue, 9)
	# ── Scalar losses ─────────────────────────────────────────────
	"loss_cnn": loss_cnn.detach(),
	"loss_sup": loss_sup.detach(),
	"loss_blank": loss_blank.detach(),
	"loss_ic": loss_ic.detach(),
	"loss_uniq": loss_uniq.detach(),
	# Weighted contributions (what actually drives the total loss)
	"loss_blank_w": (lambda_blank * blank_scale * loss_blank).detach(),
	"loss_ic_w": (lambda_ic * ic_scale * loss_ic).detach(),
	"loss_uniq_w": (lambda_uniq * ic_scale * loss_uniq).detach(),
	# Curriculum scales for diagnostics
	"blank_scale": blank_scale,
	"ic_scale": ic_scale,
	"blank_entropy": blank_entropy.detach(),
	}