| """Norm-bounded algebraic graph reads around an otherwise frozen LM.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from strata.modeling.algebra import BoundedRelativeReZero |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class ComposeLMOutput: |
| logits: torch.Tensor |
| loss: torch.Tensor | None |
| hidden_states: torch.Tensor |
| graph_update: torch.Tensor |
| relative_update_rms: torch.Tensor |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class SparseAlgebraGraphRead: |
| """CSR-like graph values addressed to selected batch/token positions.""" |
|
|
| values: torch.Tensor |
| batch_indices: torch.Tensor |
| token_indices: torch.Tensor |
| reliability: torch.Tensor |
|
|
| def validate(self, hidden: torch.Tensor, graph_dim: int) -> None: |
| count = self.values.shape[0] |
| if self.values.ndim != 2 or self.values.shape[1] != graph_dim: |
| raise ValueError("sparse graph values must have shape [reads, graph_dim]") |
| if self.batch_indices.shape != (count,) or self.token_indices.shape != (count,): |
| raise ValueError("sparse graph indices must have shape [reads]") |
| if self.reliability.shape != (count,): |
| raise ValueError("sparse graph reliability must have shape [reads]") |
| if count and ( |
| int(self.batch_indices.min()) < 0 |
| or int(self.batch_indices.max()) >= hidden.shape[0] |
| or int(self.token_indices.min()) < 0 |
| or int(self.token_indices.max()) >= hidden.shape[1] |
| ): |
| raise ValueError("sparse graph index is out of range") |
|
|
|
|
| class AlgebraGraphReadAdapter(nn.Module): |
| """Project exact graph values and inject them through the fixed BRR bound.""" |
|
|
| def __init__(self, graph_dim: int, hidden_dim: int, *, gamma_max: float = 0.02) -> None: |
| super().__init__() |
| self.graph_projection = nn.Linear(graph_dim, hidden_dim, bias=False) |
| self.brr = BoundedRelativeReZero(hidden_dim, gamma_max=gamma_max, per_channel=False) |
|
|
| def forward( |
| self, |
| hidden: torch.Tensor, |
| graph_values: torch.Tensor, |
| reliability: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if graph_values.shape[:-1] != hidden.shape[:-1]: |
| raise ValueError("graph values must align with token hidden states") |
| if reliability.shape != hidden.shape[:-1]: |
| raise ValueError("reliability must have shape [batch, sequence]") |
| active = reliability != 0 |
| if not bool(active.any()): |
| return hidden, torch.zeros_like(hidden), hidden.new_zeros(hidden.shape[:-1]) |
| projection_dtype = self.graph_projection.weight.dtype |
| if bool(active.all()): |
| read = self.graph_projection(graph_values.to(projection_dtype)).to(hidden.dtype) |
| output = self.brr(hidden, read, reliability) |
| relative = (output.update_rms / output.hidden_rms.clamp_min(1e-8)).squeeze(-1) |
| return output.hidden, output.update, relative |
|
|
| active_hidden = hidden[active] |
| active_read = self.graph_projection(graph_values[active].to(projection_dtype)).to(hidden.dtype) |
| active_output = self.brr(active_hidden, active_read, reliability[active]) |
| updated = hidden.clone() |
| updated[active] = active_output.hidden |
| update = torch.zeros_like(hidden) |
| update[active] = active_output.update |
| relative = hidden.new_zeros(hidden.shape[:-1]) |
| relative[active] = ( |
| active_output.update_rms / active_output.hidden_rms.clamp_min(1e-8) |
| ).squeeze(-1) |
| return updated, update, relative |
|
|
| def forward_sparse( |
| self, |
| hidden: torch.Tensor, |
| read: SparseAlgebraGraphRead, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| read.validate(hidden, self.graph_projection.in_features) |
| if read.values.shape[0] == 0: |
| return hidden, torch.zeros_like(hidden), hidden.new_zeros(hidden.shape[:-1]) |
| batch = read.batch_indices.to(device=hidden.device, dtype=torch.long) |
| token = read.token_indices.to(device=hidden.device, dtype=torch.long) |
| selected_hidden = hidden[batch, token] |
| projection_dtype = self.graph_projection.weight.dtype |
| selected_read = self.graph_projection( |
| read.values.to(device=hidden.device, dtype=projection_dtype) |
| ).to(hidden.dtype) |
| selected_output = self.brr( |
| selected_hidden, |
| selected_read, |
| read.reliability.to(device=hidden.device, dtype=hidden.dtype), |
| ) |
| updated = hidden.clone() |
| updated[batch, token] = selected_output.hidden |
| update = torch.zeros_like(hidden) |
| update[batch, token] = selected_output.update |
| relative = hidden.new_zeros(hidden.shape[:-1]) |
| relative[batch, token] = ( |
| selected_output.update_rms / selected_output.hidden_rms.clamp_min(1e-8) |
| ).squeeze(-1) |
| return updated, update, relative |
|
|
|
|
| class FrozenLMWithAlgebraRead(nn.Module): |
| """Attach graph reads after a frozen decoder without changing its base path. |
| |
| The wrapped model must expose ``forward_hidden``, ``final_norm``, and |
| ``lm_head``. When graph reads are disabled, the adapter is not called. |
| """ |
|
|
| def __init__(self, base_model: nn.Module, adapter: AlgebraGraphReadAdapter) -> None: |
| super().__init__() |
| self.base_model = base_model |
| self.adapter = adapter |
| for parameter in self.base_model.parameters(): |
| parameter.requires_grad_(False) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| *, |
| attention_mask: torch.Tensor | None = None, |
| labels: torch.Tensor | None = None, |
| graph_values: torch.Tensor | None = None, |
| graph_reliability: torch.Tensor | None = None, |
| graph_enabled: bool = False, |
| ) -> ComposeLMOutput: |
| hidden, _mask = self.base_model.forward_hidden(input_ids, attention_mask=attention_mask) |
| update = torch.zeros_like(hidden) |
| relative = hidden.new_zeros(hidden.shape[:-1]) |
| if graph_enabled: |
| if graph_values is None or graph_reliability is None: |
| raise ValueError("enabled graph reads require values and reliability") |
| hidden, update, relative = self.adapter(hidden, graph_values, graph_reliability) |
| normalized = self.base_model.final_norm(hidden) |
| logits = self.base_model.lm_head(normalized) |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits[:, :-1].contiguous().view(-1, logits.shape[-1]), |
| labels[:, 1:].contiguous().view(-1), |
| ignore_index=-100, |
| ) |
| return ComposeLMOutput(logits, loss, normalized, update, relative) |
|
|
|
|
| __all__ = [ |
| "AlgebraGraphReadAdapter", |
| "ComposeLMOutput", |
| "FrozenLMWithAlgebraRead", |
| "SparseAlgebraGraphRead", |
| ] |
|
|
