Metadata-Version: 2.4
Name: cognitive-aug
Version: 0.2.1
Summary: Modular brain-inspired cognitive augmentation layers for neural networks
Author: Ashley Allen
License: MIT
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: numpy>=1.20.0
Requires-Dist: torch>=2.0.0
Provides-Extra: dev
Requires-Dist: black>=23.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Description-Content-Type: text/markdown

# Cognitive Augmentation Layers (`cognitive_aug`)

Modular, brain-inspired cognitive augmentation layers for neural networks, built atop PyTorch.

This library provides plug-and-play components to endow existing deep learning architectures with biologically inspired capabilities, starting with **Global Workspace Theory (GWT)** mechanisms for attention, integration, and broad dynamic routing.

---

## Features (Phase v0.1 MVP)

* **Module Registry**: Track and coordinate active neural networks acting as specialized cognitive modules.
* **Data Flow Manager**: High-performance, framework-integrated routing of latent representations across dynamic computation cycles.
* **Global Workspace Theory (GWT) Layer**:
  - Configurable `single-slot` (high biological fidelity) and `multi-slot` (engineering focus) workspace bottleneck.
  - Pluggable `AttentionSelector` supporting Top-Down Key-Query matching and Bottom-Up salience selection.
  - **Dynamic Ignition**: Non-linear, threshold-gated attentional selection where only highly active features are broadcast.
  - `BroadcastEngine` that distributes unified cognitive states back to all modules as context.
* **Non-Intrusive Wrappers**: `ModuleAdapter` wraps standard PyTorch `nn.Module`s using forward/backward hooks without polluting or altering original model classes.
* **Optimized Cognitive Add-ons**:
  - **Differentiable Selection**: `CosineSimilaritySelector`, `VectorizedCrossAttentionSelector` (using native FlashAttention speeds), and `EfficientGumbelSoftmaxSelector` (hard winner-take-all routing).
  - **Low-Overhead Salience**: `MagnitudeSalience` (L2 norm), `EntropySalience` (Shannon entropy confidence), and stateful `TemporalSurpriseSalience` (temporal cosine distance tracking).
  - **Decay Working Memory**: `DecayWorkingMemory` stateful wrapper with in-place exponential decay mutations and blended historical traces.
  - **Parallel Downstream Routing**: `CognitiveOutputRouter` broadcasting workspace representations back into multiple output heads in a single vectorized matrix projection pass.

---

## Installation

To install in editable mode with development dependencies:

```bash
pip install -e ".[dev]"
```

---

## Quickstart Example

Here is a simple example showing how to register an image classification module and a text processing module, enabling them to communicate via a shared Global Workspace:

```python
import torch
import torch.nn as nn
from cognitive_aug import CognitiveAugEngine, GlobalWorkspace, ModuleAdapter

# 1. Define your standard PyTorch modules
class VisualModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.fc = nn.Linear(16 * 8 * 8, 256)
    def forward(self, x):
        features = torch.relu(self.conv(x))
        features = features.view(features.size(0), -1)
        return self.fc(features)

class TextModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(1000, 64)
        self.lstm = nn.LSTM(64, 128, batch_first=True)
        self.fc = nn.Linear(128, 256)
    def forward(self, x):
        emb = self.emb(x)
        out, _ = self.lstm(emb)
        return self.fc(out[:, -1, :])

# 2. Instantiate modules
vis_mod = VisualModule()
txt_mod = TextModule()

# 3. Create the Cognitive Engine and Register Modules
engine = CognitiveAugEngine()

# Register modules with their latent dimensions (must match workspace target or be mapped)
vis_adapter = engine.register_module(
    name="vision",
    module=vis_mod,
    latent_dim=256
)

txt_adapter = engine.register_module(
    name="text",
    module=txt_mod,
    latent_dim=256
)

# 4. Instantiate Global Workspace
workspace = GlobalWorkspace(
    latent_dim=256,
    attention_type="key-query",
    ignition_threshold=0.5,
    workspace_slots=1  # Biological single-slot GWT bottleneck
)

# Attach workspace to the engine
engine.attach_workspace(workspace)

# 5. Run standard forward propagation
# Adapters will automatically capture latent representations via PyTorch forward hooks!
dummy_img = torch.randn(4, 3, 8, 8)
dummy_txt = torch.randint(0, 1000, (4, 10))

# Forward pass as usual
vis_out = vis_mod(dummy_img)
txt_out = txt_mod(dummy_txt)

# Perform GWT cycle: Selection & Global Broadcast!
workspace_state = engine.step()

print("Global Workspace broadcast vector shape:", workspace_state.shape)
```

---

## High-Performance Modular Add-ons

For advanced cognitive systems or performance-critical setups, the package includes highly optimized, zero-copy, fully vectorized components.

### 1. Advanced Attention Selectors (`selectors.py`)
Replace the default selector with one of three high-speed subclasses inheriting from `BaseSelector`:
- **`CosineSimilaritySelector`**: Uses vectorized matrix-based `cosine_similarity` to rapidly match incoming states to top-down goals.
- **`VectorizedCrossAttentionSelector`**: Uses PyTorch's native `scaled_dot_product_attention` with the Value tensor designed as an identity matrix, unlocking native FlashAttention speeds while remaining fully differentiable.
- **`EfficientGumbelSoftmaxSelector`**: A hard, winner-take-all routing mechanism using single-pass `gumbel_softmax` that retains full differentiability.

### 2. Low-Overhead Salience Metrics (`salience.py`)
Evaluate modular activation to compute confidence and decide workspace ignition:
- **`MagnitudeSalience`**: Computes L2 norms of states in a single vectorized pass using `torch.linalg.vector_norm`.
- **`EntropySalience`**: Computes Shannon entropy normalized to $[0, 1]$ confidence, penalizing noisy, unconfident representations.
- **`TemporalSurpriseSalience`**: Stateful cache that detaches gradients across iterations and uses cosine distance to measure step-to-step temporal shifts.
- **`global_pool_latent`**: A dimension-agnostic helper that automatically maps spatial/temporal latent tensor dimensions (e.g. `[B, T, D]` or `[B, C, H, W]`) into `[B, D]` vectors, avoiding sequential loops.

### 3. Short-Term Decay Working Memory (`memory.py`)
- **`DecayWorkingMemory`**: Exponentially decays inactive workspace slots in-place (`workspace_state.mul_(decay_rate)`) and returns a blended context vector of the active new winner and the decaying trace of the past.

### 4. Parallel Output Router (`routing.py`)
- **`CognitiveOutputRouter`**: Maps unified workspace representations back into dedicated output heads simultaneously using a **single parallel linear projection layer** and returns views using zero-copy slicing.

For a full demo showing how to attach and run these components in an active training loop, see [example_addons.py](file:///c:/Users/ASHLEY%20ALLEN/OneDrive/pypack/example_addons.py).

---

## License

This project is licensed under the MIT License - see the LICENSE details.
