| """Factorized, contrastively trained canonical state representation probes.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import random |
|
|
| import torch |
| from torch import Tensor, nn |
| import torch.nn.functional as F |
|
|
|
|
| CANONICAL = 0 |
| CONTRADICTED = 1 |
| HISTORICAL = 2 |
| INFERRED = 3 |
|
|
|
|
| class FactorizedStateRepresentation(nn.Module): |
| """Separate entity/relation/value/metadata fields projected to a P slot.""" |
|
|
| def __init__( |
| self, |
| entities: int, |
| relations: int, |
| values: int, |
| metadata: int = 4, |
| *, |
| field_width: int = 128, |
| slot_width: int = 512, |
| ) -> None: |
| super().__init__() |
| self.entity = nn.Embedding(entities, field_width) |
| self.relation = nn.Embedding(relations, field_width) |
| self.value = nn.Embedding(values, field_width) |
| self.metadata = nn.Embedding(metadata, field_width) |
| factor_width = field_width * 4 |
| self.slot_projection = nn.Linear(factor_width, slot_width, bias=False) |
| self.tuple_projection = nn.Linear(factor_width, slot_width, bias=False) |
| self.query_projection = nn.Linear(field_width * 3, slot_width, bias=False) |
| self.entity_decoder = nn.Linear(slot_width, entities) |
| self.relation_decoder = nn.Linear(slot_width, relations) |
| self.value_decoder = nn.Linear(slot_width, values) |
| self.metadata_decoder = nn.Linear(slot_width, metadata) |
| self.temperature = nn.Parameter(torch.tensor(0.07)) |
|
|
| def factors(self, entity: Tensor, relation: Tensor, value: Tensor, metadata: Tensor) -> Tensor: |
| return torch.cat(( |
| self.entity(entity), |
| self.relation(relation), |
| self.value(value), |
| self.metadata(metadata), |
| ), dim=-1) |
|
|
| def encode(self, entity: Tensor, relation: Tensor, value: Tensor, metadata: Tensor) -> Tensor: |
| return F.normalize(self.slot_projection(self.factors(entity, relation, value, metadata)), dim=-1) |
|
|
| def tuple_anchor(self, entity: Tensor, relation: Tensor, value: Tensor, metadata: Tensor) -> Tensor: |
| return F.normalize(self.tuple_projection(self.factors(entity, relation, value, metadata)), dim=-1) |
|
|
| def query(self, entity: Tensor, relation: Tensor, metadata: Tensor) -> Tensor: |
| fields = torch.cat((self.entity(entity), self.relation(relation), self.metadata(metadata)), dim=-1) |
| return F.normalize(self.query_projection(fields), dim=-1) |
|
|
| def scores(self, query: Tensor, slots: Tensor) -> Tensor: |
| temperature = self.temperature.clamp(0.02, 1.0) |
| return torch.einsum("bd,bkd->bk", query, slots) / temperature |
|
|
| def decode(self, slots: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: |
| return ( |
| self.entity_decoder(slots), |
| self.relation_decoder(slots), |
| self.value_decoder(slots), |
| self.metadata_decoder(slots), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class CompositionConfig: |
| entities: int = 24 |
| relations: int = 3 |
| values: int = 36 |
| candidates: int = 4 |
|
|
|
|
| def is_held_out(entity: int, relation: int, value: int) -> bool: |
| return (entity * 31 + relation * 17 + value * 13) % 5 == 0 |
|
|
|
|
| def composition_splits(config: CompositionConfig): |
| train, held_out = [], [] |
| for entity in range(config.entities): |
| for relation in range(config.relations): |
| for value in range(config.values): |
| target = held_out if is_held_out(entity, relation, value) else train |
| target.append((entity, relation, value)) |
| return train, held_out |
|
|
|
|
| def _same_split_alternative(entity, relation, value, size, want_held_out, field): |
| for offset in range(1, size): |
| if field == "value": |
| candidate = (entity, relation, (value + offset) % size) |
| else: |
| candidate = ((entity + offset) % size, relation, value) |
| if is_held_out(*candidate) == want_held_out: |
| return candidate |
| raise RuntimeError("unable to construct composition-preserving negative") |
|
|
|
|
| def candidate_tuples(positive, config: CompositionConfig, *, held_out: bool): |
| entity, relation, value = positive |
| wrong_value = _same_split_alternative( |
| entity, relation, value, config.values, held_out, "value" |
| ) |
| wrong_entity = _same_split_alternative( |
| entity, relation, value, config.entities, held_out, "entity" |
| ) |
| return ( |
| (entity, relation, value, CANONICAL, "correct"), |
| (*wrong_value, CONTRADICTED, "wrong_value"), |
| (*wrong_entity, CANONICAL, "wrong_entity"), |
| (entity, relation, value, HISTORICAL, "historical"), |
| ) |
|
|
|
|
| def make_batch(combinations, config, batch_size, rng, *, held_out, permute=True): |
| selected = [combinations[rng.randrange(len(combinations))] for _ in range(batch_size)] |
| candidates, targets, kinds = [], [], [] |
| for positive in selected: |
| rows = list(candidate_tuples(positive, config, held_out=held_out)) |
| if permute: |
| rng.shuffle(rows) |
| candidates.append([row[:4] for row in rows]) |
| kinds.append([row[4] for row in rows]) |
| targets.append(next(index for index, row in enumerate(rows) if row[4] == "correct")) |
| positive = torch.tensor(selected, dtype=torch.long) |
| return positive, torch.tensor(candidates, dtype=torch.long), torch.tensor(targets), kinds |
|
|
|
|
| def representation_loss(model, positive, candidates, targets): |
| entity, relation, value = positive.T |
| flat = candidates.view(-1, 4) |
| slots = model.encode(*flat.T).view(candidates.shape[0], candidates.shape[1], -1) |
| query = model.query(entity, relation, torch.full_like(entity, CANONICAL)) |
| retrieval = F.cross_entropy(model.scores(query, slots), targets) |
| anchor = model.tuple_anchor(entity, relation, value, torch.full_like(entity, CANONICAL)) |
| contrastive = F.cross_entropy(model.scores(anchor, slots), targets) |
| decoded = model.decode(slots) |
| canonical = sum( |
| F.cross_entropy(logits.flatten(0, 1), flat[:, field]) |
| for field, logits in enumerate(decoded) |
| ) |
| return retrieval + contrastive + 0.5 * canonical |
|
|
|
|
| def evaluate_representation(model, combinations, config, *, permutations=8, seed=101): |
| model.eval() |
| totals = {"correct": 0, "wrong_value": 0, "wrong_entity": 0, "historical": 0} |
| count = 0 |
| decoded = torch.zeros(4) |
| stable = 0 |
| rng = random.Random(seed) |
| with torch.inference_mode(): |
| for positive in combinations: |
| chosen_values = [] |
| for _ in range(permutations): |
| pos, candidates, target, kinds = make_batch( |
| [positive], config, 1, rng, held_out=True, permute=True |
| ) |
| flat = candidates.view(-1, 4) |
| slots = model.encode(*flat.T).view(1, config.candidates, -1) |
| query = model.query(pos[:, 0], pos[:, 1], torch.zeros(1, dtype=torch.long)) |
| scores = model.scores(query, slots)[0] |
| correct_index = int(target[0]) |
| prediction = int(scores.argmax()) |
| selected_slot = slots[0, prediction] |
| selected_value = int(model.value_decoder(selected_slot).argmax()) |
| chosen_values.append(selected_value) |
| if prediction == correct_index and selected_value == positive[2]: |
| totals["correct"] += 1 |
| for index, kind in enumerate(kinds[0]): |
| if kind != "correct" and scores[correct_index] > scores[index]: |
| totals[kind] += 1 |
| decoded_logits = model.decode(slots[0, correct_index]) |
| truth = candidates[0, correct_index] |
| decoded += torch.tensor([ |
| int(logits.argmax() == truth[field]) |
| for field, logits in enumerate(decoded_logits) |
| ]) |
| count += 1 |
| stable += int(len(set(chosen_values)) == 1 and chosen_values[0] == positive[2]) |
| return { |
| "p_only_state_recovery": totals["correct"] / count, |
| "hard_negative_accuracy": { |
| kind: totals[kind] / count for kind in ("wrong_value", "wrong_entity", "historical") |
| }, |
| "canonical_decode_accuracy": { |
| field: float(decoded[index] / count) |
| for index, field in enumerate(("entity", "relation", "value", "metadata")) |
| }, |
| "permutation_stability": stable / len(combinations), |
| "held_out_combinations": len(combinations), |
| "permutations_per_combination": permutations, |
| } |
|
|
|
|
| def train_and_probe_representation( |
| *, steps=600, batch_size=64, slot_width=512, seed=97, evaluation_limit=256 |
| ): |
| torch.manual_seed(seed) |
| config = CompositionConfig() |
| train, held_out = composition_splits(config) |
| model = FactorizedStateRepresentation( |
| config.entities, config.relations, config.values, slot_width=slot_width |
| ) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3) |
| rng = random.Random(seed) |
| losses = [] |
| model.train() |
| for _ in range(steps): |
| positive, candidates, targets, _ = make_batch( |
| train, config, batch_size, rng, held_out=False |
| ) |
| loss = representation_loss(model, positive, candidates, targets) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| losses.append(float(loss.detach())) |
| probe = evaluate_representation( |
| model, held_out[:evaluation_limit], config, permutations=8, seed=seed + 1 |
| ) |
| probe.update({ |
| "training_loss_first": losses[0], |
| "training_loss_last": losses[-1], |
| "train_combinations": len(train), |
| "total_held_out_combinations": len(held_out), |
| "slot_width": slot_width, |
| }) |
| return model, config, probe |
|
|
