import * as tf from '@tensorflow/tfjs-node';
import { promises as fs } from 'fs';
import * as path from 'path';
import type {
IMemoryModel,
IMemoryState,
IMemoryUpdateResult,
IAttentionBlock,
ISurpriseMetrics,
IModelGradients,
HopeMemoryConfig,
SerializedAuxiliaryMemoryState
} from '../types.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,
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 readonly continuumMemory: ContinuumMemory;
private readonly selectiveFilter: SelectiveStateSpace;
private readonly retentiveCore: RetentiveCore;
private readonly memoryRouter: MemoryRouter;
private readonly compressionHook: DeltaCompressionHook;
private readonly layerScheduler: LayerScheduler;
private readonly updateBuffer: UpdateBuffer;
private readonly outputKernel: tf.Variable<tf.Rank.R2>;
private readonly outputBias: tf.Variable<tf.Rank.R1>;
private optimizer: tf.AdamOptimizer;
private retentionState?: RetentionState;
private latestMemoryState: HopeMemoryState;
private tokenFlowState: TokenFlowState;
constructor(config: Partial<HopeMemoryConfig> = {}) {
this.config = { ...DEFAULT_CONFIG, ...config };
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,
// HOPE paper features
momentumDecay: 0.9,
enableMomentum: this.config.enableMomentum ?? true,
enableForgettingGate: this.config.enableForgettingGate ?? true,
baseForgettingRate: 0.1,
surpriseForgettingWeight: 0.3
};
this.continuumMemory = new ContinuumMemory(memoryConfig);
const filter = new SelectiveStateSpace({
hiddenDim: this.config.hiddenDim,
contextDim: this.config.hiddenDim,
dropoutRate: this.config.dropoutRate
});
this.selectiveFilter = filter;
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, filter);
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
};
}
public async initialize(config?: Partial<HopeMemoryConfig>): Promise<void> {
if (config) {
this.config = { ...this.config, ...config };
}
this.optimizer = tf.train.adam(this.config.learningRate);
this.retentionState = this.retentiveCore.initState(1);
this.latestMemoryState = this.continuumMemory.initialize();
}
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);
}
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 async encodeText(text: string): Promise<tf.Tensor2D> {
const tokens = Array.from(text).map(char => char.codePointAt(0) ?? 0);
const normalized = new Array(this.config.inputDim).fill(0);
for (let i = 0; i < Math.min(tokens.length, this.config.inputDim); i += 1) {
normalized[i] = (tokens[i] % 1024) / 1024;
}
return tf.tensor2d([normalized]);
}
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 };
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 as tf.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);
}
// 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: number = 5): Promise<tf.Tensor2D[]> {
const queryTensor = await this.encodeText(query);
const queryTensor2d = queryTensor.expandDims(0) as tf.Tensor2D;
// 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 as tf.Tensor2D];
}
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) as tf.Tensor2D;
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;
}
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]) as tf.Tensor2D;
}
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]) as tf.Tensor2D;
}
if (typeof x_next === 'string') {
targetTensor = await this.encodeText(x_next);
} else {
targetTensor = tf.tensor2d([x_next]) as tf.Tensor2D;
}
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
};
}
}
// Export types for external use
export type { HopeMemoryConfig, HopeMemoryState, HierarchicalStats, RetentionState, RoutingDecision };
