ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Extract reusable higher-level patterns from toroidal aggregates.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class AbstractionEngine(nn.Module): | |
| def __init__(self, d_model: int = 256, abstraction_capacity: int = 64) -> None: | |
| super().__init__() | |
| self.d_model = d_model | |
| self.abstraction_capacity = abstraction_capacity | |
| self.register_buffer("abstraction_memory", torch.zeros(abstraction_capacity, d_model)) | |
| self.register_buffer("abstraction_usage", torch.zeros(abstraction_capacity)) | |
| self.register_buffer("abstraction_age", torch.zeros(abstraction_capacity, dtype=torch.long)) | |
| def compute_pattern(self, aggregates: list[dict]) -> torch.Tensor | None: | |
| if not aggregates: | |
| return None | |
| # Structural signature: retain primitive atom information without a | |
| # learned feature projection. Normalize each property before fusion. | |
| parts = [] | |
| for key in ("r", "phi", "omega", "E", "kappa", "M", "tau"): | |
| x = torch.stack([a[key] for a in aggregates]).mean(dim=0) | |
| parts.append(F.normalize(x, dim=-1) if key in ("r", "M") else x) | |
| pattern = sum(parts) / len(parts) | |
| return F.normalize(pattern, dim=-1) | |
| def match_abstraction(self, pattern: torch.Tensor) -> tuple[int, float]: | |
| if self.abstraction_capacity == 0: | |
| return -1, -1.0 | |
| p = F.normalize(pattern.reshape(1, -1), dim=-1) | |
| mem = F.normalize(self.abstraction_memory, dim=-1) | |
| similarities = (p @ mem.T).squeeze(0) | |
| # Empty slots are never selected as existing abstractions. | |
| similarities = similarities.masked_fill(self.abstraction_usage <= 0, -1.0) | |
| idx = similarities.argmax().item() | |
| return idx, similarities[idx].item() | |
| def create_or_update_abstraction(self, aggregates: list[dict], threshold: float = 0.8) -> int: | |
| pattern = self.compute_pattern(aggregates) | |
| if pattern is None: | |
| return -1 | |
| best_idx, best_sim = self.match_abstraction(pattern) | |
| with torch.no_grad(): | |
| if best_idx >= 0 and best_sim >= threshold: | |
| self.abstraction_memory[best_idx].mul_(0.9).add_(0.1 * pattern) | |
| self.abstraction_usage[best_idx].add_(1) | |
| self.abstraction_age[best_idx].zero_() | |
| return best_idx | |
| free = torch.nonzero(self.abstraction_usage <= 0, as_tuple=False).flatten() | |
| if free.numel(): | |
| idx = free[0].item() | |
| else: | |
| # Replace the least-used/oldest representation, never overwrite all memory. | |
| score = self.abstraction_usage + 0.01 * self.abstraction_age.float() | |
| idx = score.argmin().item() | |
| self.abstraction_memory[idx].copy_(pattern) | |
| self.abstraction_usage[idx] = 1 | |
| self.abstraction_age[idx] = 0 | |
| return idx | |
| def step_age(self) -> None: | |
| with torch.no_grad(): | |
| self.abstraction_age.add_(1) | |
| def get_active_abstractions(self, min_similarity: float = 0.5) -> list[dict]: | |
| active = torch.nonzero(self.abstraction_usage > 0, as_tuple=False).flatten().tolist() | |
| return [{"id": i, "pattern": self.abstraction_memory[i].detach().cpu()} for i in active] | |