| """PAT: query-conditioned, relation-aware Path-Aware Transformer.""" |
|
|
| from __future__ import annotations |
|
|
| import itertools |
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass |
| from enum import IntEnum |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional |
| from torch.utils.checkpoint import checkpoint |
|
|
| from graph.config import PATConfig |
| from graph.types import ConstraintState, GraphPath |
| from graph.utils import qk_normalize, soft_cap, stable_hash |
|
|
|
|
| class PATTokenType(IntEnum): |
| QUERY = 0 |
| ENTITY = 1 |
| RELATION = 2 |
| EDGE = 3 |
| TYPE = 4 |
| PATH_START = 5 |
| PATH_STEP = 6 |
| PATH_END = 7 |
| HOP = 8 |
| CONSTRAINT = 9 |
|
|
|
|
| @dataclass(frozen=True) |
| class PATInput: |
| token_features: Tensor |
| token_type_ids: Tensor |
| hop_positions: Tensor |
| attention_mask: Tensor |
|
|
|
|
| @dataclass(frozen=True) |
| class PATOutput: |
| path_states: Tensor |
| relation_sequence_states: Tensor |
| transition_states: Tensor |
| compatibility: Tensor |
| token_states: Tensor |
| endpoint_states: Tensor |
| hop_states: Tensor |
| constraint_compatibility: Tensor |
|
|
|
|
| class RelationAwareAttention(nn.Module): |
| def __init__( |
| self, |
| hidden_dim: int, |
| num_heads: int, |
| num_token_types: int, |
| dropout: float, |
| logit_cap: float = 0.0, |
| qk_norm: bool = False, |
| ) -> None: |
| super().__init__() |
| if hidden_dim % num_heads: |
| raise ValueError("PAT hidden_dim must be divisible by num_heads") |
| self.logit_cap = logit_cap |
| self.qk_norm = qk_norm |
| self.hidden_dim = hidden_dim |
| self.num_heads = num_heads |
| self.head_dim = hidden_dim // num_heads |
| self.q_proj = nn.Linear(hidden_dim, hidden_dim) |
| self.k_proj = nn.Linear(hidden_dim, hidden_dim) |
| self.v_proj = nn.Linear(hidden_dim, hidden_dim) |
| self.out_proj = nn.Linear(hidden_dim, hidden_dim) |
| self.pair_bias = nn.Parameter( |
| torch.zeros(num_token_types, num_token_types, num_heads) |
| ) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward( |
| self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor |
| ) -> Tensor: |
| batch, length, _ = states.shape |
| q = ( |
| self.q_proj(states) |
| .view(batch, length, self.num_heads, self.head_dim) |
| .transpose(1, 2) |
| ) |
| k = ( |
| self.k_proj(states) |
| .view(batch, length, self.num_heads, self.head_dim) |
| .transpose(1, 2) |
| ) |
| v = ( |
| self.v_proj(states) |
| .view(batch, length, self.num_heads, self.head_dim) |
| .transpose(1, 2) |
| ) |
| if self.qk_norm: |
| q, k = qk_normalize(q), qk_normalize(k) |
| logits = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim**0.5) |
| pair = self.pair_bias[ |
| token_type_ids.unsqueeze(-1), token_type_ids.unsqueeze(-2) |
| ] |
| logits = logits + pair.permute(0, 3, 1, 2) |
| query_key = query_state.view(batch, self.num_heads, self.head_dim) |
| logits = logits + torch.einsum("bhd,bhld->bhl", query_key, k).unsqueeze(-2) / ( |
| self.head_dim**0.5 |
| ) |
| logits = soft_cap(logits, self.logit_cap) |
| valid = mask.to(torch.bool) |
| logits = logits.masked_fill( |
| ~valid[:, None, None, :], torch.finfo(logits.dtype).min |
| ) |
| weights = torch.softmax(logits, dim=-1).to(v.dtype) |
| weights = self.dropout(weights) |
| result = ( |
| torch.matmul(weights, v) |
| .transpose(1, 2) |
| .contiguous() |
| .view(batch, length, self.hidden_dim) |
| ) |
| return self.out_proj(result) |
|
|
|
|
| class PATBlock(nn.Module): |
| def __init__( |
| self, config: PATConfig, logit_cap: float = 0.0, qk_norm: bool = False |
| ) -> None: |
| super().__init__() |
| self.attn = RelationAwareAttention( |
| config.hidden_dim, |
| config.num_heads, |
| config.num_token_types, |
| config.dropout, |
| logit_cap, |
| qk_norm, |
| ) |
| self.norm1 = nn.LayerNorm(config.hidden_dim) |
| self.norm2 = nn.LayerNorm(config.hidden_dim) |
| self.ff = nn.Sequential( |
| nn.Linear(config.hidden_dim, config.hidden_dim * 4), |
| nn.GELU(), |
| nn.Dropout(config.dropout), |
| nn.Linear(config.hidden_dim * 4, config.hidden_dim), |
| ) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward( |
| self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor |
| ) -> Tensor: |
| attended = self.attn(self.norm1(states), token_type_ids, mask, query_state) |
| states = states + self.dropout(attended) |
| return states + self.dropout(self.ff(self.norm2(states))) |
|
|
|
|
| class PathAwareTransformer(nn.Module): |
| """PAT operates on structured vectors, never on a textual graph dump.""" |
|
|
| def __init__( |
| self, |
| config: PATConfig, |
| gradient_checkpointing: bool = False, |
| logit_cap: float = 0.0, |
| qk_norm: bool = False, |
| ) -> None: |
| super().__init__() |
| self.config = config |
| self.gradient_checkpointing = gradient_checkpointing |
| self.token_type_embedding = nn.Embedding( |
| config.num_token_types, config.hidden_dim |
| ) |
| self.path_position_embedding = nn.Embedding( |
| config.max_path_length * 5 + 5, config.hidden_dim |
| ) |
| self.hop_embedding = nn.Embedding(config.max_path_length + 2, config.hidden_dim) |
| self.type_embedding = nn.Embedding(config.type_vocab_size, config.hidden_dim) |
| self.attribute_embedding = nn.Embedding( |
| config.attribute_vocab_size, config.hidden_dim |
| ) |
| self.query_gate = nn.Linear(config.hidden_dim, config.hidden_dim) |
| self.blocks = nn.ModuleList( |
| [PATBlock(config, logit_cap, qk_norm) for _ in range(config.num_layers)] |
| ) |
| self.final_norm = nn.LayerNorm(config.hidden_dim) |
| self.query_projection = nn.Linear(config.hidden_dim, config.hidden_dim) |
|
|
| def forward(self, inputs: PATInput, query_state: Tensor) -> PATOutput: |
| token_types = inputs.token_type_ids.clamp(0, self.config.num_token_types - 1) |
| positions = torch.arange( |
| inputs.token_features.shape[1], device=inputs.token_features.device |
| ) |
| positions = positions.clamp_max(self.config.max_path_length * 5 + 4).unsqueeze( |
| 0 |
| ) |
| hops = inputs.hop_positions.clamp(0, self.config.max_path_length + 1) |
| states = inputs.token_features + self.token_type_embedding(token_types) |
| if query_state.ndim == 1: |
| query_state = query_state.unsqueeze(0) |
| if query_state.shape[0] == 1 and states.shape[0] != 1: |
| |
| query_state = query_state.expand(states.shape[0], -1).contiguous() |
| states = ( |
| states + self.path_position_embedding(positions) + self.hop_embedding(hops) |
| ) |
| attention_query = ( |
| query_state |
| if self.config.query_conditioning |
| else torch.zeros_like(query_state) |
| ) |
| if self.config.query_conditioning: |
| states = states + self.query_gate(query_state).unsqueeze(1) |
| for block in self.blocks: |
| if self.gradient_checkpointing and self.training: |
| states = checkpoint( |
| lambda values, layer=block: layer( |
| values, token_types, inputs.attention_mask, attention_query |
| ), |
| states, |
| use_reentrant=False, |
| ) |
| else: |
| states = block( |
| states, token_types, inputs.attention_mask, attention_query |
| ) |
| states = self.final_norm(states) |
| mask = inputs.attention_mask.to(states.dtype).unsqueeze(-1) |
| path_states = (states * mask).sum(1) / mask.sum(1).clamp_min(1.0) |
| path_states = functional.normalize(path_states, dim=-1) |
| relation_mask = ( |
| token_types == int(PATTokenType.RELATION) |
| ) & inputs.attention_mask.bool() |
| max_relations = max(1, int(relation_mask.sum(dim=1).max().item())) |
| |
| |
| rows, columns = relation_mask.nonzero(as_tuple=True) |
| slots = relation_mask.long().cumsum(1)[rows, columns] - 1 |
| relation_sequence_states = states.new_zeros( |
| (states.shape[0], max_relations, states.shape[-1]) |
| ).index_put((rows, slots), states[rows, columns]) |
| transition_states = relation_sequence_states |
| compatibility = ( |
| path_states |
| * functional.normalize(self.query_projection(attention_query), dim=-1) |
| ).sum(-1) |
| entity_mask = ( |
| token_types == int(PATTokenType.ENTITY) |
| ) & inputs.attention_mask.bool() |
| entity_positions = ( |
| torch.arange(states.shape[1], device=states.device) |
| .unsqueeze(0) |
| .expand(states.shape[0], -1) |
| ) |
| last_entity = ( |
| torch.where( |
| entity_mask, entity_positions, torch.zeros_like(entity_positions) |
| ) |
| .max(dim=1) |
| .values |
| ) |
| batch_indices = torch.arange(states.shape[0], device=states.device) |
| endpoint_states = states[batch_indices, last_entity] |
| hop_states = relation_sequence_states.mean(1) |
| constraint_compatibility = functional.cosine_similarity( |
| path_states, self.query_projection(attention_query), dim=-1 |
| ) |
| return PATOutput( |
| path_states, |
| relation_sequence_states, |
| transition_states, |
| compatibility, |
| states, |
| endpoint_states, |
| hop_states, |
| constraint_compatibility, |
| ) |
|
|
| def build_input( |
| self, |
| paths: Sequence[GraphPath], |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| query_state: Tensor, |
| constraint_state: ConstraintState | None = None, |
| ) -> PATInput: |
| if query_state.ndim == 2: |
| query_state = query_state[0] |
| if not paths: |
| raise ValueError("PAT requires at least one candidate path") |
| feature_rows: list[Tensor] = [] |
| type_rows: list[Tensor] = [] |
| hop_rows: list[Tensor] = [] |
| for path in paths: |
| if not path.steps: |
| raise ValueError("PAT cannot encode an empty path") |
| features = [torch.zeros_like(query_state)] |
| types = [int(PATTokenType.PATH_START)] |
| hops = [0] |
| if constraint_state is not None: |
| constraint_signal = float( |
| constraint_state.max_hops - constraint_state.min_hops |
| ) |
| if constraint_state.exact_hops is not None: |
| constraint_signal += float(constraint_state.exact_hops) |
| features.append( |
| query_state |
| * (1.0 + constraint_signal / max(1.0, self.config.max_path_length)) |
| ) |
| types.append(int(PATTokenType.CONSTRAINT)) |
| hops.append(0) |
| first = path.steps[0] |
| features.append(node_states[first.source_id]) |
| types.append(int(PATTokenType.ENTITY)) |
| hops.append(0) |
| features.append(self._type_state(first.source_type, query_state)) |
| types.append(int(PATTokenType.TYPE)) |
| hops.append(0) |
| for hop, step in enumerate(path.steps, start=1): |
| features.append( |
| relation_states.get(step.relation_id, torch.zeros_like(query_state)) |
| ) |
| types.append(int(PATTokenType.RELATION)) |
| hops.append(hop) |
| features.append( |
| self._attribute_state(step.edge_attributes, query_state) |
| ) |
| types.append(int(PATTokenType.EDGE)) |
| hops.append(hop) |
| features.append(node_states[step.target_id]) |
| types.append(int(PATTokenType.ENTITY)) |
| hops.append(hop) |
| features.append(self._type_state(step.target_type, query_state)) |
| types.append(int(PATTokenType.TYPE)) |
| hops.append(hop) |
| features.append( |
| query_state * (float(hop) / max(1.0, self.config.max_path_length)) |
| ) |
| types.append(int(PATTokenType.HOP)) |
| hops.append(hop) |
| features.append(torch.zeros_like(query_state)) |
| types.append(int(PATTokenType.PATH_END)) |
| hops.append(len(path.steps)) |
| feature_rows.append(torch.stack(features)) |
| type_rows.append( |
| torch.tensor(types, dtype=torch.long, device=query_state.device) |
| ) |
| hop_rows.append( |
| torch.tensor(hops, dtype=torch.long, device=query_state.device) |
| ) |
| length = max(row.shape[0] for row in feature_rows) |
| hidden = query_state.shape[-1] |
| padded_features = torch.zeros( |
| (len(paths), length, hidden), |
| dtype=query_state.dtype, |
| device=query_state.device, |
| ) |
| padded_types = torch.full( |
| (len(paths), length), |
| int(PATTokenType.PATH_END), |
| dtype=torch.long, |
| device=query_state.device, |
| ) |
| padded_hops = torch.zeros( |
| (len(paths), length), dtype=torch.long, device=query_state.device |
| ) |
| padded_mask = torch.zeros( |
| (len(paths), length), dtype=torch.long, device=query_state.device |
| ) |
| for i, (feature_row, type_row, hop_row) in enumerate( |
| zip(feature_rows, type_rows, hop_rows, strict=True) |
| ): |
| size = feature_row.shape[0] |
| padded_features[i, :size] = feature_row |
| padded_types[i, :size] = type_row |
| padded_hops[i, :size] = hop_row |
| padded_mask[i, :size] = 1 |
| return PATInput(padded_features, padded_types, padded_hops, padded_mask) |
|
|
| def build_batch_input( |
| self, |
| paths: Sequence[GraphPath], |
| owners: Sequence[int], |
| node_rows: Sequence[Mapping[str, int]], |
| relation_rows: Sequence[Mapping[str, int]], |
| node_table: Tensor, |
| relation_table: Tensor, |
| query_states: Tensor, |
| constraints: Sequence[ConstraintState | None], |
| ) -> PATInput: |
| """:meth:`build_input` for paths drawn from many examples at once. |
| |
| Path ``i`` belongs to example ``owners[i]``, whose node and relation |
| states are rows ``node_rows[owner][id]`` of ``node_table`` and |
| ``relation_rows[owner][id]`` of ``relation_table``. Every token is a row |
| of one feature table plus a multiple of its example's query state, so the |
| whole batch is assembled by a single gather instead of one small tensor |
| operation per token. Token for token the result matches |
| :meth:`build_input`, padded to the batch's longest path. |
| """ |
| if not paths: |
| raise ValueError("PAT requires at least one candidate path") |
| max_length = max(1.0, self.config.max_path_length) |
| type_ids: dict[int, int] = {} |
| attribute_sets: dict[tuple[int, ...], int] = {} |
| node_offset = 1 |
| relation_offset = node_offset + node_table.shape[0] |
| |
| |
| token_rows: list[list[int | tuple[str, int]]] = [] |
| token_scales: list[list[float]] = [] |
| token_types: list[list[int]] = [] |
| token_hops: list[list[int]] = [] |
|
|
| def type_slot(type_name: str | None) -> tuple[str, int]: |
| index = stable_hash(type_name or "unknown", self.config.type_vocab_size) |
| return ("t", type_ids.setdefault(index, len(type_ids))) |
|
|
| def attribute_slot(attributes: Mapping[str, str]) -> int | tuple[str, int]: |
| if not attributes: |
| return 0 |
| key = tuple( |
| stable_hash(f"{name}={value}", self.config.attribute_vocab_size) |
| for name, value in sorted(attributes.items()) |
| ) |
| return ("a", attribute_sets.setdefault(key, len(attribute_sets))) |
|
|
| for path, owner in zip(paths, owners, strict=True): |
| if not path.steps: |
| raise ValueError("PAT cannot encode an empty path") |
| nodes, relations = node_rows[owner], relation_rows[owner] |
| rows: list[int | tuple[str, int]] = [0] |
| scales = [0.0] |
| types = [int(PATTokenType.PATH_START)] |
| hops = [0] |
| constraint = constraints[owner] |
| if constraint is not None: |
| signal = float(constraint.max_hops - constraint.min_hops) |
| if constraint.exact_hops is not None: |
| signal += float(constraint.exact_hops) |
| rows.append(0) |
| scales.append(1.0 + signal / max_length) |
| types.append(int(PATTokenType.CONSTRAINT)) |
| hops.append(0) |
| first = path.steps[0] |
| rows += [node_offset + nodes[first.source_id], type_slot(first.source_type)] |
| scales += [0.0, 0.0] |
| types += [int(PATTokenType.ENTITY), int(PATTokenType.TYPE)] |
| hops += [0, 0] |
| for hop, step in enumerate(path.steps, start=1): |
| relation = relations.get(step.relation_id) |
| rows += [ |
| 0 if relation is None else relation_offset + relation, |
| attribute_slot(step.edge_attributes), |
| node_offset + nodes[step.target_id], |
| type_slot(step.target_type), |
| 0, |
| ] |
| scales += [0.0, 0.0, 0.0, 0.0, float(hop) / max_length] |
| types += [ |
| int(PATTokenType.RELATION), |
| int(PATTokenType.EDGE), |
| int(PATTokenType.ENTITY), |
| int(PATTokenType.TYPE), |
| int(PATTokenType.HOP), |
| ] |
| hops += [hop] * 5 |
| rows.append(0) |
| scales.append(0.0) |
| types.append(int(PATTokenType.PATH_END)) |
| hops.append(len(path.steps)) |
| token_rows.append(rows) |
| token_scales.append(scales) |
| token_types.append(types) |
| token_hops.append(hops) |
|
|
| device = query_states.device |
| type_offset = relation_offset + relation_table.shape[0] |
| attribute_offset = type_offset + len(type_ids) |
| parts = [ |
| query_states.new_zeros((1, query_states.shape[-1])), |
| node_table.to(query_states.dtype), |
| relation_table.to(query_states.dtype), |
| ] |
| if type_ids: |
| parts.append( |
| self.type_embedding( |
| torch.tensor(list(type_ids), dtype=torch.long, device=device) |
| ).to(query_states.dtype) |
| ) |
| if attribute_sets: |
| flat = [index for key in attribute_sets for index in key] |
| offsets = torch.tensor( |
| [0, *itertools.accumulate(len(key) for key in attribute_sets)][:-1], |
| dtype=torch.long, |
| device=device, |
| ) |
| parts.append( |
| functional.embedding_bag( |
| torch.tensor(flat, dtype=torch.long, device=device), |
| self.attribute_embedding.weight, |
| offsets, |
| mode="mean", |
| ).to(query_states.dtype) |
| ) |
| table = torch.cat(parts) |
|
|
| length = max(len(rows) for rows in token_rows) |
| count = len(paths) |
| index = torch.zeros((count, length), dtype=torch.long) |
| scale = torch.zeros((count, length), dtype=query_states.dtype) |
| padded_types = torch.full( |
| (count, length), int(PATTokenType.PATH_END), dtype=torch.long |
| ) |
| padded_hops = torch.zeros((count, length), dtype=torch.long) |
| padded_mask = torch.zeros((count, length), dtype=torch.long) |
| for row, (rows, scales, types, hops) in enumerate( |
| zip(token_rows, token_scales, token_types, token_hops, strict=True) |
| ): |
| size = len(rows) |
| index[row, :size] = torch.tensor( |
| [ |
| slot |
| if isinstance(slot, int) |
| else (type_offset if slot[0] == "t" else attribute_offset) + slot[1] |
| for slot in rows |
| ], |
| dtype=torch.long, |
| ) |
| scale[row, :size] = torch.tensor(scales, dtype=query_states.dtype) |
| padded_types[row, :size] = torch.tensor(types, dtype=torch.long) |
| padded_hops[row, :size] = torch.tensor(hops, dtype=torch.long) |
| padded_mask[row, :size] = 1 |
| index, scale = index.to(device), scale.to(device) |
| owner_index = torch.tensor(list(owners), dtype=torch.long, device=device) |
| features = table[index] + scale.unsqueeze(-1) * query_states[ |
| owner_index |
| ].unsqueeze(1) |
| return PATInput( |
| features, |
| padded_types.to(device), |
| padded_hops.to(device), |
| padded_mask.to(device), |
| ) |
|
|
| def _type_state(self, type_name: str | None, query_state: Tensor) -> Tensor: |
| index = stable_hash(type_name or "unknown", self.config.type_vocab_size) |
| return self.type_embedding(torch.tensor(index, device=query_state.device)) |
|
|
| def _attribute_state( |
| self, attributes: Mapping[str, str], query_state: Tensor |
| ) -> Tensor: |
| if not attributes: |
| return torch.zeros_like(query_state) |
| indices = [ |
| stable_hash(f"{key}={value}", self.config.attribute_vocab_size) |
| for key, value in sorted(attributes.items()) |
| ] |
| ids = torch.tensor(indices, dtype=torch.long, device=query_state.device) |
| return self.attribute_embedding(ids).mean(0) |
|
|
| def encode_paths( |
| self, |
| paths: Sequence[GraphPath], |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| query_state: Tensor, |
| constraint_state: ConstraintState | None = None, |
| ) -> PATOutput: |
| return self( |
| self.build_input( |
| paths, node_states, relation_states, query_state, constraint_state |
| ), |
| query_state if query_state.ndim == 2 else query_state.unsqueeze(0), |
| ) |
|
|
