manutej Polynomial Functors for Learner Composition
Implementation of Spivak-Niu polynomial functor framework for compositional learning systems.
Core Concepts
Polynomial functors model systems with:
- Position: Internal state space
- Direction: Interface/interaction type
- Lens: Bidirectional data flow (get/put)
- Chart: Dynamical update rules
Basic Polynomial Structure
from dataclasses import dataclass
from typing import Generic, TypeVar, Callable, Tuple
S = TypeVar('S') # State
A = TypeVar('A') # Input
B = TypeVar('B') # Output
@dataclass
class Polynomial(Generic[S, A, B]):
"""
Polynomial functor p(y) = Σ_{s∈S} y^{A(s)}
Represents a system with:
- Positions S (internal states)
- Directions A(s) → B (interface at each state)
"""
positions: set # S
directions: Callable[[S], Tuple[type, type]] # s ↦ (A(s), B(s))
Lenses as Morphisms
@dataclass
class Lens(Generic[S, A, B]):
"""
Lens: bidirectional transformation.
get: S → A (forward/observation)
put: S × B → S (backward/update)
"""
get: Callable[[S], A]
put: Callable[[S, B], S]
def compose(self, other: 'Lens[A, C, D]') -> 'Lens[S, C, D]':
"""Lens composition (Kleisli-like)."""
return Lens(
get=lambda s: other.get(self.get(s)),
put=lambda s, d: self.put(s, other.put(self.get(s), d))
)
# Example: Neural network layer as lens
@dataclass
class LayerLens:
weights: 'np.ndarray'
def forward(self, x: 'np.ndarray') -> 'np.ndarray':
return self.weights @ x
def backward(self, x: 'np.ndarray', grad: 'np.ndarray') -> Tuple['np.ndarray', 'np.ndarray']:
weight_grad = np.outer(grad, x)
input_grad = self.weights.T @ grad
return weight_grad, input_grad
Learners as Polynomial Coalgebras
@dataclass
class Learner(Generic[S, A, B]):
"""
Learner in the sense of Spivak-Niu.
A learner is a lens equipped with:
- Parameter space P
- Implementation: P × A → B
- Update: P × A × B → P
"""
param_space: type
implement: Callable # (P, A) → B
update: Callable # (P, A, B) → P
request: Callable # (P, A) → A (optional, for backprop)
def compose_learners(l1: Learner, l2: Learner) -> Learner:
"""
Compose learners sequentially.
l1: A → B with params P1
l2: B → C with params P2
Result: A → C with params (P1, P2)
"""
return Learner(
param_space=(l1.param_space, l2.param_space),
implement=lambda p, a: l2.implement(p[1], l1.implement(p[0], a)),
update=lambda p, a, c: (
l1.update(p[0], a, l2.request(p[1], l1.implement(p[0], a))),
l2.update(p[1], l1.implement(p[0], a), c)
),
request=lambda p, a: l1.request(p[0], a)
)
Charts for Dynamical Systems
@dataclass
class Chart(Generic[S, A, B]):
"""
Chart: dynamical system on polynomial.
A chart assigns to each position a way to:
- Observe output from state
- Update state given input
"""
readout: Callable[[S], B] # Observe
update: Callable[[S, A], S] # Transition
def run(self, initial: S, inputs: list) -> list:
"""Execute dynamical system."""
state = initial
outputs = []
for inp in inputs:
outputs.append(self.readout(state))
state = self.update(state, inp)
return outputs
Wiring Diagrams
@dataclass
class WiringDiagram:
"""
Wiring diagram: composition pattern for polynomials.
Specifies how outputs connect to inputs across components.
"""
boxes: list # Component polynomials
wires: list # Connections (box_idx, port) → (box_idx, port)
def compose(self) -> 'Polynomial':
"""Evaluate wiring diagram to single polynomial."""
# Trace through connections
pass
# Example: Feedback loop
def feedback(p: Polynomial) -> Polynomial:
"""Create feedback loop: connect output to input."""
return WiringDiagram(
boxes=[p],
wires=[(0, 'out', 0, 'in')] # Connect output back to input
).compose()
Neural Network as Polynomial
import numpy as np
class NeuralLearner:
"""Neural network layer as polynomial functor learner."""
def __init__(self, in_dim: int, out_dim: int):
self.weights = np.random.randn(out_dim, in_dim) * 0.01
self.bias = np.zeros(out_dim)
def implement(self, x: np.ndarray) -> np.ndarray:
"""Forward pass."""
return np.tanh(self.weights @ x + self.bias)
def request(self, x: np.ndarray, grad_out: np.ndarray) -> np.ndarray:
"""Backward pass - compute gradient request."""
output = self.implement(x)
dtanh = 1 - output ** 2
return self.weights.T @ (grad_out * dtanh)
def update(self, x: np.ndarray, grad_out: np.ndarray, lr: float = 0.01):
"""Parameter update."""
output = self.implement(x)
dtanh = 1 - output ** 2
delta = grad_out * dtanh
self.weights -= lr * np.outer(delta, x)
self.bias -= lr * delta
def compose_layers(layers: list) -> 'ComposedNetwork':
"""Compose neural layers as polynomial functor composition."""
class ComposedNetwork:
def __init__(self, layers):
self.layers = layers
def forward(self, x):
for layer in self.layers:
x = layer.implement(x)
return x
def backward(self, x, grad):
# Store activations
activations = [x]
for layer in self.layers:
x = layer.implement(x)
activations.append(x)
# Backpropagate
for i, layer in reversed(list(enumerate(self.layers))):
grad = layer.request(activations[i], grad)
layer.update(activations[i], grad)
return ComposedNetwork(layers)
Categorical Properties
def verify_lens_laws(lens: Lens, s, b):
"""Verify lens satisfies get-put and put-get laws."""
# Get-Put: put(s, get(s)) = s
assert lens.put(s, lens.get(s)) == s, "Get-Put law violated"
# Put-Get: get(put(s, b)) = b
assert lens.get(lens.put(s, b)) == b, "Put-Get law violated"
def verify_composition_associativity(l1, l2, l3, s, d):
"""Verify lens composition is associative."""
left = l1.compose(l2).compose(l3)
right = l1.compose(l2.compose(l3))
assert left.get(s) == right.get(s), "Associativity violated (get)"
assert left.put(s, d) == right.put(s, d), "Associativity violated (put)"
Categorical Guarantees
Polynomial functors provide:
- Compositionality: Learners compose via lens composition
- Bidirectionality: Forward/backward passes as lens get/put
- Modularity: Wiring diagrams specify composition patterns
- Dynamical Semantics: Charts model system evolution
- Type Safety: Polynomial types ensure interface compatibility