Add PyTorch neural network modules for GPU-accelerated eigen similarity#3
Conversation
Implements three key PyTorch modules building on Lorentz-invariant similarity: 1. gpu_similarity.py: - Batched eigen_similarity for GPU acceleration - Supports 2D (B,D)x(N,D) and 3D (B,L,D)x(B,L,D) tensors - Modified Minkowski embedding: (u, 1) instead of (u, ||u||) - Explicit self-similarity detection returns 0.0 (loop prevention) - Different vectors get informative non-zero similarities 2. eigen_memory.py (EigenMemory): - External memory bank with Lorentz-invariant retrieval - Ring-buffer storage (configurable capacity) - Top-k retrieval with softmax attention weights - Temporal decay favoring recent entries - Loop prevention via epsilon threshold 3. eigen_attention.py (EigenAttention): - Multi-head attention using eigen similarity instead of dot-product - Self-attention suppression via loop_epsilon - Optional causal masking for autoregressive models - Negative similarity suppression (disconnected tokens) - Fully differentiable with gradient flow Test Suite (test_pytorch_modules.py): - 33 tests covering all modules (30 passed, 3 skipped on CPU) - GPU similarity validation (self-sim = 0.0) - Memory read/write operations - Attention mechanics (causal masking, external masks) - Gradient flow verification - Integration tests (memory + attention pipeline) - Consistency with NumPy implementation Key Innovation: Modified Lorentz embedding (u, 1) makes <u,u>_L = ||u||² - 1 (non-zero), allowing informative cross-similarities while still detecting true self-reference (same vector) and setting it to 0.0 for loop prevention. Dependencies: - torch>=2.0.0 added to requirements.txt - Backward compatible with existing NumPy implementation Documentation: - README.md updated with PyTorch module usage examples - Comprehensive docstrings in all modules - Test suite documentation
There was a problem hiding this comment.
Pull Request Overview
This PR adds GPU-accelerated PyTorch neural network modules implementing Lorentz-invariant (EigenFunction) similarity for deep learning applications. The similarity metric naturally yields zero self-similarity, preventing loop amplification in recursive architectures like attention mechanisms and memory retrieval.
Key changes:
- Implements batched, differentiable similarity functions with automatic differentiation support
- Adds EigenMemory module with ring-buffer storage and temporal decay
- Adds EigenAttention multi-head attention layer with causal masking support
Reviewed Changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| test_pytorch_modules.py | Comprehensive test suite covering GPU similarity functions, EigenMemory, EigenAttention, and integration tests |
| gpu_similarity.py | Core GPU-accelerated similarity computations supporting 2D and 3D tensors |
| eigen_memory.py | External memory module with top-k retrieval and loop prevention |
| eigen_attention.py | Multi-head attention using Lorentz-invariant similarity instead of dot-product |
| requirements.txt | Adds PyTorch 2.0+ dependency |
| README.md | Documentation for new PyTorch modules with usage examples |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| u -> (u, ||u||), v -> (v, ||v||) | ||
|
|
||
| Lorentz inner product: | ||
| <u, v>_L = u·v - ||u|| * ||v|| | ||
|
|
||
| Self-product (lightlike): | ||
| <u, u>_L = ||u||² - ||u||² = 0 |
There was a problem hiding this comment.
The documentation describes the embedding as u -> (u, ||u||) (lines 30-31), but the actual implementation uses u -> (u, 1) with a constant time component (see lines 111, 135-137). The documentation should be updated to match the implementation: u -> (u, 1) in Minkowski space, making <u,u>_L = ||u||² - 1 instead of <u,u>_L = ||u||² - ||u||².
| u -> (u, ||u||), v -> (v, ||v||) | |
| Lorentz inner product: | |
| <u, v>_L = u·v - ||u|| * ||v|| | |
| Self-product (lightlike): | |
| <u, u>_L = ||u||² - ||u||² = 0 | |
| u -> (u, 1), v -> (v, 1) | |
| Lorentz inner product: | |
| <u, v>_L = u·v - 1 | |
| Self-product: | |
| <u, u>_L = ||u||² - 1 |
|
|
||
| def compare_self_similarity_torch( | ||
| v: torch.Tensor, | ||
| ) -> dict[str, torch.Tensor]: |
There was a problem hiding this comment.
The return type annotation dict[str, torch.Tensor] uses PEP 585 syntax requiring Python 3.9+, while the union operator | on line 27 of eigen_memory.py requires Python 3.10+. Since PyTorch 2.0 supports Python 3.8+, this creates a compatibility issue. Consider using from typing import Dict, Union and writing Dict[str, torch.Tensor] to support Python 3.8+, or document that Python 3.10+ is required.
| sim_threshold: float = 0.0, | ||
| loop_epsilon: float = 1e-3, | ||
| decay: float = 0.99, | ||
| device: str | torch.device | None = None, |
There was a problem hiding this comment.
The type annotation uses the | union operator which requires Python 3.10+. Since PyTorch 2.0 supports Python 3.8+, this creates a compatibility issue. Consider using from typing import Union, Optional and writing Optional[Union[str, torch.device]] to support Python 3.8+, or document that Python 3.10+ is required.
| if D != D_k: | ||
| raise ValueError(f"Dimension mismatch: q has dim {D}, k has dim {D_k}") | ||
|
|
||
| B = B_q |
There was a problem hiding this comment.
Variable B is not used.
| B = B_q |
|
|
||
| import pytest | ||
| import torch | ||
| import torch.nn as nn |
There was a problem hiding this comment.
Import of 'nn' is not used.
| import torch.nn as nn |
This merge resolves conflicts by providing BOTH GPU acceleration options: 1. PyTorch Implementation (kept from this branch): - gpu_similarity.py: PyTorch tensors with automatic differentiation - eigen_attention.py: Multi-head attention module - eigen_memory.py: External memory module - test_pytorch_modules.py: Comprehensive PyTorch tests 2. CuPy Implementation (from main, renamed): - cupy_similarity.py: NumPy-compatible CUDA acceleration - examples_gpu.py: CuPy usage examples - test_gpu_similarity.py: CuPy GPU tests Resolution Details: - Renamed main's gpu_similarity.py -> cupy_similarity.py - Kept this branch's gpu_similarity.py (PyTorch version) - Updated requirements.txt: Both PyTorch and CuPy documented - Updated README.md: Two GPU acceleration options documented Users can now choose: - PyTorch: For deep learning with attention/memory modules - CuPy: For raw CUDA performance with NumPy API Both implementations provide Lorentz-invariant similarity with self-similarity = 0.0 for loop prevention. Files added from main: - cupy_similarity.py (renamed from gpu_similarity.py) - examples_gpu.py - test_gpu_similarity.py Files preserved from this branch: - gpu_similarity.py (PyTorch) - eigen_attention.py - eigen_memory.py - test_pytorch_modules.py
No description provided.