monomind 1.10.0 → 1.10.2

package/packages/@monomind/cli/dist/src/ruvector/moe-router.js

@@ -0,0 +1,649 @@
+/**
+ * Mixture of Experts (MoE) Router for Dynamic Agent Routing
+ *
+ * Features:
+ * - 8 expert slots for specialized agent types
+ * - Gating network for soft expert selection (top-k)
+ * - Online weight updates via reward signals
+ * - Load balancing with auxiliary loss
+ * - Weight persistence to .swarm/moe-weights.json
+ *
+ * Architecture:
+ * - Input: 384-dim task embedding (from ONNX)
+ * - Hidden: 128-dim layer with ReLU
+ * - Output: 8-dim softmax weights
+ *
+ * @module moe-router
+ */
+import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync } from 'fs';
+import { dirname, join } from 'path';
+/**
+ * Expert names in order (index corresponds to expert slot)
+ */
+export const EXPERT_NAMES = [
+    'coder',
+    'tester',
+    'reviewer',
+    'architect',
+    'security',
+    'performance',
+    'researcher',
+    'coordinator',
+];
+/**
+ * Number of experts (fixed at 8)
+ */
+export const NUM_EXPERTS = 8;
+/**
+ * Input dimension (384 from ONNX MiniLM-L6-v2)
+ */
+export const INPUT_DIM = 384;
+/**
+ * Hidden layer dimension
+ */
+export const HIDDEN_DIM = 128;
+/**
+ * Default configuration
+ */
+const DEFAULT_CONFIG = {
+    topK: 2,
+    learningRate: 0.01,
+    temperature: 1.0,
+    loadBalanceCoef: 0.01,
+    weightsPath: '.swarm/moe-weights.json',
+    autoSaveInterval: 50,
+    enableNoise: true,
+    noiseStd: 0.1,
+};
+// ============================================================================
+// Matrix Operations
+// ============================================================================
+/**
+ * Initialize weights using Xavier/Glorot initialization
+ */
+function xavierInit(fanIn, fanOut) {
+    const std = Math.sqrt(2.0 / (fanIn + fanOut));
+    const weights = new Float32Array(fanIn * fanOut);
+    for (let i = 0; i < weights.length; i++) {
+        // Box-Muller transform for normal distribution
+        const u1 = Math.random();
+        const u2 = Math.random();
+        const z = Math.sqrt(-2 * Math.log(u1 + 1e-8)) * Math.cos(2 * Math.PI * u2);
+        weights[i] = z * std;
+    }
+    return weights;
+}
+/**
+ * Matrix-vector multiplication: y = Wx
+ * W is stored row-major: [rows * cols]
+ */
+function matmul(W, x, rows, cols, out) {
+    for (let i = 0; i < rows; i++) {
+        let sum = 0;
+        const rowOffset = i * cols;
+        // 4x loop unrolling for SIMD-friendly access
+        let j = 0;
+        for (; j + 3 < cols; j += 4) {
+            sum +=
+                W[rowOffset + j] * x[j] +
+                    W[rowOffset + j + 1] * x[j + 1] +
+                    W[rowOffset + j + 2] * x[j + 2] +
+                    W[rowOffset + j + 3] * x[j + 3];
+        }
+        // Handle remainder
+        for (; j < cols; j++) {
+            sum += W[rowOffset + j] * x[j];
+        }
+        out[i] = sum;
+    }
+}
+/**
+ * Vector addition: y = x + b
+ */
+function addBias(x, b, out) {
+    for (let i = 0; i < x.length; i++) {
+        out[i] = x[i] + b[i];
+    }
+}
+/**
+ * ReLU activation: y = max(0, x)
+ */
+function relu(x, out) {
+    for (let i = 0; i < x.length; i++) {
+        out[i] = x[i] > 0 ? x[i] : 0;
+    }
+}
+/**
+ * Softmax with temperature: y_i = exp(x_i/T) / sum(exp(x_j/T))
+ */
+function softmax(x, temperature, out) {
+    // Find max for numerical stability
+    let maxVal = x[0];
+    for (let i = 1; i < x.length; i++) {
+        if (x[i] > maxVal)
+            maxVal = x[i];
+    }
+    // Compute exp and sum
+    let sum = 0;
+    for (let i = 0; i < x.length; i++) {
+        out[i] = Math.exp((x[i] - maxVal) / temperature);
+        sum += out[i];
+    }
+    // Normalize
+    const invSum = 1.0 / (sum + 1e-8);
+    for (let i = 0; i < x.length; i++) {
+        out[i] *= invSum;
+    }
+}
+/**
+ * Compute entropy of distribution: H = -sum(p * log(p))
+ */
+function entropy(p) {
+    let h = 0;
+    for (let i = 0; i < p.length; i++) {
+        if (p[i] > 1e-8) {
+            h -= p[i] * Math.log(p[i]);
+        }
+    }
+    return h;
+}
+/**
+ * Add Gaussian noise for exploration
+ */
+function addNoise(x, std, out) {
+    for (let i = 0; i < x.length; i++) {
+        const u1 = Math.random();
+        const u2 = Math.random();
+        const z = Math.sqrt(-2 * Math.log(u1 + 1e-8)) * Math.cos(2 * Math.PI * u2);
+        out[i] = x[i] + z * std;
+    }
+}
+// ============================================================================
+// MoE Router Implementation
+// ============================================================================
+/**
+ * Mixture of Experts Router
+ *
+ * Implements a two-layer gating network:
+ * - Layer 1: Linear(384, 128) + ReLU
+ * - Layer 2: Linear(128, 8) + Softmax
+ *
+ * Uses top-k expert selection with load balancing.
+ */
+export class MoERouter {
+    config;
+    // Network weights (pre-allocated Float32Arrays)
+    W1; // [HIDDEN_DIM x INPUT_DIM]
+    b1; // [HIDDEN_DIM]
+    W2; // [NUM_EXPERTS x HIDDEN_DIM]
+    b2; // [NUM_EXPERTS]
+    // Intermediate buffers (pre-allocated, no GC pressure)
+    hidden; // [HIDDEN_DIM]
+    hiddenWithBias; // [HIDDEN_DIM]
+    hiddenActivated; // [HIDDEN_DIM]
+    logits; // [NUM_EXPERTS]
+    logitsWithBias; // [NUM_EXPERTS]
+    noisyLogits; // [NUM_EXPERTS]
+    probs; // [NUM_EXPERTS]
+    // Gradient buffers for online learning
+    gradW2; // [NUM_EXPERTS x HIDDEN_DIM]
+    gradb2; // [NUM_EXPERTS]
+    gradW1; // [HIDDEN_DIM x INPUT_DIM]
+    gradb1; // [HIDDEN_DIM]
+    gradHidden; // [HIDDEN_DIM]
+    // Statistics
+    routingCounts; // [NUM_EXPERTS]
+    totalRoutings = 0;
+    updateCount = 0;
+    avgReward = 0;
+    // Cache for last input (for gradient computation)
+    lastInput = null;
+    lastHiddenActivated = null;
+    lastProbs = null;
+    lastSelectedExperts = [];
+    constructor(config = {}) {
+        this.config = { ...DEFAULT_CONFIG, ...config };
+        // Initialize weights
+        this.W1 = xavierInit(INPUT_DIM, HIDDEN_DIM);
+        this.b1 = new Float32Array(HIDDEN_DIM);
+        this.W2 = xavierInit(HIDDEN_DIM, NUM_EXPERTS);
+        this.b2 = new Float32Array(NUM_EXPERTS);
+        // Pre-allocate intermediate buffers
+        this.hidden = new Float32Array(HIDDEN_DIM);
+        this.hiddenWithBias = new Float32Array(HIDDEN_DIM);
+        this.hiddenActivated = new Float32Array(HIDDEN_DIM);
+        this.logits = new Float32Array(NUM_EXPERTS);
+        this.logitsWithBias = new Float32Array(NUM_EXPERTS);
+        this.noisyLogits = new Float32Array(NUM_EXPERTS);
+        this.probs = new Float32Array(NUM_EXPERTS);
+        // Pre-allocate gradient buffers
+        this.gradW2 = new Float32Array(NUM_EXPERTS * HIDDEN_DIM);
+        this.gradb2 = new Float32Array(NUM_EXPERTS);
+        this.gradW1 = new Float32Array(HIDDEN_DIM * INPUT_DIM);
+        this.gradb1 = new Float32Array(HIDDEN_DIM);
+        this.gradHidden = new Float32Array(HIDDEN_DIM);
+        // Statistics
+        this.routingCounts = new Float32Array(NUM_EXPERTS);
+    }
+    /**
+     * Initialize router, loading persisted weights if available
+     */
+    async initialize() {
+        await this.loadWeights();
+    }
+    /**
+     * Route task to top-k experts based on embedding
+     *
+     * @param taskEmbedding - 384-dim task embedding from ONNX
+     * @returns Routing result with selected experts and weights
+     */
+    route(taskEmbedding) {
+        // Convert to Float32Array if needed
+        const input = taskEmbedding instanceof Float32Array
+            ? taskEmbedding
+            : new Float32Array(taskEmbedding);
+        // Validate input dimension
+        if (input.length !== INPUT_DIM) {
+            throw new Error(`Expected embedding dimension ${INPUT_DIM}, got ${input.length}`);
+        }
+        // Forward pass through gating network
+        // Layer 1: Linear + ReLU
+        matmul(this.W1, input, HIDDEN_DIM, INPUT_DIM, this.hidden);
+        addBias(this.hidden, this.b1, this.hiddenWithBias);
+        relu(this.hiddenWithBias, this.hiddenActivated);
+        // Layer 2: Linear
+        matmul(this.W2, this.hiddenActivated, NUM_EXPERTS, HIDDEN_DIM, this.logits);
+        addBias(this.logits, this.b2, this.logitsWithBias);
+        // Add noise for exploration if enabled
+        if (this.config.enableNoise) {
+            addNoise(this.logitsWithBias, this.config.noiseStd, this.noisyLogits);
+        }
+        else {
+            this.noisyLogits.set(this.logitsWithBias);
+        }
+        // Softmax to get probabilities
+        softmax(this.noisyLogits, this.config.temperature, this.probs);
+        // Select top-k experts
+        const expertIndices = this.selectTopK(this.probs, this.config.topK);
+        // Compute load balance loss
+        const loadBalanceLoss = this.computeLoadBalanceLoss();
+        // Compute entropy
+        const routingEntropy = entropy(this.probs);
+        // Update statistics
+        for (const idx of expertIndices) {
+            this.routingCounts[idx]++;
+        }
+        this.totalRoutings++;
+        // Cache for gradient computation
+        this.lastInput = new Float32Array(input);
+        this.lastHiddenActivated = new Float32Array(this.hiddenActivated);
+        this.lastProbs = new Float32Array(this.probs);
+        this.lastSelectedExperts = expertIndices;
+        // Build result
+        const totalWeight = expertIndices.reduce((sum, idx) => sum + this.probs[idx], 0);
+        const experts = expertIndices.map((idx) => ({
+            name: EXPERT_NAMES[idx],
+            index: idx,
+            weight: this.probs[idx] / (totalWeight + 1e-8), // Normalize weights
+            score: this.probs[idx],
+        }));
+        return {
+            experts,
+            allScores: Array.from(this.probs),
+            loadBalanceLoss,
+            entropy: routingEntropy,
+        };
+    }
+    /**
+     * Update expert weights based on reward signal
+     *
+     * Uses REINFORCE-style gradient update:
+     * gradient = reward * d_log_prob / d_weights
+     *
+     * @param expert - Expert that received the reward
+     * @param reward - Reward signal (-1 to 1, positive = good)
+     */
+    updateExpertWeights(expert, reward) {
+        const expertIdx = typeof expert === 'number' ? expert : EXPERT_NAMES.indexOf(expert);
+        if (expertIdx < 0 || expertIdx >= NUM_EXPERTS) {
+            console.warn(`[MoE] Invalid expert: ${expert}`);
+            return;
+        }
+        if (!this.lastInput || !this.lastHiddenActivated || !this.lastProbs) {
+            console.warn('[MoE] No cached forward pass for gradient computation');
+            return;
+        }
+        // Clamp reward to [-1, 1]
+        const clampedReward = Math.max(-1, Math.min(1, reward));
+        // Compute gradients using REINFORCE
+        // For softmax: d_log_p_i / d_logit_j = delta_ij - p_j
+        // gradient = reward * (1 - p_expert) for selected expert
+        // gradient = reward * (-p_j) for other experts
+        // Clear gradient buffers
+        this.gradW2.fill(0);
+        this.gradb2.fill(0);
+        this.gradW1.fill(0);
+        this.gradb1.fill(0);
+        this.gradHidden.fill(0);
+        // Gradient w.r.t. logits (before softmax)
+        for (let i = 0; i < NUM_EXPERTS; i++) {
+            if (i === expertIdx) {
+                this.gradb2[i] = clampedReward * (1 - this.lastProbs[i]);
+            }
+            else {
+                this.gradb2[i] = clampedReward * (-this.lastProbs[i]);
+            }
+        }
+        // Gradient w.r.t. W2: outer product of gradb2 and hiddenActivated
+        for (let i = 0; i < NUM_EXPERTS; i++) {
+            const rowOffset = i * HIDDEN_DIM;
+            for (let j = 0; j < HIDDEN_DIM; j++) {
+                this.gradW2[rowOffset + j] = this.gradb2[i] * this.lastHiddenActivated[j];
+            }
+        }
+        // Backprop through W2 to get gradient w.r.t. hidden
+        for (let j = 0; j < HIDDEN_DIM; j++) {
+            let sum = 0;
+            for (let i = 0; i < NUM_EXPERTS; i++) {
+                sum += this.gradb2[i] * this.W2[i * HIDDEN_DIM + j];
+            }
+            this.gradHidden[j] = sum;
+        }
+        // Backprop through ReLU: gradient is 0 where activation was 0
+        for (let j = 0; j < HIDDEN_DIM; j++) {
+            if (this.lastHiddenActivated[j] <= 0) {
+                this.gradHidden[j] = 0;
+            }
+        }
+        // Gradient w.r.t. b1
+        this.gradb1.set(this.gradHidden);
+        // Gradient w.r.t. W1: outer product of gradHidden and input
+        for (let i = 0; i < HIDDEN_DIM; i++) {
+            const rowOffset = i * INPUT_DIM;
+            for (let j = 0; j < INPUT_DIM; j++) {
+                this.gradW1[rowOffset + j] = this.gradHidden[i] * this.lastInput[j];
+            }
+        }
+        // Apply gradients with learning rate
+        const lr = this.config.learningRate;
+        for (let i = 0; i < this.W2.length; i++) {
+            this.W2[i] += lr * this.gradW2[i];
+        }
+        for (let i = 0; i < this.b2.length; i++) {
+            this.b2[i] += lr * this.gradb2[i];
+        }
+        for (let i = 0; i < this.W1.length; i++) {
+            this.W1[i] += lr * this.gradW1[i];
+        }
+        for (let i = 0; i < this.b1.length; i++) {
+            this.b1[i] += lr * this.gradb1[i];
+        }
+        // Update statistics
+        this.updateCount++;
+        this.avgReward = (this.avgReward * (this.updateCount - 1) + clampedReward) / this.updateCount;
+        // Auto-save
+        if (this.config.autoSaveInterval > 0 &&
+            this.updateCount % this.config.autoSaveInterval === 0) {
+            this.saveWeights().catch(() => { }); // Fire and forget
+        }
+    }
+    /**
+     * Get load balance statistics across all experts
+     */
+    getLoadBalance() {
+        const counts = {};
+        const utilization = {};
+        const total = this.totalRoutings || 1;
+        const idealUtilization = 1 / NUM_EXPERTS;
+        for (let i = 0; i < NUM_EXPERTS; i++) {
+            const name = EXPERT_NAMES[i];
+            counts[name] = this.routingCounts[i];
+            utilization[name] = this.routingCounts[i] / total;
+        }
+        // Compute Gini coefficient
+        const gini = this.computeGiniCoefficient();
+        // Compute coefficient of variation
+        const mean = total / NUM_EXPERTS;
+        let variance = 0;
+        for (let i = 0; i < NUM_EXPERTS; i++) {
+            variance += Math.pow(this.routingCounts[i] - mean, 2);
+        }
+        variance /= NUM_EXPERTS;
+        const cv = Math.sqrt(variance) / (mean + 1e-8);
+        return {
+            utilization,
+            totalRoutings: this.totalRoutings,
+            routingCounts: counts,
+            giniCoefficient: gini,
+            coefficientOfVariation: cv,
+        };
+    }
+    /**
+     * Get router statistics
+     */
+    getStats() {
+        return {
+            totalRoutings: this.totalRoutings,
+            updateCount: this.updateCount,
+            avgReward: this.avgReward,
+            topK: this.config.topK,
+            temperature: this.config.temperature,
+            learningRate: this.config.learningRate,
+            giniCoefficient: this.computeGiniCoefficient(),
+        };
+    }
+    /**
+     * Reset all statistics and routing counts
+     */
+    resetStats() {
+        this.routingCounts.fill(0);
+        this.totalRoutings = 0;
+        this.updateCount = 0;
+        this.avgReward = 0;
+    }
+    /**
+     * Load weights from persistence file
+     */
+    async loadWeights(path) {
+        const weightsPath = path ? path : join(process.cwd(), this.config.weightsPath);
+        try {
+            if (!existsSync(weightsPath)) {
+                return false;
+            }
+            const data = readFileSync(weightsPath, 'utf-8');
+            const model = JSON.parse(data);
+            // Validate version
+            if (!model.version || !model.version.startsWith('1.')) {
+                console.warn(`[MoE] Incompatible model version: ${model.version}`);
+                return false;
+            }
+            // Load weights with dimension validation
+            const w1 = new Float32Array(model.weights.W1.flat());
+            const b1 = new Float32Array(model.weights.b1);
+            const w2 = new Float32Array(model.weights.W2.flat());
+            const b2 = new Float32Array(model.weights.b2);
+            if (w1.length !== INPUT_DIM * HIDDEN_DIM || b1.length !== HIDDEN_DIM ||
+                w2.length !== HIDDEN_DIM * NUM_EXPERTS || b2.length !== NUM_EXPERTS) {
+                console.warn('[MoE] Weight dimensions mismatch — reinitializing');
+                return false;
+            }
+            this.W1 = w1;
+            this.b1 = b1;
+            this.W2 = w2;
+            this.b2 = b2;
+            // Load stats
+            this.updateCount = model.stats.updateCount || 0;
+            this.avgReward = model.stats.avgReward || 0;
+            this.routingCounts = new Float32Array(model.stats.routingCounts || new Array(NUM_EXPERTS).fill(0));
+            this.totalRoutings = this.routingCounts.reduce((a, b) => a + b, 0);
+            return true;
+        }
+        catch (err) {
+            console.warn(`[MoE] Failed to load weights: ${err}`);
+            return false;
+        }
+    }
+    /**
+     * Save weights to persistence file
+     */
+    async saveWeights(path) {
+        const weightsPath = path ? path : join(process.cwd(), this.config.weightsPath);
+        try {
+            // Ensure directory exists
+            const dir = dirname(weightsPath);
+            if (!existsSync(dir)) {
+                mkdirSync(dir, { recursive: true });
+            }
+            // Convert Float32Arrays to nested arrays for JSON
+            const W1_2d = [];
+            for (let i = 0; i < HIDDEN_DIM; i++) {
+                W1_2d.push(Array.from(this.W1.slice(i * INPUT_DIM, (i + 1) * INPUT_DIM)));
+            }
+            const W2_2d = [];
+            for (let i = 0; i < NUM_EXPERTS; i++) {
+                W2_2d.push(Array.from(this.W2.slice(i * HIDDEN_DIM, (i + 1) * HIDDEN_DIM)));
+            }
+            const model = {
+                version: '1.0.0',
+                config: {
+                    topK: this.config.topK,
+                    temperature: this.config.temperature,
+                    learningRate: this.config.learningRate,
+                    loadBalanceCoef: this.config.loadBalanceCoef,
+                },
+                weights: {
+                    W1: W1_2d,
+                    b1: Array.from(this.b1),
+                    W2: W2_2d,
+                    b2: Array.from(this.b2),
+                },
+                stats: {
+                    updateCount: this.updateCount,
+                    routingCounts: Array.from(this.routingCounts),
+                    avgReward: this.avgReward,
+                },
+                metadata: {
+                    savedAt: new Date().toISOString(),
+                    expertNames: [...EXPERT_NAMES],
+                },
+            };
+            const tmp = weightsPath + '.tmp';
+            writeFileSync(tmp, JSON.stringify(model, null, 2));
+            renameSync(tmp, weightsPath);
+            return true;
+        }
+        catch (err) {
+            console.warn(`[MoE] Failed to save weights: ${err}`);
+            return false;
+        }
+    }
+    /**
+     * Reset weights to random initialization
+     */
+    resetWeights() {
+        this.W1 = xavierInit(INPUT_DIM, HIDDEN_DIM);
+        this.b1.fill(0);
+        this.W2 = xavierInit(HIDDEN_DIM, NUM_EXPERTS);
+        this.b2.fill(0);
+        this.resetStats();
+    }
+    // ============================================================================
+    // Private Methods
+    // ============================================================================
+    /**
+     * Select top-k indices from probabilities
+     */
+    selectTopK(probs, k) {
+        // Create index-value pairs and sort by value descending
+        const pairs = [];
+        for (let i = 0; i < probs.length; i++) {
+            pairs.push([i, probs[i]]);
+        }
+        pairs.sort((a, b) => b[1] - a[1]);
+        // Return top-k indices
+        return pairs.slice(0, k).map((p) => p[0]);
+    }
+    /**
+     * Compute load balance loss for regularization
+     *
+     * Uses auxiliary loss from Switch Transformer:
+     * L_balance = N * sum(f_i * P_i)
+     * where f_i = fraction of tokens routed to expert i
+     *       P_i = average routing probability to expert i
+     */
+    computeLoadBalanceLoss() {
+        if (this.totalRoutings === 0)
+            return 0;
+        let loss = 0;
+        for (let i = 0; i < NUM_EXPERTS; i++) {
+            const fraction = this.routingCounts[i] / this.totalRoutings;
+            const avgProb = this.probs[i]; // Current routing prob
+            loss += fraction * avgProb;
+        }
+        return NUM_EXPERTS * loss * this.config.loadBalanceCoef;
+    }
+    /**
+     * Compute Gini coefficient for load distribution
+     */
+    computeGiniCoefficient() {
+        if (this.totalRoutings === 0)
+            return 0;
+        // Sort counts
+        const sorted = Array.from(this.routingCounts).sort((a, b) => a - b);
+        const n = sorted.length;
+        const mean = this.totalRoutings / n;
+        // Compute Gini using the formula: G = (2 * sum(i * x_i) - (n+1) * sum(x_i)) / (n * sum(x_i))
+        let weightedSum = 0;
+        for (let i = 0; i < n; i++) {
+            weightedSum += (i + 1) * sorted[i];
+        }
+        const gini = (2 * weightedSum - (n + 1) * this.totalRoutings) /
+            (n * this.totalRoutings + 1e-8);
+        return Math.max(0, gini);
+    }
+}
+// ============================================================================
+// Singleton Instance
+// ============================================================================
+let moeRouterInstance = null;
+let _moeInitPromise = null;
+/**
+ * Get singleton MoE router instance
+ *
+ * @param config - Optional configuration (only used on first call)
+ * @returns MoE router instance
+ */
+export function getMoERouter(config) {
+    if (!moeRouterInstance) {
+        moeRouterInstance = new MoERouter(config);
+        _moeInitPromise = moeRouterInstance.initialize().catch((err) => {
+            console.warn('[MoE] Failed to initialize router:', err);
+        });
+    }
+    return moeRouterInstance;
+}
+/**
+ * Get singleton MoE router instance, waiting for initialization to complete
+ *
+ * @param config - Optional configuration (only used on first call)
+ * @returns Promise resolving to fully initialized MoE router instance
+ */
+export async function getMoERouterReady(config) {
+    const router = getMoERouter(config);
+    if (_moeInitPromise)
+        await _moeInitPromise;
+    return router;
+}
+/**
+ * Reset singleton instance (for testing)
+ */
+export function resetMoERouter() {
+    moeRouterInstance = null;
+}
+/**
+ * Factory function to create new router
+ */
+export function createMoERouter(config) {
+    return new MoERouter(config);
+}
+//# sourceMappingURL=moe-router.js.map