""" caff/model.py ============= Full CAFFModel — assembles encoder, CSV, per-hop scorers into a single trainable module. Per paper §8.4, learnable capacity must be < 12M params (~3.5% of the 340M-param frozen backbone). This is asserted at __init__. Architecture summary (paper Algorithm 1): for ℓ = 1 to L: C_ℓ ← BFS(G, S(Q), ℓ) # Eq. 1 C_ℓ ← FreqCap(C_ℓ, K_r) # Eq. 13 g_ℓ ← σ(P_ℓ z_{ℓ-1} + b^g_ℓ) # Eq. 16 Δ_ℓ ← U_ℓ diag(g_ℓ) V_ℓ^T # Eq. 17 W^ctx_ℓ ← W_0 + A_ℓ B_ℓ^T + Δ_ℓ # Eq. 18 for (h, r, t) ∈ C_ℓ: s ← σ(q^T W^ctx_ℓ e_r + v^T(q ⊙ e_r) + β_ℓ) # Eq. 19 if s ≥ θ: S_ℓ ← S_ℓ ∪ {(h,r,t)} z_ℓ ← CSV(S_ℓ) # Eq. 14 """ from __future__ import annotations import logging from dataclasses import dataclass import torch import torch.nn as nn from .config import AblationFlags, CAFFConfig from .csv import CSV from .encoders import RelationEmbeddingCache from .scorer import HopScorer logger = logging.getLogger(__name__) @dataclass class HopOutput: """Output of scoring a single hop.""" logits: torch.Tensor # pre-sigmoid scores, shape (N,) scores: torch.Tensor # sigmoid-activated, shape (N,) retained_mask: torch.Tensor # bool, shape (N,) — True if score >= θ z_next: torch.Tensor # CSV computed from retained set, shape (d,) class CAFFModel(nn.Module): """Full CAFF filtering model (paper §6). Parameters ---------- config : CAFFConfig Hyperparameters from paper §8.4. relation_cache : RelationEmbeddingCache Pre-encoded frozen relation embeddings. ablation : AblationFlags Component toggles for ablation study (paper §10.1). """ def __init__( self, config: CAFFConfig, relation_cache: RelationEmbeddingCache, ablation: AblationFlags | None = None, ) -> None: super().__init__() self.config = config self.ablation = ablation or AblationFlags() self.relation_cache = relation_cache d = config.d rho = config.rho L = config.L # Shared parameters self.W0 = nn.Parameter(torch.empty(d, d)) nn.init.xavier_uniform_(self.W0) self.v = nn.Parameter(torch.zeros(d)) # element-wise interaction nn.init.normal_(self.v, mean=0.0, std=0.01) # Per-hop scorers (one HopScorer per ℓ ∈ {1, ..., L}) self.hop_scorers = nn.ModuleList([ HopScorer( d=d, rho=rho, gate_activation=self.ablation.gate_activation, use_dbm=self.ablation.use_dbm, ) for _ in range(L) ]) # CSV is parameter-free, but holds a reference to the cache self.csv = CSV(relation_cache, pool=self.ablation.csv_pool) # Sanity check: paper §8.4 says learnable params < 12M n_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) logger.info( f"CAFFModel initialized: {n_trainable / 1e6:.2f}M trainable params " f"(paper budget: <12M)" ) if n_trainable > 12_000_000: logger.warning( f"Trainable param count {n_trainable / 1e6:.2f}M exceeds " f"paper's <12M budget. Check d, rho, L." ) @property def device(self) -> torch.device: return next(self.parameters()).device def get_hop_W_ctx( self, hop_idx: int, z_prev: torch.Tensor, ) -> torch.Tensor: """Compute W^ctx_ℓ for a batch of items at hop ℓ (Eq. 18). Parameters ---------- hop_idx : int in {0, 1, ..., L-1} z_prev : Tensor of shape (B, d) CSV from hop ℓ-1 (or zeros at hop 0 / when set is empty). If ablation.use_csv is False, z_prev is forcibly zeroed. Returns ------- W_ctx : Tensor of shape (B, d, d) """ if not self.ablation.use_csv: # Ablation §10.1: w/o CSV → z_{ℓ-1} ≡ 0 → reduces to DepthBilinear z_prev = torch.zeros_like(z_prev) return self.hop_scorers[hop_idx].compute_W_ctx(self.W0, z_prev) def score_hop_candidates( self, hop_idx: int, W_ctx_single: torch.Tensor, q: torch.Tensor, relation_names: list[str], return_logits: bool = False, ) -> torch.Tensor: """Score all candidates of a single (query, hop) — Eq. 19. Smart Engineering S3: W_ctx_single is computed ONCE per (query, hop), then this method scores all N_ℓ candidates against it in a single matmul. Parameters ---------- hop_idx : int W_ctx_single : Tensor of shape (d, d) Pre-computed W^ctx for one item. q : Tensor of shape (d,) ℓ2-normalized query embedding. relation_names : list of str, length N Relation names of the candidate triples. return_logits : bool If True, return pre-sigmoid logits (for BCEWithLogitsLoss). Returns ------- Tensor of shape (N,). """ E_r = self.relation_cache.get_batch(relation_names) # (N, d) scorer = self.hop_scorers[hop_idx] if return_logits: return scorer.score_logits(W_ctx_single, self.v, q, E_r) return scorer.score_candidates(W_ctx_single, self.v, q, E_r) @torch.no_grad() def filter_query( self, q: torch.Tensor, candidate_sets: list[list[tuple[str, str, str]]], ) -> list[list[tuple[str, str, str]]]: """Apply CAFF filtering to a single query — Algorithm 1 (paper). Parameters ---------- q : Tensor of shape (d,) ℓ2-normalized query embedding. candidate_sets : list of length L candidate_sets[ℓ] is the BFS-extracted (and FreqCapped) list of (h, r, t) tuples at hop ℓ+1. Returns ------- retained_sets : list of length L retained_sets[ℓ] is S_{ℓ+1} = {(h,r,t) : s_ℓ(...) ≥ θ}. """ L = self.config.L theta = self.config.theta d = self.config.d # Initialize: z_0 = 0, S_0 = ∅ (Algorithm 1, line 2) z_prev = torch.zeros(1, d, device=self.device) # (1, d) — single query retained_sets: list[list[tuple[str, str, str]]] = [] for hop_idx in range(L): candidates = candidate_sets[hop_idx] if len(candidates) == 0: retained_sets.append([]) # z stays the same — empty retained set carries no info continue # Eq. 18: compute W^ctx for this hop given z_prev W_ctx = self.get_hop_W_ctx(hop_idx, z_prev).squeeze(0) # (d, d) # Eq. 19: score all candidates relation_names = [r for (_, r, _) in candidates] scores = self.score_hop_candidates( hop_idx, W_ctx, q.squeeze(0) if q.dim() > 1 else q, relation_names ) # Eq. 2: retain if score >= θ retained = [ triple for triple, s in zip(candidates, scores.tolist()) if s >= theta ] retained_sets.append(retained) # Eq. 14: update z for next hop retained_relations = [r for (_, r, _) in retained] z_prev = self.csv([retained_relations]) # (1, d) return retained_sets