docs/MMEP_THEORY.md

Matrix-Memory Equilibrium Propagation: Mathematical Foundation

Project: BURT-IMMA
Contact: jessica@collectivekitty.com
License: BSL-1.1

---

1. Overview

Matrix-Memory Equilibrium Propagation (MMEP) is a biologically plausible training algorithm that combines equilibrium propagation with structured matrix memory. Unlike backpropagation, MMEP relies on local Hebbian-style learning rules derived from the difference between free and nudged equilibrium states, while maintaining hard constraints on spectral norm, entropy, and sparsity.

MMEP replaces the implicit differentiation of standard EP with explicit matrix-memory dynamics governed by CIFG (Coupled Input-Forget Gate) cells, enabling the network to store and retrieve knowledge through invertible matrix operations.

---

2. Energy Function

The total energy of the system is defined as:

E(C, H, alpha; theta) = E_pred + E_constraint + E_entropy + E_sparsity

where:

• C is the set of memory matrices (global + per-expert)
• H is the set of hidden state activations
• alpha is the expert gating distribution
• theta is the set of all learnable parameters

2.1 Prediction Energy

E_pred = ||y - f(x; theta)||^2

Standard squared prediction error. The function f(x; theta) represents the full forward pass through the BURT architecture, including Boolean constraint enforcement and memory retrieval.

2.2 Spectral Constraint Energy

E_constraint = lambda * max(0, sigma_max(W) - lambda_max)

where:

• sigma_max(W) is the largest singular value of weight matrix W
• lambda_max is the spectral radius bound (default: 0.95)
• lambda is the constraint penalty strength

This term ensures all weight matrices remain within a bounded spectral radius, preventing gradient explosion and guaranteeing convergence of the equilibrium dynamics.

2.3 Entropy Energy

E_entropy = -sum(alpha_k * log(alpha_k))    for k = 1, ..., K

Bounded by 0.20: The entropy of the expert gating distribution is constrained to not exceed 0.20 nats. This prevents the router from distributing load uniformly across all experts (which would defeat the purpose of sparse mixture-of-experts) while ensuring at least some diversity in expert selection.

The bound is enforced via projection after each update:

if H(alpha) > 0.20:
    alpha <- project_entropy(alpha, bound=0.20)

2.4 Sparsity Energy

E_sparsity = ||alpha||_1

The L1 norm of the gating vector encourages sparse expert activation. Combined with Gates normalization, this ensures that only a small subset of experts are active for any given input, maintaining computational efficiency and interpretability.

---

3. Two-Phase Dynamics

MMEP operates in two distinct phases that together define the learning signal.

3.1 Free Phase

In the free phase, the network relaxes to its natural equilibrium without external forcing:

dh/dt = -dE/dh

The hidden states h evolve according to the negative gradient of the energy function with respect to themselves. This is iterated for T_free steps (default: 20) until approximate convergence:

for t in range(T_free):
    h <- h - eta_h * dE/dh

At convergence, the network reaches a fixed point h_free that represents its best prediction given the current parameters.

Convergence criterion: ||dh/dt|| < epsilon where epsilon = 1e-5.

3.2 Nudged Phase

In the nudged phase, the output layer is clamped toward the target with strength beta:

dh/dt = -dE/dh - beta * (h_output - y_target)

Only the output layer receives the nudging signal. The rest of the network adjusts to accommodate this external forcing, reaching a new equilibrium h_nudged after T_nudge steps (default: 4):

for t in range(T_nudge):
    h <- h - eta_h * (dE/dh + beta * clamp_signal)

The nudging strength beta controls the trade-off between biological plausibility (small beta) and learning speed (large beta). In the limit beta -> 0, the learning rule becomes exactly the gradient of the energy function.

---

4. Local Learning Rule

The weight update is computed from the difference between nudged and free equilibrium states:

dW = (1/beta) * (h_nudged * h_nudged^T - h_free * h_free^T)

This is a purely local, Hebbian-style rule: each synapse only needs access to the pre- and post-synaptic activities in both phases. No backpropagation of error signals through the network is required.

Properties:

• In the limit beta -> 0: dW converges to the true gradient dE/dtheta
• For finite beta: dW is a biased but consistent estimator of the gradient
• The rule is symmetric: dW = dW^T for symmetric connections
• Memory matrices are updated by the same rule applied to their input/output pairs

4.1 Bias Updates

db = (1/beta) * (h_nudged - h_free)

4.2 Memory Matrix Updates

For each memory matrix C:

dC = (1/beta) * (h_nudged * h_nudged^T - h_free * h_free^T) projected onto rank-1 update

---

5. CIFG Freeze Dynamics

The Coupled Input-Forget Gate (CIFG) controls how memory matrices are updated, preventing catastrophic forgetting:

C_new = f * C_old + (1 - f) * candidate

where:

f = sigma(W_f * [h, x])

• f is the forget gate (values in [0, 1])
• sigma is the sigmoid function
• [h, x] is the concatenation of hidden state and input
• candidate is the proposed new memory content

Key insight: When f -> 1, the memory is frozen (old content preserved). When f -> 0, the memory is fully overwritten. The CIFG couples the input and forget gates (input gate = 1 - f), halving the parameter count compared to standard LSTM gates while maintaining expressivity.

5.1 Freeze Schedule

During training, the forget gate bias is gradually increased:

bias_f(epoch) = bias_f_init + freeze_rate * epoch

This encourages early plasticity followed by later consolidation, mimicking biological memory formation.

---

6. Memory Matrices

6.1 Global Memory: C_global

A single shared memory matrix accessible by all experts and all layers:

• Shape: [d_model, d_model]
• Stores corpus-level knowledge (facts, patterns, invariants)
• Updated during Phase 1 (memorization) and Phase 3 (convergence)
• Protected by high forget gate bias after initial training

6.2 Expert Memory: C_expert_k

Per-expert memory matrices (one per expert in the MoE layer):

• Shape: [d_expert, d_expert] for each expert k = 1, ..., K
• Stores expert-specific knowledge
• Updated only when the expert is selected by the router
• Lower forget gate bias allows more rapid adaptation

6.3 Memory Retrieval

Given a query vector q, retrieval from memory C is:

v = C * q / ||C * q||

Normalized retrieval ensures bounded activations regardless of memory content magnitude.

---

7. Convergence Guarantee

Theorem (Lyapunov Convergence): Under the MMEP dynamics with spectral constraint sigma_max(W) <= lambda_max < 1, the energy function E(C, H, alpha; theta) is a Lyapunov function for the free-phase dynamics.

Proof sketch:

Define V(h) = E(C, H, alpha; theta) evaluated at the current state.

1. Bounded below: V(h) >= 0 since all terms are non-negative (squared error, max with 0, bounded entropy, L1 norm).

2. Monotonically decreasing: Along trajectories of dh/dt = -dE/dh:

   dV/dt = (dE/dh)^T * (dh/dt) = -||dE/dh||^2 <= 0
   

3. Strict decrease: dV/dt = 0 only at fixed points where dE/dh = 0.

4. Bounded trajectories: The spectral constraint ensures ||h(t)|| <= M for some finite M, since:

   ||Wh|| <= sigma_max(W) * ||h|| <= lambda_max * ||h|| < ||h||
   

   for lambda_max < 1, guaranteeing contraction.

By LaSalle's invariance principle, all trajectories converge to the set of equilibria. The spectral constraint guarantees uniqueness of the equilibrium for each input, completing the proof.

Convergence rate: Exponential with rate 1 - lambda_max. For the default lambda_max = 0.95, convergence to machine precision requires approximately T_free = ceil(log(epsilon) / log(0.95)) = 98 steps, though in practice T_free = 20 suffices for the learning rule.

---

8. Constraint Projection

After each parameter update, constraints are enforced by projection:

8.1 Spectral Projection

def project_spectral(W, lambda_max=0.95):
    U, S, V = svd(W)
    S_clipped = minimum(S, lambda_max)
    return U @ diag(S_clipped) @ V^T

8.2 Entropy Projection

def project_entropy(alpha, bound=0.20):
    if entropy(alpha) <= bound:
        return alpha
    # Binary search for temperature tau such that
    # H(softmax(log(alpha) / tau)) = bound
    tau = binary_search_temperature(alpha, bound)
    return softmax(log(alpha) / tau)

8.3 Combined Projection

def project_all(params):
    for W in params.weight_matrices:
        W <- project_spectral(W)
    for alpha in params.gating_vectors:
        alpha <- project_entropy(alpha)
    # Sparsity is enforced via L1 penalty, not projection

---

9. Relationship to Standard EP

Property	Standard EP	MMEP
Energy function	`E =
Memory	Implicit in weights	Explicit matrix memory (C_global, C_expert)
Constraints	None	Spectral, entropy, sparsity
Convergence	Assumed	Guaranteed via Lyapunov + spectral bound
Gating	None	CIFG freeze dynamics
Expert routing	None	Sparse MoE with Gates normalization
Learning rule	Same	Same form: (1/beta)(nudged - free)

---

10. Notation Summary

Symbol	Meaning
E	Total energy
C	Memory matrix (global or expert)
H	Hidden states
alpha	Expert gating distribution
theta	All learnable parameters
sigma_max(W)	Largest singular value of W
lambda_max	Spectral radius bound (default 0.95)
beta	Nudging strength
f	Forget gate value
T_free	Free phase steps (default 20)
T_nudge	Nudged phase steps (default 4)
K	Number of experts
d_model	Model dimension
d_expert	Expert dimension

---

References

1. Scellier, B. & Bengio, Y. (2017). "Equilibrium Propagation: Bridging the Gap between Energy-Based Models and Backpropagation." Frontiers in Computational Neuroscience.
2. Greff, K. et al. (2017). "LSTM: A Search Space Odyssey." IEEE TNNLS.
3. Fedus, W. et al. (2022). "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity."
4. Miyato, T. et al. (2018). "Spectral Normalization for Generative Adversarial Networks." ICLR.