mlx-pc

A straightforward, minimal Predictive Coding (PC) implementation for the Apple MLX framework.

mlx-pc provides stateless predictive coding layers and a network container to manage relaxation dynamics. It is designed to be completely decoupled from the optimizer, allowing you to plug in standard MLX optimizers or custom implementations.

Installation

pip install mlx-pc

API Reference

The library exposes two main components: PCNetwork and PCLayer.

1. PCNetwork (Container) The primary module that stacks PCLayers and handles the temporal state management (relaxation iterations) during the forward pass.

• Initialization:

model = PCNetwork(layer_dims: list[int], bias=True)

layer_dims: A list of integers defining the feature dimensions of each layer (e.g., [128, 256, 512]).

• Forward Pass (call):

predictions, layer_errors = model(sensory_x, max_iters=10, base_eta=0.1, alpha=1.0, beta=0.5)

x: Input tensor of shape (batch_size, seq_len, layer_dims[0]). max_iters: Number of internal relaxation steps to perform for the fast variables. Returns: A tuple containing the final predictions and a list of layer_errors tensors.

2. PCLayer (Stateless Block) A single predictive coding layer. It acts as a pure function, computing local surprisal and the updated state based on bottom-up inputs and top-down priors.

• Initialization:

layer = PCLayer(in_dim: int, out_dim: int, bias=True)

• Forward Pass (call):

current_state, error = layer(x, higher_state, prev_state, base_eta=0.1, alpha=1.0, beta=0.5)

x: The bottom-up input from the lower layer. higher_state: The expected state from the higher layer. prev_state: The state of this layer at iteration t-1. Returns: The updated current_state and the calculated error (residual).

Quick Integration Example

mlx-pc relies on standard MLX auto-grad to update weights. You extract the total squared error (Free Energy) and pass it to any optimizer.

import mlx.core as mx
import mlx.optimizers as optim
from mlx_pc import PCNetwork

# 1. Init
model = PCNetwork(layer_dims=[128, 256, 512])
optimizer = optim.AdamW(learning_rate=1e-3)

# 2. Define Loss
def loss_fn(model, x):
    predictions, layer_errors = model(x)
    return mx.sum(mx.array([mx.sum(mx.square(e)) for e in layer_errors]))

loss_and_grad_fn = mx.value_and_grad(model, loss_fn)

# 3. Step
x = mx.random.normal((4, 32, 128))
loss, grads = loss_and_grad_fn(model, x)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state)