nested-learning 0.2.0__py3-none-any.whl

nested_learning/optim/deep.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class DeepMomentumState:
+    grad_avg: Optional[torch.Tensor] = None
+    sq_avg: Optional[torch.Tensor] = None
+
+
+class DeepMomentum(nn.Module):
+    """Implements momentum variants described in the NL paper."""
+
+    def __init__(
+        self,
+        *,
+        beta: float = 0.9,
+        beta2: float = 0.999,
+        eps: float = 1e-8,
+        variant: str = "preconditioned",
+    ) -> None:
+        super().__init__()
+        self.beta = beta
+        self.beta2 = beta2
+        self.eps = eps
+        self.variant = variant
+        self.state: dict[str, DeepMomentumState] = {}
+        self.nonlinearity = nn.Tanh() if variant in {"dmgd", "muon"} else nn.Identity()
+        self.last_metrics: dict[str, float] = {}
+
+    def reset_state(self) -> None:
+        self.state.clear()
+
+    def _precondition(self, grad: torch.Tensor, state: DeepMomentumState) -> torch.Tensor:
+        if state.sq_avg is None or state.sq_avg.shape != grad.shape:
+            state.sq_avg = torch.zeros_like(grad)
+        state.sq_avg.mul_(self.beta2).addcmul_(grad, grad, value=1 - self.beta2)
+        denom = state.sq_avg.sqrt().add_(self.eps)
+        return grad / denom
+
+    def _nl_precondition(
+        self,
+        grad: torch.Tensor,
+        context: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, dict[str, float]]:
+        metrics: dict[str, float] = {
+            "ctx_norm": 0.0,
+            "proj_norm": 0.0,
+            "proj_skipped": 0.0,
+        }
+        if context is None:
+            return grad, metrics
+        ctx = context
+        if ctx.ndim > 1:
+            ctx = ctx.reshape(-1, ctx.shape[-1]).mean(dim=0)
+        ctx_norm = torch.norm(ctx)
+        metrics["ctx_norm"] = ctx_norm.item()
+
+        if ctx_norm > 0:
+            if grad.ndim == 0 or grad.shape[-1] != ctx.shape[-1]:
+                metrics["proj_skipped"] = 1.0
+                return grad, metrics
+            unit = ctx / (ctx_norm + self.eps)
+            # Project grad orthogonal to context (rank-1 projector).
+            projection = (grad * unit).sum(dim=-1, keepdim=True) * unit
+            update = grad - projection
+            metrics["proj_norm"] = torch.norm(update).item()
+            return update, metrics
+        return grad, metrics
+
+    def forward(  # type: ignore[override]
+        self,
+        grad: torch.Tensor,
+        *,
+        context: torch.Tensor | None = None,
+        param_key: str | None = None,
+    ) -> torch.Tensor:
+        key = param_key or "__default__"
+        state = self.state.get(key)
+        if state is None:
+            state = DeepMomentumState()
+            self.state[key] = state
+        if state.grad_avg is None or state.grad_avg.shape != grad.shape:
+            state.grad_avg = torch.zeros_like(grad)
+        self.last_metrics = {}
+        update = grad
+        if self.variant in {"preconditioned", "muon"}:
+            update = self._precondition(grad, state)
+        if self.variant == "l2_objective":
+            update = grad + 0.1 * torch.mean(grad, dim=-1, keepdim=True)
+        if self.variant == "nl_l2_precond":
+            update, metrics = self._nl_precondition(grad, context)
+            self.last_metrics.update(metrics)
+        if self.variant in {"dmgd", "muon"}:
+            update = self.nonlinearity(update)
+        state.grad_avg.mul_(self.beta).add_(update, alpha=1 - self.beta)
+        return state.grad_avg

nested_learning/optim/factory.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from .deep import DeepMomentum
+
+
+def build_optimizer(config: Dict[str, Any]) -> DeepMomentum:
+    opt_type = config.get("type", "deep_momentum").lower()
+    if opt_type != "deep_momentum":
+        raise ValueError(f"Unsupported optimizer type {opt_type}")
+    params = config.get("params", {})
+    return DeepMomentum(**params)

nested_learning/optim/m3.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+
+
+def _newton_schulz(matrix: torch.Tensor, steps: int, eps: float = 1e-6) -> torch.Tensor:
+    if matrix.ndim != 2:
+        raise ValueError("Newton-Schulz expects a 2D matrix")
+    dtype = matrix.dtype
+    device = matrix.device
+    m, n = matrix.shape
+    x = matrix
+    norm = torch.linalg.norm(x)
+    x = x / (norm + eps)
+    eye = torch.eye(n, device=device, dtype=dtype)
+    for _ in range(steps):
+        x = 0.5 * x @ (3.0 * eye - x.T @ x)
+    return x
+
+
+def _orthogonalize(tensor: torch.Tensor, steps: int, eps: float) -> torch.Tensor:
+    if tensor.ndim < 2:
+        return tensor
+    mat = tensor.reshape(tensor.shape[0], -1)
+    ortho = _newton_schulz(mat, steps=steps, eps=eps)
+    return ortho.reshape_as(tensor)
+
+
+class M3(torch.optim.Optimizer):
+    """
+    Multi-scale Momentum Muon (M3) optimizer (Nested Learning paper, Algorithm 1).
+
+    This is a paper-faithful implementation for 2D weight tensors:
+      - M1: fast momentum
+      - M2: slow momentum (updated every `slow_chunk` steps)
+      - V: second moment
+      - O1/O2: Newton-Schulz orthogonalized momenta
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        *,
+        lr: float = 1e-3,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        beta3: float = 0.9,
+        alpha: float = 1.0,
+        eps: float = 1e-8,
+        ns_steps: int = 3,
+        slow_chunk: int = 100,
+        weight_decay: float = 0.0,
+    ) -> None:
+        defaults = dict(
+            lr=lr,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            alpha=alpha,
+            eps=eps,
+            ns_steps=ns_steps,
+            slow_chunk=slow_chunk,
+            weight_decay=weight_decay,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):  # type: ignore[override]
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            lr = group["lr"]
+            beta1 = group["beta1"]
+            beta2 = group["beta2"]
+            beta3 = group["beta3"]
+            alpha = group["alpha"]
+            eps = group["eps"]
+            ns_steps = group["ns_steps"]
+            slow_chunk = group["slow_chunk"]
+            weight_decay = group["weight_decay"]
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if weight_decay != 0.0:
+                    grad = grad.add(p, alpha=weight_decay)
+                state = self.state[p]
+                if not state:
+                    state["step"] = 0
+                    state["m1"] = torch.zeros_like(p)
+                    state["m2"] = torch.zeros_like(p)
+                    state["v"] = torch.zeros_like(p)
+                    state["slow_buffer"] = torch.zeros_like(p)
+                    state["o2"] = torch.zeros_like(p)
+                state["step"] += 1
+                m1 = state["m1"]
+                m2 = state["m2"]
+                v = state["v"]
+                slow_buffer = state["slow_buffer"]
+
+                m1.add_(grad, alpha=beta1)
+                v.addcmul_(grad, grad, value=beta2)
+                slow_buffer.add_(grad)
+
+                o1 = _orthogonalize(m1, steps=ns_steps, eps=eps)
+                o2 = state["o2"]
+                denom = v.sqrt().add_(eps)
+                update = (o1 + alpha * o2) / denom
+                p.add_(update, alpha=-lr)
+
+                if slow_chunk > 0 and state["step"] % slow_chunk == 0:
+                    # Paper Algorithm 1 uses the updated slow momentum term in the *next* chunk.
+                    # Compute it after applying the current step update to avoid off-by-one usage.
+                    m2.add_(slow_buffer, alpha=beta3)
+                    slow_buffer.zero_()
+                    state["o2"] = _orthogonalize(m2, steps=ns_steps, eps=eps)
+        return loss

nested_learning/optim/manager.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Sequence, Tuple
+
+import torch
+from torch import nn
+
+from ..levels import LevelClock, LevelSpec
+from .factory import build_optimizer
+
+
+@dataclass
+class LevelConfig:
+    specs: Sequence[LevelSpec]
+    optimizer_configs: Dict[str, dict]
+    default_lr: float
+
+
+class LevelOptimizerManager:
+    def __init__(self, config: LevelConfig):
+        self.clock = LevelClock(config.specs)
+        self.learning_rates: Dict[str, float] = {}
+        self.optimizers = {}
+        self._last_metrics: Dict[str, Dict[str, float]] = {}
+        for spec in config.specs:
+            key = spec.optimizer_key or "default"
+            optim_cfg = config.optimizer_configs.get(key, {"type": "deep_momentum", "params": {}})
+            lr = optim_cfg.get("lr", config.default_lr)
+            params_cfg = optim_cfg.get("params", {})
+            optimizer = build_optimizer(
+                {"type": optim_cfg.get("type", "deep_momentum"), "params": params_cfg}
+            )
+            self.optimizers[spec.name] = optimizer
+            self.learning_rates[spec.name] = lr
+
+    def should_update(self, level: str) -> bool:
+        return self.clock.should_update(level)
+
+    def optimize(
+        self,
+        level: str,
+        module: nn.Module,
+        loss: torch.Tensor,
+        *,
+        context: torch.Tensor | None = None,
+        force: bool = False,
+    ) -> float:
+        if (not force) and (not self.should_update(level)):
+            return 0.0
+        named_params: Tuple[Tuple[str, torch.nn.Parameter], ...] = tuple(
+            (name, param) for name, param in module.named_parameters() if param.requires_grad
+        )
+        if not named_params:
+            return 0.0
+        params = tuple(param for _, param in named_params)
+        grads = torch.autograd.grad(loss, params, retain_graph=False, allow_unused=True)
+        grads_dict: Dict[str, torch.Tensor] = {}
+        for (name, _), grad in zip(named_params, grads, strict=True):
+            if grad is None:
+                continue
+            grads_dict[name] = grad
+        return self.apply_module_grads(
+            level,
+            module,
+            grads_dict,
+            context=context,
+            force=True,
+        )
+
+    def apply_module_grads(
+        self,
+        level: str,
+        module: nn.Module,
+        grads: Dict[str, torch.Tensor],
+        *,
+        context: torch.Tensor | None = None,
+        force: bool = False,
+    ) -> float:
+        if (not force) and (not self.should_update(level)):
+            return 0.0
+        optimizer = self.optimizers[level]
+        lr = self.learning_rates[level]
+        total_norm = 0.0
+        with torch.no_grad():
+            for name, param in module.named_parameters():
+                if not param.requires_grad:
+                    continue
+                grad = grads.get(name)
+                if grad is None:
+                    continue
+                update = optimizer(grad, context=context, param_key=name)
+                param.add_(update, alpha=-lr)
+                total_norm += grad.norm().item()
+        self.clock.record_update(level)
+        metrics = getattr(optimizer, "last_metrics", None)
+        if metrics:
+            self._last_metrics[level] = dict(metrics)
+        else:
+            self._last_metrics[level] = {}
+        return total_norm
+
+    def tick(self) -> None:
+        self.clock.tick()
+
+    def pop_last_metrics(self, level: str) -> Dict[str, float]:
+        return self._last_metrics.pop(level, {})
+
+    def apply_grads(
+        self,
+        level: str,
+        params: Dict[str, torch.Tensor],
+        grads: Dict[str, torch.Tensor],
+        *,
+        context: torch.Tensor | None = None,
+        force: bool = False,
+        differentiable: bool = False,
+    ) -> tuple[Dict[str, torch.Tensor], float]:
+        if (not force) and (not self.should_update(level)):
+            return params, 0.0
+        optimizer = self.optimizers[level]
+        lr = self.learning_rates[level]
+        updated: Dict[str, torch.Tensor] = {}
+        total_norm = 0.0
+        if differentiable:
+            for name, param in params.items():
+                grad = grads.get(name)
+                if grad is None:
+                    updated[name] = param
+                    continue
+                updated[name] = param - lr * grad
+                total_norm += float(grad.detach().norm().item())
+            self.clock.record_update(level)
+            self._last_metrics[level] = {"differentiable_updates": 1.0}
+            return updated, total_norm
+        with torch.no_grad():
+            for name, param in params.items():
+                grad = grads.get(name)
+                if grad is None:
+                    updated[name] = param
+                    continue
+                update = optimizer(grad, context=context, param_key=name)
+                updated[name] = (param - lr * update).detach()
+                total_norm += grad.norm().item()
+        self.clock.record_update(level)
+        metrics = getattr(optimizer, "last_metrics", None)
+        if metrics:
+            self._last_metrics[level] = dict(metrics)
+        else:
+            self._last_metrics[level] = {}
+        return updated, total_norm

nested_learning/titan/memory.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+import torch
+import torch.nn as nn
+
+from ..assoc_memory import AssocMemory
+
+
+@dataclass
+class TitanMemoryConfig:
+    dim: int
+    hidden_multiplier: int = 4
+    layers: int = 2
+    activation: str = "gelu"
+
+
+def _activation(name: str) -> nn.Module:
+    if name.lower() == "relu":
+        return nn.ReLU()
+    if name.lower() == "gelu":
+        return nn.GELU()
+    if name.lower() == "silu":
+        return nn.SiLU()
+    msg = f"Unsupported activation {name}"
+    raise ValueError(msg)
+
+
+class TitanMemory(AssocMemory):
+    """Simplified TITAN-style associative memory."""
+
+    def __init__(self, config: TitanMemoryConfig):
+        super().__init__()
+        self.config = config
+        hidden = config.dim * config.hidden_multiplier
+        blocks = []
+        activation = _activation(config.activation)
+        for layer_idx in range(config.layers - 1):
+            blocks.extend([nn.Linear(config.dim if layer_idx == 0 else hidden, hidden), activation])
+        blocks.append(nn.Linear(hidden if config.layers > 1 else config.dim, config.dim))
+        self.net = nn.Sequential(*blocks)
+        self.norm = nn.LayerNorm(config.dim)
+        self.grad_clip = 1.0
+
+    def forward(self, query: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        attn = self.net(query)
+        if self.training and self.grad_clip > 0:
+            with torch.no_grad():
+                norm = attn.norm(dim=-1, keepdim=True)
+                scale = torch.clamp(norm / self.grad_clip, min=1.0)
+            attn = attn / scale
+        return self.norm(attn)
+
+    def surprise(self, residual: torch.Tensor) -> torch.Tensor:
+        return residual.norm(dim=-1, keepdim=True)
+
+    @torch.no_grad()
+    def update(
+        self,
+        *,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        error_signal: torch.Tensor | None = None,
+        lr: float = 1e-3,
+    ) -> None:
+        with torch.enable_grad():
+            key_detached = key.detach().requires_grad_(True)
+            prediction = self.forward(key_detached)
+            target = value.detach()
+            if error_signal is None:
+                loss = torch.mean((prediction - target) ** 2)
+            else:
+                loss = torch.mean(error_signal * prediction)
+        grads = torch.autograd.grad(loss, list(self.net.parameters()), retain_graph=False)
+        for param, grad in zip(self.net.parameters(), grads, strict=False):
+            if grad is None:
+                continue
+            param.add_(grad, alpha=-lr)
+
+    @torch.no_grad()
+    def apply_deltas(self, deltas: Dict[str, torch.Tensor], scale: float = 1.0) -> None:
+        for name, tensor in deltas.items():
+            target = dict(self.named_parameters()).get(name)
+            if target is None:
+                continue
+            target.add_(tensor, alpha=scale)