import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Optional, Tuple
class ConstraintConfig:
def __init__(
self,
names: List[str] = (“reciprocity”, “testifiability”, “decidability”),
heads_per_constraint: int = 1,
trace_token_scores: bool = True,
trace_attn_maps: bool = False,
router_hidden: int = 128,
loss_weights: Dict[str, float] = None,
variant: str = “A”, # “A” = additive capacity; “B” = constant capacity
):
self.names = list(names)
self.heads_per_constraint = int(heads_per_constraint)
self.trace_token_scores = trace_token_scores
self.trace_attn_maps = trace_attn_maps
self.router_hidden = router_hidden
self.loss_weights = loss_weights or {name: 1.0 for name in names}
assert variant in (“A”, “B”)
self.variant = variant
class MultiHeadAttentionWithConstraints(nn.Module):
“””
MHA with extra ‘constraint heads’ dedicated to lawful reasoning.
Returns the standard MHA output plus an audit ‘traces’ dict.
“””
def __init__(
self,
d_model: int,
n_heads_base: int,
n_layers: int = 1, # not used here, but kept for parity with block builders
constraint: Optional[ConstraintConfig] = None,
dropout: float = 0.0,
):
super().__init__()
self.d_model = d_model
self.n_heads_base = n_heads_base
self.constraint = constraint or ConstraintConfig()
self.dropout = nn.Dropout(dropout)
# — head accounting —
self.n_constraint_heads = self.constraint.heads_per_constraint * len(self.constraint.names)
self.n_total_heads = self.n_heads_base + self.n_constraint_heads
# Per-head dimensions.
# Variant A: keep dk=dv=d_model//n_heads_base for base heads; use SAME for constraint heads → more concat width.
# Variant B: set dk=dv so that (n_total_heads * dv) ≈ d_model → constant concat width.
if self.constraint.variant == “A”:
= self.dv = d_model // self.n_heads_base
concat_width = self.dv * self.n_total_heads # grows with extra heads
else:
= self.dv = d_model // self.n_total_heads
concat_width = self.dv * self.n_total_heads # ~== d_model
# — base head projections —
self.Wq_base = nn.Linear(d_model,
* self.n_heads_base, bias=False)
self.Wk_base = nn.Linear(d_model,
* self.n_heads_base, bias=False)
self.Wv_base = nn.Linear(d_model, self.dv * self.n_heads_base, bias=False)
# — constraint head projections (separate parameterization) —
if self.n_constraint_heads > 0:
self.Wq_con = nn.Linear(d_model,
* self.n_constraint_heads, bias=False)
self.Wk_con = nn.Linear(d_model,
* self.n_constraint_heads, bias=False)
self.Wv_con = nn.Linear(d_model, self.dv * self.n_constraint_heads, bias=False)
else:
self.Wq_con = self.Wk_con = self.Wv_con = None
# One output projection over concatenated head outputs
self.Wo = nn.Linear(concat_width, d_model, bias=False)
# — constraint routers & scorers —
# A simple router that gates constraint heads from a pooled token (e.g., first token or mean pool)
route_in = d_model
self.router = nn.Sequential(
nn.Linear(route_in, self.constraint.router_hidden),
nn.ReLU(),
nn.Linear(self.constraint.router_hidden, self.n_constraint_heads),
)
# Per constraint head token scorer (for audit + auxiliary loss)
# We use an MLP->scalar per token for each constraint head.
self.token_scorers = nn.ModuleList([
nn.Sequential(
nn.Linear(self.dv, self.dv),
nn.ReLU(),
nn.Linear(self.dv, 1) # scalar per token
) for _ in range(self.n_constraint_heads)
])
# Helper: map constraint names to contiguous head ranges
self._constraint_spans = {}
idx = 0
for name in self.constraint.names:
self._constraint_spans[name] = (idx, idx + self.constraint.heads_per_constraint)
idx += self.constraint.heads_per_constraint
self.scale = (
) ** -0.5 # 1/sqrt(dk)
def _split_heads(self, x: torch.Tensor, n_heads: int, head_dim: int) -> torch.Tensor:
# x: [B, T, n_heads * head_dim] -> [B, n_heads, T, head_dim]
B, T, _ = x.shape
x = x.view(B, T, n_heads, head_dim).transpose(1, 2)
return x # [B, H, T, D]
def _combine_heads(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, H, T, D] -> [B, T, H*D]
B, H, T, D = x.shape
return x.transpose(1, 2).contiguous().view(B, T, H * D)
def forward(
self,
x: torch.Tensor,
attn_mask: Optional[torch.Tensor] = None, # shape broadcastable to [B, 1, T, T]
need_weights: bool = False,
router_hint: Optional[torch.Tensor] = None, # optional [B, d_model] routing context
) -> Tuple[torch.Tensor, Dict]:
“””
x: [B, T, d_model]
returns: (y: [B, T, d_model], traces: dict)
“””
B, T, _ = x.shape
# — projections —
Qb = self._split_heads(self.Wq_base(x), self.n_heads_base,
) # [B, Hb, T, dk]
Kb = self._split_heads(self.Wk_base(x), self.n_heads_base,
)
Vb = self._split_heads(self.Wv_base(x), self.n_heads_base, self.dv)
if self.n_constraint_heads > 0:
Qc = self._split_heads(self.Wq_con(x), self.n_constraint_heads,
) # [B, Hc, T, dk]
Kc = self._split_heads(self.Wk_con(x), self.n_constraint_heads,
)
Vc = self._split_heads(self.Wv_con(x), self.n_constraint_heads, self.dv)
else:
Qc = Kc = Vc = None
# — router gates for constraint heads —
if self.n_constraint_heads > 0:
if router_hint is None:
# Use mean pool over sequence as a cheap context
router_hint = x.mean(dim=1) # [B, d_model]
gates = torch.sigmoid(self.router(router_hint)) # [B, Hc]
gates = gates.view(B, self.n_constraint_heads, 1, 1) # broadcast over T,T
else:
gates = None
# — scaled dot-product attention —
def attn(Q, K, V, gate=None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
# Q,K,V: [B, H, T, D]
scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale # [B, H, T, T]
if attn_mask is not None:
scores = scores + attn_mask # mask contains 0 or -inf
if gate is not None:
# Light-handed way: add log(gate) to diagonal of scores to boost self-focus when gate is low/high.
# Alternatively, multiply V by gate (next line). We do both minimally:
V = V * gate # [B, H, T, D]
P = torch.softmax(scores, dim=-1)
P = self.dropout(P)
out = torch.matmul(P, V) # [B, H, T, D]
return out, (P if need_weights else None)
Hb, Pb = attn(Qb, Kb, Vb, None)
if self.n_constraint_heads > 0:
Hc, Pc = attn(Qc, Kc, Vc, gates)
else:
Hc, Pc = None, None
# — concat heads and project —
if Hc is not None:
H_all =
([Hb, Hc], dim=1) # [B, Hb+Hc, T, D]
else:
H_all = Hb
Z = self._combine_heads(H_all) # [B, T, (Hb+Hc)*Dv]
y = self.Wo(Z) # [B, T, d_model]
# — traces (audit) —
traces = {“constraint”: {}, “need_weights”: need_weights}
if self.n_constraint_heads > 0:
# token_scores per head: [B, Hc, T]
head_scores = []
for h in range(self.n_constraint_heads):
# Take that head’s token states: [B, 1, T, Dv] -> [B, T, Dv]
h_states = Hc[:, h, :, :] # [B, T, Dv]
s = self.token_scorers[h](h_states).squeeze(-1) # [B, T]
head_scores.append(s)
head_scores = torch.stack(head_scores, dim=1) # [B, Hc, T]
traces[“constraint”][“head_token_scores”] = head_scores # raw per-head token scalars
# Aggregate to named constraints
for name in self.constraint.names:
lo, hi = self._constraint_spans[name]
# mean across heads in that group
token_scores = head_scores[:, lo:hi, :].mean(dim=1) # [B, T]
entry = {“token_scores”: token_scores}
if need_weights and self.constraint.trace_attn_maps:
entry[“attn_maps”] = Pc[:, lo:hi, :, :] # [B, Hc_g, T, T]
traces[“constraint”][name] = entry
# Also expose router gates (per-batch, per-head)
traces[“constraint”][“router_gates”] = gates.squeeze(-1).squeeze(-1) # [B, Hc]
# Optionally expose base attention maps
if need_weights:
traces[“base_attn_maps”] = Pb # [B, Hb, T, T]
return y, traces
.no_grad()
def explain(self, traces: Dict, tokens: List[str], constraint_name: str = “reciprocity”) -> str:
“””
Human-readable audit line from token_scores.
“””
if “constraint” not in traces or constraint_name not in traces[“constraint”]:
return “No constraint trace.”
ts = traces[“constraint”][constraint_name][“token_scores”] # [B, T]
ts = ts[0] # first in batch
# Make a short explanation by selecting top-k contributing tokens
k = min(5, len(tokens))
top_idx = torch.topk(ts, k=k).indices.tolist()
parts = [f”{tokens[i]}:{ts[i]:.2f}” for i in top_idx]
return f”{constraint_name} focus → ” + “, “.join(parts)