Titan Cognitive Memory

🚀 Titan Memory Server Implementation

This is a collaborative project with @jasonkneen and @ExpressionsBot.

Follow us on Twitter:

• jasonkneen
• megaprompt

Inspired by the Google Research paper "Generative AI for Programming: A Common Task Framework", this server implements a neural memory system. It can learn and predict sequences, maintain states through memory vectors, and adheres to the principles proposed in the paper to enhance code generation and understanding.

📚 Research Background

This implementation draws on the concepts proposed in the paper:

• Enhanced sequence learning and memory capabilities
• Use of the Surprise mechanism for novelty detection
• Adoption of manifold optimization for a stable training process
• Maintenance of states through memory vectors

These features align with the paper's goal of improving code understanding and generation through better memory and state management.

✨ Features

• A neural memory model with configurable dimensions
• Sequence learning and prediction capabilities
• Surprise calculation for novelty detection
• Model persistence (save/load)
• State memory management
• Full integration of MCP tools

📦 Installation

# Install dependencies
npm install

# Build the project
npm run build

# Run tests
npm test

🛠️ Available MCP Tools

1. 🎯 init_model

Initialize the Titan memory model with custom configurations.

{
  inputDim?: number;  // Input dimension (default: 64)
  outputDim?: number; // Output/memory dimension (default: 64)
}

2. 📚 train_step

Perform a single-step training using the current and next state vectors.

{
  x_t: number[];    // Current state vector
  x_next: number[]; // Next state vector
}

3. 💡 forward_pass

Run a forward pass through the model using an input vector.

{
  x: number[]; // Input vector
}

4. 💾 save_model

Save the model to a specified path.

{
  path: string; // Path to save the model
}

5. 📂 load_model

Load the model from a specified path.

{
  path: string; // Path to load the model from
}

6. ℹ️ get_status

Get the current model status and configuration information.

{} // No parameter required

7. 🔄 train_sequence

Train the model using a sequence of vectors.

{
  sequence: number[][]; // Array of vectors to train on
}

💻 Usage Examples

Basic Usage

// Initialize the model
await callTool('init_model', { inputDim: 64, outputDim: 64 });

// Train with a sequence
const sequence = [
  [1, 0, 0, /* ... */],
  [0, 1, 0, /* ... */],
  [0, 0, 1, /* ... */]
];
await callTool('train_sequence', { sequence });

// Run a forward pass
const result = await callTool('forward_pass', {
  x: [1, 0, 0, /* ... */]
});

🔧 Technical Details

• Efficient tensor operations are implemented using TensorFlow.js.
• Manifold optimization is adopted for a stable training process.
• The Surprise mechanism is implemented for novelty detection.
• Appropriate tensor cleanup is used for memory management.
• Type-safe MCP tool definitions are provided.
• Operations between memory states are maintained.

🧪 Testing

The project includes comprehensive tests covering:

• Model initialization and configuration
• Training and forward pass operations
• Memory state management
• Model persistence
• Edge cases and error handling
• Tensor cleanup and memory management

To run the tests, use:

npm test

🔍 Implementation Notes

• All tensor operations are wrapped in tf.tidy() for proper memory management.
• Appropriate error handling is implemented with detailed error messages.
• Type-safe MCP tool definitions are used.
• Memory states are maintained during the training process.
• Floating-point precision issues are handled through epsilon tolerance.

📄 License

This project is under the MIT License. You are welcome to use and modify it freely!