MCP Titan

import * as tf from '@tensorflow/tfjs-node'; import { promises as fs } from 'fs'; import * as path from 'path'; import * as crypto from 'crypto'; import type { IMemoryModel, IMemoryState, IMemoryUpdateResult, IAttentionBlock, ISurpriseMetrics, IModelGradients, HopeMemoryConfig, SerializedAuxiliaryMemoryState } from '../types.js'; import type { AdvancedTokenizer } from '../tokenizer/index.js'; import { ContinuumMemory, type ContinuumMemoryConfig, type HopeMemoryState, type HierarchicalStats } from './continuum_memory.js'; import { RetentiveCore, type RetentiveCoreConfig, type RetentionState } from './retention_core.js'; import { SelectiveStateSpace } from './mamba_filters.js'; import { MemoryRouter, type RoutingDecision } from './memory_router.js'; import { DeltaCompressionHook, LayerScheduler, UpdateBuffer } from './optimizer_hooks.js'; import { tidyMemoryState } from './type_utils.js'; const DEFAULT_CONFIG: HopeMemoryConfig = { inputDim: 256, hiddenDim: 192, memoryDim: 256, shortTermSlots: 64, longTermSlots: 256, archiveSlots: 512, learningRate: 1e-3, dropoutRate: 0.1, promotionThreshold: 0.05, surpriseRetention: 0.85, routerTopK: 2, // Backward compatibility fields maxSequenceLength: 512, memorySlots: 256, transformerLayers: 6, enableMomentum: true, enableTokenFlow: true, enableForgettingGate: true, baseForgettingRate: 0.1, surpriseForgettingWeight: 0.3, consolidationInterval: 100, enableHierarchicalMemory: true, useHierarchicalMemory: true }; interface ForwardArtifacts { logits: tf.Tensor2D; memoryState: HopeMemoryState; retentionState: RetentionState; decision: RoutingDecision; } /** * HOPE Paper: Token Flow Tracking * Captures sequential dependencies beyond momentary surprise */ export interface TokenFlowState { history: number[][]; // Recent token embeddings weights: number[]; // Recency × similarity weights windowSize: number; // Sliding window (default 32) decay: number; // Temporal decay (default 0.95) } export class HopeMemoryModel implements IMemoryModel { private config: HopeMemoryConfig; private continuumMemory: ContinuumMemory; private selectiveFilter: SelectiveStateSpace; private retentiveCore: RetentiveCore; private memoryRouter: MemoryRouter; private compressionHook: DeltaCompressionHook; private layerScheduler: LayerScheduler; private updateBuffer: UpdateBuffer; private outputKernel: tf.Variable<tf.Rank.R2>; private outputBias: tf.Variable<tf.Rank.R1>; private optimizer: tf.AdamOptimizer; private retentionState?: RetentionState; private latestMemoryState: HopeMemoryState; private tokenFlowState: TokenFlowState; private tokenizer?: AdvancedTokenizer; private consolidationCounter = 0; private consolidationInterval = 100; private consolidationStats = { runs: 0, lastRun: 0 }; constructor(config: Partial<HopeMemoryConfig> = {}) { this.config = { ...DEFAULT_CONFIG, ...config }; this.buildComponents(); } public async initialize(config?: Partial<HopeMemoryConfig>): Promise<void> { if (config) { this.config = { ...this.config, ...config }; } this.consolidationInterval = this.config.consolidationInterval ?? 100; this.buildComponents(); } public createInitialState(): HopeMemoryState { this.retentionState = this.retentiveCore.initState(1); this.latestMemoryState = this.continuumMemory.initialize(); return this.latestMemoryState; } public forward(x: tf.Tensor2D, memoryState: IMemoryState): { predicted: tf.Tensor2D; memoryUpdate: IMemoryUpdateResult; } { const hopeState = memoryState as HopeMemoryState; const result = this.computeForward(x, hopeState, true); this.latestMemoryState = result.memoryState; return this.buildForwardResult(result, hopeState); } public trainStep(x_t: tf.Tensor2D, x_next: tf.Tensor2D, memoryState: IMemoryState): { loss: tf.Tensor; gradients: IModelGradients; memoryUpdate: IMemoryUpdateResult; } { const hopeState = memoryState as HopeMemoryState; const clonedState = this.continuumMemory.clone(hopeState); const { value: loss, grads } = tf.variableGrads(() => { const forwardResult = this.computeForward(x_t, clonedState, false); const target = this.ensure2d(x_next); const prediction = forwardResult.logits; const mse = tf.losses.meanSquaredError(target, prediction); return mse.mean(); }); const gradientEntries = Object.entries(grads); const gradientTensors = gradientEntries.map(([, tensor]) => tensor); const payload = this.compressionHook.compress(gradientTensors); const decompressed = this.compressionHook.decompress(payload); const activeLayers = this.layerScheduler.selectActiveLayers(decompressed); const gradientsToApply: Record<string, tf.Tensor> = {}; gradientEntries.forEach(([name, tensor], index) => { if (activeLayers.includes(index)) { gradientsToApply[name] = tensor; this.updateBuffer.push(name, tensor.clone()); } }); if (Object.keys(gradientsToApply).length === 0 && gradientEntries.length > 0) { const [fallbackName, fallbackTensor] = gradientEntries[0]; gradientsToApply[fallbackName] = fallbackTensor; } this.optimizer.applyGradients(gradientsToApply as any); const forwardResult = this.computeForward(x_t, hopeState, true); this.latestMemoryState = forwardResult.memoryState; const memoryUpdate = this.buildForwardResult(forwardResult, hopeState).memoryUpdate; return { loss, gradients: { shortTerm: tf.zerosLike(hopeState.shortTerm), longTerm: tf.zerosLike(hopeState.longTerm), meta: tf.zerosLike(hopeState.meta) }, memoryUpdate }; } public getTrainableVariables(): tf.Variable[] { return [ ...this.retentiveCore.getTrainableVariables(), ...this.memoryRouter.getTrainableVariables(), this.outputKernel, this.outputBias ]; } public applyGradients?(gradients: Map<string, tf.Tensor>): void { const grads: Record<string, tf.Tensor> = {}; gradients.forEach((tensor, key) => { grads[key] = tensor; }); this.optimizer.applyGradients(grads as any); } public getConfig(): HopeMemoryConfig { return { ...this.config }; } public resetGradients(): void { this.compressionHook.reset(); this.updateBuffer.clear(); this.optimizer = tf.train.adam(this.config.learningRate); } /** * Attach a tokenizer for higher-quality text embeddings. * Tokenizer embeddings are pooled and adapted to model inputDim. */ public attachTokenizer(tokenizer: AdvancedTokenizer): void { this.tokenizer = tokenizer; } public hydrateMemoryState(state: IMemoryState): void { this.latestMemoryState = state as HopeMemoryState; } public async pruneMemoryByInformationGain(threshold: number): Promise<HopeMemoryState> { this.latestMemoryState = this.continuumMemory.prune(this.latestMemoryState, threshold); return this.latestMemoryState; } public getPruningStats(): HierarchicalStats { return this.continuumMemory.getStats(this.latestMemoryState); } public getConsolidationStats(): { runs: number; lastRun: number } { return { ...this.consolidationStats }; } public async encodeText(text: string): Promise<tf.Tensor2D> { const cleaned = text.trim().slice(0, 4096); // Prefer the advanced tokenizer when available if (this.tokenizer) { try { const result = await this.tokenizer.encode(cleaned, { maxLength: this.config.maxSequenceLength, padding: true, truncation: true, addSpecialTokens: true, returnTensors: true }); const embeddings = result.embeddings; const pooled = tf.tidy(() => embeddings.mean(0)); // Dispose tokenizer-attached tensors we no longer need embeddings.dispose(); if ((result as any).attentionMask && typeof (result as any).attentionMask.dispose === 'function') { (result as any).attentionMask.dispose(); } const adjusted = tf.tidy(() => { const targetDim = this.config.inputDim; const pooled1d = pooled as tf.Tensor1D; if (pooled1d.shape[0] === targetDim) { return pooled1d.expandDims(0); } if (pooled1d.shape[0] > targetDim) { const sliced = pooled1d.slice([0], [targetDim]); return sliced.expandDims(0); } const padAmount = targetDim - pooled1d.shape[0]; const padded = tf.concat([pooled1d, tf.zeros([padAmount])]) as tf.Tensor1D; return padded.expandDims(0); }); pooled.dispose(); return adjusted; } catch (error) { // Fallback to hashed encoding below console.warn('Tokenizer encode failed, falling back to hashed encoding:', error); } } // Fallback: hashed bag-of-tokens mapped into inputDim with L2 normalization const tokens = cleaned .toLowerCase() .split(/[^a-z0-9]+/i) .filter(Boolean) .slice(0, this.config.inputDim * 4); // cap to avoid huge loops const vec = new Float32Array(this.config.inputDim).fill(0); const spreadsPerToken = 4; if (tokens.length === 0) { // Fallback to simple char codes if no tokens extracted const chars = Array.from(cleaned).map(c => c.codePointAt(0) ?? 0); for (let i = 0; i < Math.min(chars.length, vec.length); i += 1) { vec[i] = (chars[i] % 1024) / 1024; } } else { for (const tok of tokens) { const digest = crypto.createHash('sha256').update(tok).digest(); for (let i = 0; i < spreadsPerToken; i += 1) { const offset = (i * 4) % digest.length; const idx = digest.readUInt32BE(offset) % this.config.inputDim; vec[idx] += 1; } } } // L2 normalize to keep scale stable let norm = 0; for (let i = 0; i < vec.length; i += 1) { norm += vec[i] * vec[i]; } norm = Math.sqrt(norm) || 1; for (let i = 0; i < vec.length; i += 1) { vec[i] = vec[i] / norm; } return tf.tensor2d([Array.from(vec)], [1, this.config.inputDim]); } public async storeMemory(text: string): Promise<void> { const embedding = await this.encodeText(text); const baseState = this.latestMemoryState; const decision = this.memoryRouter.route(this.retentionState?.hidden ?? embedding, baseState); this.latestMemoryState = this.continuumMemory.write(baseState, embedding, { surprise: decision.surprise, timestamp: Date.now(), routeWeights: decision.weights }); } public exportAuxiliaryState(): SerializedAuxiliaryMemoryState | undefined { const snapshot = this.updateBuffer.flush(); if (snapshot.size === 0) { return undefined; } const tensors: Record<string, { data: number[]; shape: number[] }> = {}; snapshot.forEach((tensor, name) => { tensors[name] = { data: Array.from(tensor.dataSync()), shape: tensor.shape as number[] }; tensor.dispose(); }); return { extendedMemory: { tensors } }; } public restoreAuxiliaryState(_: SerializedAuxiliaryMemoryState | undefined): void { // No-op for now – HOPE recomputes auxiliary state during initialization. } // Alias for backward compatibility public async load(directory: string): Promise<void> { await this.loadModel(directory); } public async loadModel(directory: string): Promise<void> { const filePath = path.join(directory, 'hope_model.json'); const exists = await fs.access(filePath).then(() => true).catch(() => false); if (!exists) { return; } const raw = await fs.readFile(filePath, 'utf8'); const payload = JSON.parse(raw) as { config: HopeMemoryConfig; weights: number[][]; shapes: number[][]; }; this.config = { ...this.config, ...payload.config }; this.buildComponents(); const variables = this.getTrainableVariables(); payload.weights.forEach((values, index) => { const shape = payload.shapes[index]; if (!variables[index]) { return; } const tensor = tf.tensor(values, shape); variables[index].assign(tensor); }); } public async save(directory: string): Promise<void> { await fs.mkdir(directory, { recursive: true }); const variables = this.getTrainableVariables(); const weights = await Promise.all(variables.map(async variable => Array.from(await variable.data()))); const shapes = variables.map(variable => variable.shape as number[]); const payload = { config: this.config, weights, shapes }; await fs.writeFile(path.join(directory, 'hope_model.json'), JSON.stringify(payload)); } // Alias for IMemoryModel compatibility public async saveModel(path: string): Promise<void> { await this.save(path); } public async snapshot(): Promise<{ config: HopeMemoryConfig; weights: number[][]; shapes: number[][]; optimizer?: { weights: number[][]; shapes: number[][] }; }> { const variables = this.getTrainableVariables(); const weights = await Promise.all(variables.map(async variable => Array.from(await variable.data()))); const shapes = variables.map(variable => variable.shape as number[]); let optimizerSnapshot: { weights: number[][]; shapes: number[][] } | undefined; const optimizerWeights = await this.optimizer.getWeights(); if (optimizerWeights.length > 0) { optimizerSnapshot = { weights: await Promise.all(optimizerWeights.map(async tensor => Array.from(await tensor.data()))), shapes: optimizerWeights.map(tensor => tensor.shape as number[]) }; } return { config: { ...this.config }, weights, shapes, optimizer: optimizerSnapshot }; } public async restoreSnapshot(payload: { config: HopeMemoryConfig; weights: number[][]; shapes: number[][]; optimizer?: { weights: number[][]; shapes: number[][] }; }): Promise<void> { this.config = { ...this.config, ...payload.config }; this.buildComponents(); const variables = this.getTrainableVariables(); payload.weights.forEach((values, index) => { const shape = payload.shapes[index]; if (!variables[index]) { return; } const tensor = tf.tensor(values, shape); variables[index].assign(tensor); }); if (payload.optimizer) { const tensors = payload.optimizer.weights.map((values, index) => tf.tensor(values, payload.optimizer?.shapes[index]) ); await this.optimizer.setWeights(tensors); tensors.forEach(t => t.dispose()); } } // IMemoryModel required methods public getMemoryState(): HopeMemoryState { return this.latestMemoryState; } public resetMemory(): void { this.latestMemoryState = this.createInitialState(); this.retentionState = undefined; } public updateMetaMemory(surprise: ISurpriseMetrics, context: tf.Tensor): tf.Tensor { // For HOPE, meta memory is managed within ContinuumMemory // Return the context unchanged as this is handled internally return context; } public pruneMemory(memoryState: IMemoryState, threshold: number): IMemoryState { // Convert IMemoryState to HopeMemoryState and prune const hopeState = memoryState as unknown as HopeMemoryState; const pruned = this.continuumMemory.prune(hopeState, threshold); return pruned as unknown as IMemoryState; } public manifoldStep(base: tf.Tensor, velocity: tf.Tensor): tf.Tensor { // Simple Euler step on the manifold (base + velocity) // In future, this could implement geodesic stepping return tf.add(base, velocity); } public getMemorySnapshot(): Record<string, tf.Tensor> { const state = this.latestMemoryState; return { shortTerm: state.shortTerm, longTerm: state.longTerm, archive: state.archive || tf.zeros([1, this.config.memoryDim]), surpriseHistory: state.surpriseHistory, accessCounts: state.accessCounts }; } public restoreMemoryState(state: IMemoryState): void { this.latestMemoryState = state as unknown as HopeMemoryState; } public async recallMemory(query: string, topK = 5): Promise<tf.Tensor2D[]> { const queryTensor = await this.encodeText(query); const queryTensor2d = queryTensor.expandDims(0); // Read from memory using the router const decision = this.memoryRouter.route(queryTensor2d, this.latestMemoryState); const memoryRead = this.continuumMemory.read(this.latestMemoryState, queryTensor2d, decision.weights); // Return the memory read as a single-element array (simplified recall) return [memoryRead]; } public distillMemories(similarMemories: tf.Tensor2D[]): tf.Tensor2D { return tf.tidy(() => tf.mean(tf.stack(similarMemories), 0)); } public dispose(): void { this.getTrainableVariables().forEach(variable => variable.dispose()); this.retentionState?.hidden.dispose(); this.retentionState?.filter.carry.dispose(); this.retentionState?.filter.bandwidth.dispose(); } private computeForward(input: tf.Tensor2D, memoryState: HopeMemoryState, updateState: boolean): ForwardArtifacts { return tidyMemoryState<ForwardArtifacts>(() => { const normalizedInput = this.ensure2d(input); // HOPE Paper: Update token flow before routing this.updateTokenFlow(normalizedInput); const readWeights = this.memoryRouter.route(this.retentionState?.hidden ?? normalizedInput, memoryState); const memoryRead = this.continuumMemory.read(memoryState, normalizedInput, readWeights.weights); const coreInput = tf.concat([normalizedInput, memoryRead], 1); const retentionState = this.retentionState ?? this.retentiveCore.initState(1); const { outputs, state } = this.retentiveCore.forwardSequence(coreInput, retentionState); const logits = tf.add(tf.matMul(outputs, this.outputKernel), this.outputBias); let updatedState = memoryState; if (updateState) { // HOPE Paper: Weight surprise by token flow for sequential dependencies const weightedSurprise = this.weightSurpriseByTokenFlow(readWeights.surprise); updatedState = this.continuumMemory.write(memoryState, outputs.slice([outputs.shape[0] - 1, 0], [1, -1]), { surprise: weightedSurprise, timestamp: Date.now(), routeWeights: readWeights.weights }); this.retentionState = state; this.latestMemoryState = updatedState; this.consolidationCounter += 1; if (this.consolidationCounter % this.consolidationInterval === 0) { this.latestMemoryState = this.continuumMemory.promote(this.latestMemoryState); this.latestMemoryState = this.continuumMemory.prune(this.latestMemoryState, this.config.promotionThreshold); this.consolidationStats.runs += 1; this.consolidationStats.lastRun = Date.now(); } } return { logits, memoryState: updatedState, retentionState: state, decision: readWeights }; }); } private buildForwardResult(artifacts: ForwardArtifacts, originalState: HopeMemoryState): { predicted: tf.Tensor2D; memoryUpdate: IMemoryUpdateResult; } { const attention: IAttentionBlock = { keys: tf.zeros([1, this.config.memoryDim]), values: tf.zeros([1, this.config.memoryDim]), scores: artifacts.decision.weights }; const surprise: ISurpriseMetrics = { immediate: tf.tensor1d([artifacts.decision.surprise]), accumulated: tf.tensor1d([originalState.surpriseHistory.shape[0]]), totalSurprise: tf.tensor1d([artifacts.decision.surprise]) }; return { predicted: artifacts.logits, memoryUpdate: { newState: artifacts.memoryState, attention, surprise } }; } private ensure2d(tensor: tf.Tensor2D): tf.Tensor2D { if (tensor.rank === 2) { return tensor; } return tensor.reshape([tensor.shape[0] ?? 1, this.config.inputDim]); } /** * HOPE Paper: Update token flow tracking * Maintains a sliding window of recent embeddings with recency-weighted similarity */ private updateTokenFlow(currentEmbedding: tf.Tensor2D): void { if (!this.config.enableTokenFlow) { return; } const embedding = currentEmbedding.arraySync()[0]; // Add to history with sliding window this.tokenFlowState.history.push(embedding); if (this.tokenFlowState.history.length > this.tokenFlowState.windowSize) { this.tokenFlowState.history.shift(); } // Compute recency × similarity weights const weights = this.tokenFlowState.history.map((histEmb, i) => { const recency = Math.pow( this.tokenFlowState.decay, this.tokenFlowState.history.length - i - 1 ); const similarity = this.cosineSimilarity(embedding, histEmb); return recency * similarity; }); this.tokenFlowState.weights = weights; } /** * Compute cosine similarity between two embeddings */ private cosineSimilarity(a: number[], b: number[]): number { const minLen = Math.min(a.length, b.length); let dotProduct = 0; let normA = 0; let normB = 0; for (let i = 0; i < minLen; i++) { dotProduct += a[i] * b[i]; normA += a[i] * a[i]; normB += b[i] * b[i]; } const denominator = Math.sqrt(normA) * Math.sqrt(normB); return denominator > 0 ? dotProduct / denominator : 0; } /** * HOPE Paper: Weight surprise by token flow strength * Integrates sequential dependency into surprise calculation */ private weightSurpriseByTokenFlow(surprise: number): number { if (!this.config.enableTokenFlow || this.tokenFlowState.weights.length === 0) { return surprise; } const flowStrength = this.tokenFlowState.weights.reduce((a, b) => a + b, 0) / this.tokenFlowState.weights.length; // Flow weight factor: how much sequence context affects surprise const flowWeightFactor = 0.3; return surprise * (1 + flowWeightFactor * flowStrength); } /** * Get current token flow metrics for debugging/analysis */ public getTokenFlowMetrics(): { historySize: number; averageWeight: number; flowStrength: number } { const averageWeight = this.tokenFlowState.weights.length > 0 ? this.tokenFlowState.weights.reduce((a, b) => a + b, 0) / this.tokenFlowState.weights.length : 0; return { historySize: this.tokenFlowState.history.length, averageWeight, flowStrength: averageWeight }; } // MCP Server compatibility methods public async init_model(config: any): Promise<{ status: string }> { await this.initialize(config); return { status: 'initialized' }; } public async forward_pass(x: string | number[], memoryState?: IMemoryState): Promise<any> { let inputTensor: tf.Tensor2D; if (typeof x === 'string') { inputTensor = await this.encodeText(x); } else { inputTensor = tf.tensor2d([x]); } const state = (memoryState as HopeMemoryState) ?? this.latestMemoryState; const result = this.forward(inputTensor, state); return { predicted: await result.predicted.array(), memoryState: result.memoryUpdate.newState }; } public async train_step(x_t: string | number[], x_next: string | number[]): Promise<{ loss: number }> { let inputTensor: tf.Tensor2D; let targetTensor: tf.Tensor2D; if (typeof x_t === 'string') { inputTensor = await this.encodeText(x_t); } else { inputTensor = tf.tensor2d([x_t]); } if (typeof x_next === 'string') { targetTensor = await this.encodeText(x_next); } else { targetTensor = tf.tensor2d([x_next]); } const result = this.trainStep(inputTensor, targetTensor, this.latestMemoryState); const lossValue = await result.loss.data(); return { loss: lossValue[0] }; } public get_memory_state(): any { return { shortTerm: this.latestMemoryState.shortTerm.shape, longTerm: this.latestMemoryState.longTerm.shape, archive: this.latestMemoryState.archive.shape, surpriseHistory: this.latestMemoryState.surpriseHistory.shape }; } /** * Rebuild model components to reflect the current configuration. * Used during initialization and when loading checkpoints/configs. */ private buildComponents(): void { this.disposeTrainables(); const memoryConfig: ContinuumMemoryConfig = { memoryDim: this.config.memoryDim, shortTermSlots: this.config.shortTermSlots, longTermSlots: this.config.longTermSlots, archiveSlots: this.config.archiveSlots, promotionThreshold: this.config.promotionThreshold, surpriseRetention: this.config.surpriseRetention, momentumDecay: 0.9, enableMomentum: this.config.enableMomentum ?? true, enableForgettingGate: this.config.enableForgettingGate ?? true, baseForgettingRate: this.config.baseForgettingRate ?? 0.1, surpriseForgettingWeight: this.config.surpriseForgettingWeight ?? 0.3 }; this.continuumMemory = new ContinuumMemory(memoryConfig); this.selectiveFilter = new SelectiveStateSpace({ hiddenDim: this.config.hiddenDim, contextDim: this.config.hiddenDim, dropoutRate: this.config.dropoutRate }); const coreConfig: RetentiveCoreConfig = { inputDim: this.config.inputDim + this.config.memoryDim, hiddenDim: this.config.hiddenDim, dropoutRate: this.config.dropoutRate, chunkSize: 64 }; this.retentiveCore = new RetentiveCore(coreConfig, this.selectiveFilter); this.memoryRouter = new MemoryRouter({ hiddenDim: this.config.hiddenDim, numExperts: 3, topK: this.config.routerTopK }); this.compressionHook = new DeltaCompressionHook(); this.layerScheduler = new LayerScheduler({ maxActiveLayers: 4 }); this.updateBuffer = new UpdateBuffer(); this.outputKernel = tf.variable(tf.randomNormal([this.config.hiddenDim, this.config.inputDim])); this.outputBias = tf.variable(tf.zeros([this.config.inputDim])); this.optimizer = tf.train.adam(this.config.learningRate); this.retentionState = this.retentiveCore.initState(1); this.latestMemoryState = this.continuumMemory.initialize(); this.tokenFlowState = { history: [], weights: [], windowSize: 32, decay: 0.95 }; this.consolidationInterval = this.config.consolidationInterval ?? 100; } private disposeTrainables(): void { try { this.outputKernel?.dispose(); this.outputBias?.dispose(); if (this.retentionState?.hidden) { tf.dispose(this.retentionState.hidden); } if (this.retentionState?.cell) { tf.dispose(this.retentionState.cell); } } catch { // best-effort cleanup } } } // Export types for external use export type { HopeMemoryConfig, HopeMemoryState, HierarchicalStats, RetentionState, RoutingDecision };