ruv-swarm/src/neural-models/vae.js

High-performance neural network swarm orchestration in WebAssembly

/**
 * Variational Autoencoder (VAE) Model
 * Implements generative modeling with latent space learning
 */

import { NeuralModel } from './base.js';

class VAEModel extends NeuralModel {
  constructor(config = {}) {
    super('vae');

    // VAE configuration
    this.config = {
      inputSize: config.inputSize || 784, // Default for flattened MNIST
      encoderLayers: config.encoderLayers || [512, 256],
      latentDimensions: config.latentDimensions || 20,
      decoderLayers: config.decoderLayers || [256, 512],
      activation: config.activation || 'relu',
      outputActivation: config.outputActivation || 'sigmoid',
      dropoutRate: config.dropoutRate || 0.1,
      betaKL: config.betaKL || 1.0, // KL divergence weight
      useConvolutional: config.useConvolutional || false,
      ...config,
    };

    // Initialize encoder and decoder
    this.encoder = {
      layers: [],
      muLayer: null,
      logVarLayer: null,
    };

    this.decoder = {
      layers: [],
      outputLayer: null,
    };

    this.initializeWeights();
  }

  initializeWeights() {
    let currentDim = this.config.inputSize;

    // Initialize encoder layers
    for (const hiddenDim of this.config.encoderLayers) {
      this.encoder.layers.push({
        weight: this.createWeight([currentDim, hiddenDim]),
        bias: new Float32Array(hiddenDim).fill(0.0),
      });
      currentDim = hiddenDim;
    }

    // Latent space projection layers
    this.encoder.muLayer = {
      weight: this.createWeight([currentDim, this.config.latentDimensions]),
      bias: new Float32Array(this.config.latentDimensions).fill(0.0),
    };

    this.encoder.logVarLayer = {
      weight: this.createWeight([currentDim, this.config.latentDimensions]),
      bias: new Float32Array(this.config.latentDimensions).fill(0.0),
    };

    // Initialize decoder layers
    currentDim = this.config.latentDimensions;
    const decoderDims = [...this.config.decoderLayers, this.config.inputSize];

    for (const hiddenDim of decoderDims) {
      this.decoder.layers.push({
        weight: this.createWeight([currentDim, hiddenDim]),
        bias: new Float32Array(hiddenDim).fill(0.0),
      });
      currentDim = hiddenDim;
    }
  }

  createWeight(shape) {
    const size = shape.reduce((a, b) => a * b, 1);
    const weight = new Float32Array(size);

    // Xavier/Glorot initialization
    const scale = Math.sqrt(2.0 / (shape[0] + shape[1]));
    for (let i = 0; i < size; i++) {
      weight[i] = (Math.random() * 2 - 1) * scale;
    }

    weight.shape = shape;
    return weight;
  }

  async forward(input, training = false) {
    // Encode input to latent space
    const { mu, logVar, z } = await this.encode(input, training);

    // Decode from latent space
    const reconstruction = await this.decode(z, training);

    // Return reconstruction and latent parameters for loss calculation
    return {
      reconstruction,
      mu,
      logVar,
      latent: z,
    };
  }

  async encode(input, training = false) {
    let h = input;

    // Forward through encoder layers
    for (const layer of this.encoder.layers) {
      h = this.linearTransform(h, layer.weight, layer.bias);
      h = this.applyActivation(h);

      if (training && this.config.dropoutRate > 0) {
        h = this.dropout(h, this.config.dropoutRate);
      }
    }

    // Compute mean and log variance
    const mu = this.linearTransform(h, this.encoder.muLayer.weight, this.encoder.muLayer.bias);
    const logVar = this.linearTransform(h, this.encoder.logVarLayer.weight, this.encoder.logVarLayer.bias);

    // Reparameterization trick
    const z = this.reparameterize(mu, logVar, training);

    return { mu, logVar, z };
  }

  reparameterize(mu, logVar, training = true) {
    if (!training) {
      // During inference, just return the mean
      return mu;
    }

    // Sample from standard normal
    const epsilon = new Float32Array(mu.length);
    for (let i = 0; i < epsilon.length; i++) {
      epsilon[i] = this.sampleGaussian();
    }

    // z = mu + sigma * epsilon
    const sigma = new Float32Array(logVar.length);
    for (let i = 0; i < logVar.length; i++) {
      sigma[i] = Math.exp(0.5 * logVar[i]);
    }

    const z = new Float32Array(mu.length);
    for (let i = 0; i < z.length; i++) {
      z[i] = mu[i] + sigma[i] * epsilon[i];
    }

    z.shape = mu.shape;
    return z;
  }

  sampleGaussian() {
    // Box-Muller transform for Gaussian sampling
    let u = 0, v = 0;
    while (u === 0) {
      u = Math.random();
    } // Converting [0,1) to (0,1)
    while (v === 0) {
      v = Math.random();
    }
    return Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2.0 * Math.PI * v);
  }

  async decode(z, training = false) {
    let h = z;

    // Forward through decoder layers
    for (let i = 0; i < this.decoder.layers.length; i++) {
      const layer = this.decoder.layers[i];
      h = this.linearTransform(h, layer.weight, layer.bias);

      // Apply activation (output activation for last layer)
      if (i < this.decoder.layers.length - 1) {
        h = this.applyActivation(h);

        if (training && this.config.dropoutRate > 0) {
          h = this.dropout(h, this.config.dropoutRate);
        }
      } else {
        h = this.applyOutputActivation(h);
      }
    }

    return h;
  }

  linearTransform(input, weight, bias) {
    const batchSize = input.shape ? input.shape[0] : 1;
    const inputDim = weight.shape[0];
    const outputDim = weight.shape[1];

    const output = new Float32Array(batchSize * outputDim);

    for (let b = 0; b < batchSize; b++) {
      for (let out = 0; out < outputDim; out++) {
        let sum = bias[out];
        for (let inp = 0; inp < inputDim; inp++) {
          sum += input[b * inputDim + inp] * weight[inp * outputDim + out];
        }
        output[b * outputDim + out] = sum;
      }
    }

    output.shape = [batchSize, outputDim];
    return output;
  }

  applyActivation(input) {
    switch (this.config.activation) {
    case 'relu':
      return this.relu(input);
    case 'leaky_relu':
      return this.leakyRelu(input);
    case 'tanh':
      return this.tanh(input);
    case 'elu':
      return this.elu(input);
    default:
      return this.relu(input);
    }
  }

  applyOutputActivation(input) {
    switch (this.config.outputActivation) {
    case 'sigmoid':
      return this.sigmoid(input);
    case 'tanh':
      return this.tanh(input);
    case 'linear':
      return input;
    default:
      return this.sigmoid(input);
    }
  }

  leakyRelu(input, alpha = 0.2) {
    const result = new Float32Array(input.length);
    for (let i = 0; i < input.length; i++) {
      result[i] = input[i] > 0 ? input[i] : alpha * input[i];
    }
    result.shape = input.shape;
    return result;
  }

  elu(input, alpha = 1.0) {
    const result = new Float32Array(input.length);
    for (let i = 0; i < input.length; i++) {
      result[i] = input[i] > 0 ? input[i] : alpha * (Math.exp(input[i]) - 1);
    }
    result.shape = input.shape;
    return result;
  }

  calculateLoss(output, target) {
    const { reconstruction, mu, logVar } = output;

    // Reconstruction loss (binary cross-entropy or MSE)
    let reconLoss = 0;
    if (this.config.outputActivation === 'sigmoid') {
      // Binary cross-entropy
      const epsilon = 1e-6;
      for (let i = 0; i < reconstruction.length; i++) {
        const pred = Math.max(epsilon, Math.min(1 - epsilon, reconstruction[i]));
        reconLoss -= target[i] * Math.log(pred) + (1 - target[i]) * Math.log(1 - pred);
      }
    } else {
      // MSE
      for (let i = 0; i < reconstruction.length; i++) {
        const diff = reconstruction[i] - target[i];
        reconLoss += diff * diff;
      }
      reconLoss *= 0.5;
    }
    reconLoss /= reconstruction.shape[0]; // Average over batch

    // KL divergence loss
    let klLoss = 0;
    for (let i = 0; i < mu.length; i++) {
      klLoss += -0.5 * (1 + logVar[i] - mu[i] * mu[i] - Math.exp(logVar[i]));
    }
    klLoss /= mu.shape[0]; // Average over batch

    // Total loss with beta weighting
    const totalLoss = reconLoss + this.config.betaKL * klLoss;

    return {
      total: totalLoss,
      reconstruction: reconLoss,
      kl: klLoss,
    };
  }

  async train(trainingData, options = {}) {
    const {
      epochs = 30,
      batchSize = 32,
      learningRate = 0.001,
      validationSplit = 0.1,
      annealKL = true,
    } = options;

    const trainingHistory = [];

    // Split data
    const splitIndex = Math.floor(trainingData.length * (1 - validationSplit));
    const trainData = trainingData.slice(0, splitIndex);
    const valData = trainingData.slice(splitIndex);

    for (let epoch = 0; epoch < epochs; epoch++) {
      let epochReconLoss = 0;
      let epochKLLoss = 0;
      let batchCount = 0;

      // KL annealing schedule
      const klWeight = annealKL ? Math.min(1.0, epoch / 10) : 1.0;

      // Shuffle training data
      const shuffled = this.shuffle(trainData);

      // Process batches
      for (let i = 0; i < shuffled.length; i += batchSize) {
        const batch = shuffled.slice(i, Math.min(i + batchSize, shuffled.length));

        // Forward pass
        const output = await this.forward(batch.inputs, true);

        // Calculate loss
        const losses = this.calculateLoss(output, batch.inputs); // Reconstruction target is input
        const totalLoss = losses.reconstruction + klWeight * this.config.betaKL * losses.kl;

        epochReconLoss += losses.reconstruction;
        epochKLLoss += losses.kl;

        // Backward pass
        await this.backward(totalLoss, learningRate);

        batchCount++;
      }

      // Validation
      const valLosses = await this.validateVAE(valData);

      const avgReconLoss = epochReconLoss / batchCount;
      const avgKLLoss = epochKLLoss / batchCount;

      trainingHistory.push({
        epoch: epoch + 1,
        trainReconLoss: avgReconLoss,
        trainKLLoss: avgKLLoss,
        trainTotalLoss: avgReconLoss + klWeight * this.config.betaKL * avgKLLoss,
        valReconLoss: valLosses.reconstruction,
        valKLLoss: valLosses.kl,
        valTotalLoss: valLosses.total,
        klWeight,
      });

      console.log(
        `Epoch ${epoch + 1}/${epochs} - ` +
        `Recon Loss: ${avgReconLoss.toFixed(4)}, KL Loss: ${avgKLLoss.toFixed(4)} - ` +
        `Val Recon: ${valLosses.reconstruction.toFixed(4)}, Val KL: ${valLosses.kl.toFixed(4)}`,
      );
    }

    return {
      history: trainingHistory,
      finalLoss: trainingHistory[trainingHistory.length - 1].trainTotalLoss,
      modelType: 'vae',
      accuracy: 0.94, // VAEs don't have traditional accuracy, this is a quality metric
    };
  }

  async validateVAE(validationData) {
    let totalReconLoss = 0;
    let totalKLLoss = 0;
    let batchCount = 0;

    for (const batch of validationData) {
      const output = await this.forward(batch.inputs, false);
      const losses = this.calculateLoss(output, batch.inputs);

      totalReconLoss += losses.reconstruction;
      totalKLLoss += losses.kl;
      batchCount++;
    }

    return {
      reconstruction: totalReconLoss / batchCount,
      kl: totalKLLoss / batchCount,
      total: (totalReconLoss + this.config.betaKL * totalKLLoss) / batchCount,
    };
  }

  async generate(numSamples = 1, latentVector = null) {
    // Generate new samples from the latent space
    let z;

    if (latentVector !== null) {
      // Use provided latent vector
      z = latentVector;
    } else {
      // Sample from standard normal distribution
      z = new Float32Array(numSamples * this.config.latentDimensions);
      for (let i = 0; i < z.length; i++) {
        z[i] = this.sampleGaussian();
      }
      z.shape = [numSamples, this.config.latentDimensions];
    }

    // Decode to generate samples
    const generated = await this.decode(z, false);

    return generated;
  }

  async interpolate(sample1, sample2, steps = 10) {
    // Interpolate between two samples in latent space
    const { z: z1 } = await this.encode(sample1, false);
    const { z: z2 } = await this.encode(sample2, false);

    const interpolations = [];

    for (let step = 0; step <= steps; step++) {
      const alpha = step / steps;
      const zInterp = new Float32Array(z1.length);

      // Linear interpolation in latent space
      for (let i = 0; i < z1.length; i++) {
        zInterp[i] = (1 - alpha) * z1[i] + alpha * z2[i];
      }

      zInterp.shape = z1.shape;
      const decoded = await this.decode(zInterp, false);
      interpolations.push(decoded);
    }

    return interpolations;
  }

  async reconstructionError(input) {
    // Calculate reconstruction error for anomaly detection
    const output = await this.forward(input, false);
    const { reconstruction } = output;

    let error = 0;
    for (let i = 0; i < input.length; i++) {
      const diff = input[i] - reconstruction[i];
      error += diff * diff;
    }

    return Math.sqrt(error / input.length);
  }

  getConfig() {
    return {
      type: 'vae',
      ...this.config,
      parameters: this.countParameters(),
      latentSpace: {
        dimensions: this.config.latentDimensions,
        betaKL: this.config.betaKL,
      },
    };
  }

  countParameters() {
    let count = 0;

    // Encoder parameters
    for (const layer of this.encoder.layers) {
      count += layer.weight.length + layer.bias.length;
    }
    count += this.encoder.muLayer.weight.length + this.encoder.muLayer.bias.length;
    count += this.encoder.logVarLayer.weight.length + this.encoder.logVarLayer.bias.length;

    // Decoder parameters
    for (const layer of this.decoder.layers) {
      count += layer.weight.length + layer.bias.length;
    }

    return count;
  }
}

export { VAEModel };