Initial Commit

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+3D_SEMMN_final.pdf filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/data/MNIST/raw/t10k-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/data/MNIST/raw/train-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/manifold_pca.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_clusters_3d.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_module_audio.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_module_hub.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_module_other.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_module_vision.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/neuron_modules_all.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/route_map_audio_2d.png filter=lfs diff=lfs merge=lfs -text
+3D-SEMMN/outputs/figures/route_map_vision_2d.png filter=lfs diff=lfs merge=lfs -text

3D-SEMMN/README.md

@@ -0,0 +1,105 @@
+# 3D-SEMMN (Prototype)
+
+Minimal but fully functional PyTorch 2.3+ prototype of a **3D Spatially-Embedded Multimodal Manifold Network** with:
+
+- fixed 3D neuron coordinates (`100 x 100 x 50` in `config/default.yaml`)
+- 100 columnar modules (dense intra-column, sparse inter-column)
+- exponential distance-penalized sparse connectivity (`torch.sparse`)
+- Hebbian/STDP update every 100 steps
+- shared perceptual manifold hub with cross-modal attention
+- manifold regularization (latent `d=8`) + geometric twist layer (Cayley rotation)
+- multimodal losses: classification + contrastive + reconstruction + imagination cycle
+- ablations for spatial penalty and manifold loss
+
+This code is designed for readability and modular experimentation rather than full biological fidelity.
+
+## Biological Inspirations
+
+- **seRNN (Achterberg et al., 2023)**: spatial embedding + wiring-cost constraints for modular sparse topology.
+- **MSeNN-style multimodal ideas (2021)**: cross-modal coupling and biologically motivated integration.
+- **Gallego et al. (2017)**: low-dimensional population manifolds.
+
+## Project Layout
+
+- `semmn/config.py`: dataclass config + YAML loader
+- `semmn/grid.py`: 3D coordinates, lobe masks, column assignment
+- `semmn/connectivity.py`: sparse distance-decay edges + STDP + energy proxy
+- `semmn/columns.py`: low-rank dense intra-column dynamics
+- `semmn/encoders.py`: vision CNN and audio 1D-conv encoder
+- `semmn/manifold.py`: manifold autoencoder and geometric twist
+- `semmn/hub.py`: shared manifold hub + attention + imagination decoders
+- `semmn/losses.py`: full training objective
+- `semmn/model.py`: full SEMMN forward pass
+- `semmn/dataset.py`: AV-MNIST loading with synthetic fallback
+- `train.py`: training/eval/checkpointing
+- `visualize.py`: 3D grid, manifold PCA, metrics plots
+
+## Quick Start
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+Train (light config, recommended first):
+
+```bash
+python train.py --config config/light.yaml --device cuda
+```
+
+Train (full 500K-neuron spatial grid metadata):
+
+```bash
+python train.py --config config/default.yaml --device cuda
+```
+
+Generate figures:
+
+```bash
+python visualize.py \
+  --config config/light.yaml \
+  --checkpoint outputs/best.pt \
+  --history outputs/history.json \
+  --device cuda
+```
+
+## Config Notes
+
+- `config/default.yaml`: full spatial grid (`100x100x50`)
+- `config/light.yaml`: faster profile (`100x100x10`) for rapid iteration
+- Ablations:
+  - `loss.use_spatial_penalty: false`
+  - `loss.use_manifold_loss: false`
+
+## Reported Metrics
+
+`train.py` logs:
+
+- classification loss and total loss
+- validation classification accuracy
+- cross-modal retrieval accuracy (vision->audio nearest match in batch)
+- active sparse parameters (`nnz` proxy)
+- energy proxy (`mean(|w| * distance)`)
+
+## Extension Notes
+
+### Full spiking variant (Brian2 / snnTorch)
+
+1. Replace rate-based state updates (`tanh`) with LIF/AdEx neuron dynamics.
+2. Convert STDP update from rate Hebbian correlation to spike-timing windows.
+3. Keep 3D coordinates and sparse wiring graph unchanged for fair ablations.
+
+### Neuromorphic export direction
+
+1. Map columns to compute cores (Loihi/SpiNNaker style partitioning).
+2. Convert sparse inter-column edges to routing tables.
+3. Quantize state variables and weights (8-bit or mixed precision).
+4. Use event-driven updates for power-efficient deployment.
+
+## Practical Runtime Guidance
+
+- Start with `config/light.yaml` to confirm setup and debug quickly.
+- Increase `training.batch_size` and `model.recurrent_steps` gradually.
+- Keep STDP interval at `100` to limit plasticity overhead.
+- Use mixed precision (`training.amp: true`) on CUDA devices.

3D-SEMMN/config/default.yaml

@@ -0,0 +1,51 @@
+seed: 42
+grid:
+  dims: [100, 100, 50]
+  n_columns_xy: [10, 10]
+  vision_y_ratio: 0.30
+  audio_y_ratio: 0.70
+  hub_y_range: [0.40, 0.60]
+  hub_z_range: [0.30, 0.70]
+
+model:
+  state_dim: 32
+  latent_dim: 8
+  vision_feature_dim: 256
+  audio_feature_dim: 256
+  recurrent_steps: 6
+  attention_heads: 4
+
+connectivity:
+  lambda_decay: 5.0
+  inter_density: 0.08
+  intra_density_boost: 2.0
+  min_weight: 0.0
+  max_weight: 1.5
+  init_weight_scale: 0.15
+  stdp_interval: 100
+  stdp_lr: 0.01
+  anti_hebbian: 0.05
+
+loss:
+  contrastive_weight: 1.0
+  reconstruction_weight: 0.5
+  imagination_weight: 0.3
+  manifold_weight: 0.1
+  spatial_penalty_weight: 0.0005
+  temperature: 0.1
+  use_spatial_penalty: true
+  use_manifold_loss: true
+
+training:
+  dataset_source: synthetic
+  batch_size: 64
+  epochs: 10
+  lr: 0.001
+  weight_decay: 0.0001
+  num_workers: 4
+  log_interval: 50
+  val_batches: 100
+  amp: true
+
+paths:
+  output_dir: outputs

3D-SEMMN/config/light.yaml

@@ -0,0 +1,17 @@
+seed: 42
+grid:
+  dims: [100, 100, 10]
+  n_columns_xy: [10, 10]
+
+model:
+  recurrent_steps: 4
+  state_dim: 24
+
+connectivity:
+  inter_density: 0.10
+
+training:
+  batch_size: 128
+  epochs: 10
+  num_workers: 2
+  log_interval: 20

3D-SEMMN/data/MNIST/raw/t10k-images-idx3-ubyte

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fa7898d509279e482958e8ce81c8e77db3f2f8254e26661ceb7762c4d494ce7
+size 7840016

3D-SEMMN/data/MNIST/raw/t10k-images-idx3-ubyte.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6
+size 1648877

3D-SEMMN/data/MNIST/raw/t10k-labels-idx1-ubyte.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6
+size 4542

3D-SEMMN/data/MNIST/raw/train-images-idx3-ubyte

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba891046e6505d7aadcbbe25680a0738ad16aec93bde7f9b65e87a2fc25776db
+size 47040016

3D-SEMMN/data/MNIST/raw/train-images-idx3-ubyte.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609
+size 9912422

3D-SEMMN/data/MNIST/raw/train-labels-idx1-ubyte.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c
+size 28881

3D-SEMMN/outputs/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66d1a8bbb97532ad3ea4005ad99e6dbda5b5cced507c3d5c6961fb52b0b1b1a2
+size 7554761

3D-SEMMN/outputs/figures/route_edges.csv

@@ -0,0 +1,13 @@
+step,source,target,strength,modality
+0.0,80.0,90.0,0.014531210064888,vision
+0.0,10.0,0.0,0.0044663818553090096,vision
+0.0,22.0,11.0,0.003179906401783228,vision
+1.0,80.0,90.0,0.013347183354198933,vision
+1.0,10.0,0.0,0.004102454055100679,vision
+1.0,22.0,11.0,0.002920802216976881,vision
+2.0,80.0,90.0,0.012341569177806377,vision
+2.0,10.0,0.0,0.0037933634594082832,vision
+2.0,22.0,11.0,0.0027007409371435642,vision
+3.0,80.0,90.0,0.011476870626211166,vision
+3.0,10.0,0.0,0.0035275856498628855,vision
+3.0,22.0,11.0,0.0025115162134170532,vision

3D-SEMMN/outputs/figures/route_map_interactive.html

@@ -0,0 +1,7 @@
+<html>
+<head><meta charset="utf-8" /></head>
+<body>
+    <div>                        <script type="text/javascript">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>
+        <script charset="utf-8" src="https://cdn.plot.ly/plotly-2.35.2.min.js"></script>                <div id="8c11a467-ec1f-422e-a480-88b89e6970e8" class="plotly-graph-div" style="height:100%; width:100%;"></div>            <script type="text/javascript">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById("8c11a467-ec1f-422e-a480-88b89e6970e8")) {                    Plotly.newPlot(                        "8c11a467-ec1f-422e-a480-88b89e6970e8",                        [{"hovertemplate":"column=%{text}\u003cextra\u003e\u003c\u002fextra\u003e","marker":{"color":"#D0D0D0","opacity":0.35,"size":7},"mode":"markers","name":"all_columns","text":["0","1","2","3","4","5","6","7","8","9","10","11","12","13","14","15","16","17","18","19","20","21","22","23","24","25","26","27","28","29","30","31","32","33","34","35","36","37","38","39","40","41","42","43","44","45","46","47","48","49","50","51","52","53","54","55","56","57","58","59","60","61","62","63","64","65","66","67","68","69","70","71","72","73","74","75","76","77","78","79","80","81","82","83","84","85","86","87","88","89","90","91","92","93","94","95","96","97","98","99"],"x":[4.5,4.5,4.5,4.5,4.5,4.5,4.5,4.5,4.5,4.5,14.5,14.5,14.5,14.5,14.5,14.5,14.5,14.5,14.5,14.5,24.5,24.5,24.5,24.5,24.5,24.5,24.5,24.5,24.5,24.5,34.5,34.5,34.5,34.5,34.5,34.5,34.5,34.5,34.5,34.5,44.5,44.5,44.5,44.5,44.5,44.5,44.5,44.5,44.5,44.5,54.5,54.5,54.5,54.5,54.5,54.5,54.5,54.5,54.5,54.5,64.5,64.5,64.5,64.5,64.5,64.5,64.5,64.5,64.5,64.5,74.5,74.5,74.5,74.5,74.5,74.5,74.5,74.5,74.5,74.5,84.5,84.5,84.5,84.5,84.5,84.5,84.5,84.5,84.5,84.5,94.5,94.5,94.5,94.5,94.5,94.5,94.5,94.5,94.5,94.5],"y":[4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5,4.5,14.5,24.5,34.5,44.5,54.5,64.5,74.5,84.5,94.5],"type":"scatter"},{"hoverinfo":"skip","line":{"color":"#1f77b4","width":2},"mode":"lines","name":"vision_route","opacity":0.75,"x":[84.5,94.5,null,14.5,4.5,null,24.5,14.5,null,84.5,94.5,null,14.5,4.5,null,24.5,14.5,null,84.5,94.5,null,14.5,4.5,null,24.5,14.5,null,84.5,94.5,null,14.5,4.5,null,24.5,14.5,null],"y":[4.5,4.5,null,4.5,4.5,null,24.5,14.5,null,4.5,4.5,null,4.5,4.5,null,24.5,14.5,null,4.5,4.5,null,4.5,4.5,null,24.5,14.5,null,4.5,4.5,null,4.5,4.5,null,24.5,14.5,null],"type":"scatter"},{"hoverinfo":"skip","line":{"color":"#ff7f0e","width":2},"mode":"lines","name":"audio_route","opacity":0.75,"x":[],"y":[],"type":"scatter"}],                        {"template":{"data":{"histogram2dcontour":[{"type":"histogram2dcontour","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"choropleth":[{"type":"choropleth","colorbar":{"outlinewidth":0,"ticks":""}}],"histogram2d":[{"type":"histogram2d","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"heatmap":[{"type":"heatmap","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"heatmapgl":[{"type":"heatmapgl","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"contourcarpet":[{"type":"contourcarpet","colorbar":{"outlinewidth":0,"ticks":""}}],"contour":[{"type":"contour","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"surface":[{"type":"surface","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"mesh3d":[{"type":"mesh3d","colorbar":{"outlinewidth":0,"ticks":""}}],"scatter":[{"fillpattern":{"fillmode":"overlay","size":10,"solidity":0.2},"type":"scatter"}],"parcoords":[{"type":"parcoords","line":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterpolargl":[{"type":"scatterpolargl","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"bar":[{"error_x":{"color":"#2a3f5f"},"error_y":{"color":"#2a3f5f"},"marker":{"line":{"color":"#E5ECF6","width":0.5},"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"bar"}],"scattergeo":[{"type":"scattergeo","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterpolar":[{"type":"scatterpolar","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"histogram":[{"marker":{"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"histogram"}],"scattergl":[{"type":"scattergl","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatter3d":[{"type":"scatter3d","line":{"colorbar":{"outlinewidth":0,"ticks":""}},"marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scattermapbox":[{"type":"scattermapbox","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterternary":[{"type":"scatterternary","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scattercarpet":[{"type":"scattercarpet","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"carpet":[{"aaxis":{"endlinecolor":"#2a3f5f","gridcolor":"white","linecolor":"white","minorgridcolor":"white","startlinecolor":"#2a3f5f"},"baxis":{"endlinecolor":"#2a3f5f","gridcolor":"white","linecolor":"white","minorgridcolor":"white","startlinecolor":"#2a3f5f"},"type":"carpet"}],"table":[{"cells":{"fill":{"color":"#EBF0F8"},"line":{"color":"white"}},"header":{"fill":{"color":"#C8D4E3"},"line":{"color":"white"}},"type":"table"}],"barpolar":[{"marker":{"line":{"color":"#E5ECF6","width":0.5},"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"barpolar"}],"pie":[{"automargin":true,"type":"pie"}]},"layout":{"autotypenumbers":"strict","colorway":["#636efa","#EF553B","#00cc96","#ab63fa","#FFA15A","#19d3f3","#FF6692","#B6E880","#FF97FF","#FECB52"],"font":{"color":"#2a3f5f"},"hovermode":"closest","hoverlabel":{"align":"left"},"paper_bgcolor":"white","plot_bgcolor":"#E5ECF6","polar":{"bgcolor":"#E5ECF6","angularaxis":{"gridcolor":"white","linecolor":"white","ticks":""},"radialaxis":{"gridcolor":"white","linecolor":"white","ticks":""}},"ternary":{"bgcolor":"#E5ECF6","aaxis":{"gridcolor":"white","linecolor":"white","ticks":""},"baxis":{"gridcolor":"white","linecolor":"white","ticks":""},"caxis":{"gridcolor":"white","linecolor":"white","ticks":""}},"coloraxis":{"colorbar":{"outlinewidth":0,"ticks":""}},"colorscale":{"sequential":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]],"sequentialminus":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]],"diverging":[[0,"#8e0152"],[0.1,"#c51b7d"],[0.2,"#de77ae"],[0.3,"#f1b6da"],[0.4,"#fde0ef"],[0.5,"#f7f7f7"],[0.6,"#e6f5d0"],[0.7,"#b8e186"],[0.8,"#7fbc41"],[0.9,"#4d9221"],[1,"#276419"]]},"xaxis":{"gridcolor":"white","linecolor":"white","ticks":"","title":{"standoff":15},"zerolinecolor":"white","automargin":true,"zerolinewidth":2},"yaxis":{"gridcolor":"white","linecolor":"white","ticks":"","title":{"standoff":15},"zerolinecolor":"white","automargin":true,"zerolinewidth":2},"scene":{"xaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2},"yaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2},"zaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2}},"shapedefaults":{"line":{"color":"#2a3f5f"}},"annotationdefaults":{"arrowcolor":"#2a3f5f","arrowhead":0,"arrowwidth":1},"geo":{"bgcolor":"white","landcolor":"#E5ECF6","subunitcolor":"white","showland":true,"showlakes":true,"lakecolor":"white"},"title":{"x":0.05},"mapbox":{"style":"light"}}},"legend":{"orientation":"h"},"title":{"text":"2D Route Simulation Across Column Graph (digit 7)"},"xaxis":{"title":{"text":"Column X"}},"yaxis":{"title":{"text":"Column Y"}},"updatemenus":[{"buttons":[{"args":[{"visible":[true,true,true]}],"label":"Show both","method":"update"},{"args":[{"visible":[true,true,false]}],"label":"Vision route","method":"update"},{"args":[{"visible":[true,false,true]}],"label":"Audio route","method":"update"}],"type":"dropdown","x":0.01,"y":1.12}]},                        {"responsive": true}                    )                };                            </script>        </div>
+</body>
+</html>

3D-SEMMN/outputs/figures/route_nodes.csv

@@ -0,0 +1,241 @@
+step,column,activity,modality
+0.0,0.0,0.02956058643758297,vision
+0.0,1.0,0.0409482941031456,vision
+0.0,2.0,0.026564355939626694,vision
+0.0,10.0,0.040411435067653656,vision
+0.0,11.0,0.02947550266981125,vision
+0.0,12.0,0.0309763066470623,vision
+0.0,20.0,0.027739018201828003,vision
+0.0,21.0,0.023138199001550674,vision
+0.0,22.0,0.03138161823153496,vision
+0.0,30.0,0.03375103697180748,vision
+0.0,31.0,0.034890495240688324,vision
+0.0,32.0,0.025857770815491676,vision
+0.0,40.0,0.03491085767745972,vision
+0.0,41.0,0.04038066044449806,vision
+0.0,42.0,0.030116962268948555,vision
+0.0,50.0,0.029867637902498245,vision
+0.0,51.0,0.036975305527448654,vision
+0.0,52.0,0.04041629657149315,vision
+0.0,60.0,0.03777403011918068,vision
+0.0,61.0,0.02775333821773529,vision
+0.0,62.0,0.03777413070201874,vision
+0.0,70.0,0.029612775892019272,vision
+0.0,71.0,0.03889508917927742,vision
+0.0,72.0,0.03445589914917946,vision
+0.0,80.0,0.03937682509422302,vision
+0.0,81.0,0.027864739298820496,vision
+0.0,82.0,0.027492109686136246,vision
+0.0,90.0,0.036627497524023056,vision
+0.0,91.0,0.03462231159210205,vision
+0.0,92.0,0.04038887843489647,vision
+1.0,0.0,0.04356175288558006,vision
+1.0,1.0,0.03761175647377968,vision
+1.0,2.0,0.02439984865486622,vision
+1.0,10.0,0.03711864724755287,vision
+1.0,11.0,0.038757000118494034,vision
+1.0,12.0,0.028452306985855103,vision
+1.0,20.0,0.025478797033429146,vision
+1.0,21.0,0.021252859383821487,vision
+1.0,22.0,0.028824592009186745,vision
+1.0,30.0,0.03100094571709633,vision
+1.0,31.0,0.03204755857586861,vision
+1.0,32.0,0.02375083602964878,vision
+1.0,40.0,0.032066263258457184,vision
+1.0,41.0,0.037090376019477844,vision
+1.0,42.0,0.027662981301546097,vision
+1.0,50.0,0.02743397280573845,vision
+1.0,51.0,0.033962495625019073,vision
+1.0,52.0,0.03712311014533043,vision
+1.0,60.0,0.03469613939523697,vision
+1.0,61.0,0.025491949170827866,vision
+1.0,62.0,0.03469623252749443,vision
+1.0,70.0,0.0271998792886734,vision
+1.0,71.0,0.035725854337215424,vision
+1.0,72.0,0.03164837509393692,vision
+1.0,80.0,0.03616833686828613,vision
+1.0,81.0,0.025594275444746017,vision
+1.0,82.0,0.025252006947994232,vision
+1.0,90.0,0.0870317593216896,vision
+1.0,91.0,0.03180122748017311,vision
+1.0,92.0,0.03709792718291283,vision
+2.0,0.0,0.055453140288591385,vision
+2.0,1.0,0.03477798402309418,vision
+2.0,2.0,0.02256149612367153,vision
+2.0,10.0,0.03432202339172363,vision
+2.0,11.0,0.04663990065455437,vision
+2.0,12.0,0.026308631524443626,vision
+2.0,20.0,0.023559153079986572,vision
+2.0,21.0,0.019651610404253006,vision
+2.0,22.0,0.026652866974473,vision
+2.0,30.0,0.028665248304605484,vision
+2.0,31.0,0.02963300608098507,vision
+2.0,32.0,0.021961383521556854,vision
+2.0,40.0,0.02965030074119568,vision
+2.0,41.0,0.034295883029699326,vision
+2.0,42.0,0.025578774511814117,vision
+2.0,50.0,0.025367021560668945,vision
+2.0,51.0,0.031403668224811554,vision
+2.0,52.0,0.03432615101337433,vision
+2.0,60.0,0.03208203613758087,vision
+2.0,61.0,0.0235713142901659,vision
+2.0,62.0,0.032082121819257736,vision
+2.0,70.0,0.025150563567876816,vision
+2.0,71.0,0.03303416818380356,vision
+2.0,72.0,0.029263898730278015,vision
+2.0,80.0,0.03344331309199333,vision
+2.0,81.0,0.023665931075811386,vision
+2.0,82.0,0.023349450901150703,vision
+2.0,90.0,0.12984082102775574,vision
+2.0,91.0,0.02940523438155651,vision
+2.0,92.0,0.03430286794900894,vision
+3.0,0.0,0.06567821651697159,vision
+3.0,1.0,0.03234130144119263,vision
+3.0,2.0,0.020980749279260635,vision
+3.0,10.0,0.03191728889942169,vision
+3.0,11.0,0.05341819301247597,vision
+3.0,12.0,0.024465344846248627,vision
+3.0,20.0,0.021908504888415337,vision
+3.0,21.0,0.018274743109941483,vision
+3.0,22.0,0.024785460904240608,vision
+3.0,30.0,0.026656849309802055,vision
+3.0,31.0,0.027556801214814186,vision
+3.0,32.0,0.020422684028744698,vision
+3.0,40.0,0.027572883293032646,vision
+3.0,41.0,0.03189297765493393,vision
+3.0,42.0,0.023786624893546104,vision
+3.0,50.0,0.023589707911014557,vision
+3.0,51.0,0.029203403741121292,vision
+3.0,52.0,0.03192112594842911,vision
+3.0,60.0,0.029834242537617683,vision
+3.0,61.0,0.021919816732406616,vision
+3.0,62.0,0.0298343226313591,vision
+3.0,70.0,0.023388417437672615,vision
+3.0,71.0,0.030719663947820663,vision
+3.0,72.0,0.027213554829359055,vision
+3.0,80.0,0.031100144609808922,vision
+3.0,81.0,0.02200780250132084,vision
+3.0,82.0,0.02171349711716175,vision
+3.0,90.0,0.1666511446237564,vision
+3.0,91.0,0.02734498865902424,vision
+3.0,92.0,0.03189947456121445,vision
+0.0,7.0,0.03473106026649475,audio
+0.0,8.0,0.02794949896633625,audio
+0.0,9.0,0.027408087626099586,audio
+0.0,17.0,0.038603682070970535,audio
+0.0,18.0,0.036946672946214676,audio
+0.0,19.0,0.02548990212380886,audio
+0.0,27.0,0.03835974261164665,audio
+0.0,28.0,0.04123900458216667,audio
+0.0,29.0,0.030051397159695625,audio
+0.0,37.0,0.03510997071862221,audio
+0.0,38.0,0.03655286878347397,audio
+0.0,39.0,0.024867547675967216,audio
+0.0,47.0,0.039415426552295685,audio
+0.0,48.0,0.0356949046254158,audio
+0.0,49.0,0.04190101474523544,audio
+0.0,57.0,0.033966340124607086,audio
+0.0,58.0,0.03714350610971451,audio
+0.0,59.0,0.03321348875761032,audio
+0.0,67.0,0.030020426958799362,audio
+0.0,68.0,0.034492406994104385,audio
+0.0,69.0,0.02616014890372753,audio
+0.0,77.0,0.03146708384156227,audio
+0.0,78.0,0.04283866658806801,audio
+0.0,79.0,0.02919699251651764,audio
+0.0,87.0,0.03697360306978226,audio
+0.0,88.0,0.03308973088860512,audio
+0.0,89.0,0.02546405978500843,audio
+0.0,97.0,0.03400813415646553,audio
+0.0,98.0,0.02903173305094242,audio
+0.0,99.0,0.028612902387976646,audio
+1.0,7.0,0.03473105654120445,audio
+1.0,8.0,0.0279494971036911,audio
+1.0,9.0,0.027408085763454437,audio
+1.0,17.0,0.03860367834568024,audio
+1.0,18.0,0.03694666922092438,audio
+1.0,19.0,0.02548990026116371,audio
+1.0,27.0,0.038359738886356354,audio
+1.0,28.0,0.04123900085687637,audio
+1.0,29.0,0.030051395297050476,audio
+1.0,37.0,0.03510996699333191,audio
+1.0,38.0,0.03655286505818367,audio
+1.0,39.0,0.024867547675967216,audio
+1.0,47.0,0.039415426552295685,audio
+1.0,48.0,0.035694900900125504,audio
+1.0,49.0,0.04190101474523544,audio
+1.0,57.0,0.03396633639931679,audio
+1.0,58.0,0.03714350238442421,audio
+1.0,59.0,0.03321348503232002,audio
+1.0,67.0,0.030020423233509064,audio
+1.0,68.0,0.03449240326881409,audio
+1.0,69.0,0.026160147041082382,audio
+1.0,77.0,0.03146708011627197,audio
+1.0,78.0,0.04283866286277771,audio
+1.0,79.0,0.02919699065387249,audio
+1.0,87.0,0.03697359934449196,audio
+1.0,88.0,0.03308972716331482,audio
+1.0,89.0,0.02546405978500843,audio
+1.0,97.0,0.03400813415646553,audio
+1.0,98.0,0.02903173118829727,audio
+1.0,99.0,0.028612900525331497,audio
+2.0,7.0,0.03473105654120445,audio
+2.0,8.0,0.02794949896633625,audio
+2.0,9.0,0.027408087626099586,audio
+2.0,17.0,0.038603682070970535,audio
+2.0,18.0,0.036946672946214676,audio
+2.0,19.0,0.02548990212380886,audio
+2.0,27.0,0.03835974261164665,audio
+2.0,28.0,0.04123900458216667,audio
+2.0,29.0,0.030051397159695625,audio
+2.0,37.0,0.03510997071862221,audio
+2.0,38.0,0.03655286878347397,audio
+2.0,39.0,0.024867551401257515,audio
+2.0,47.0,0.03941543027758598,audio
+2.0,48.0,0.035694900900125504,audio
+2.0,49.0,0.04190101847052574,audio
+2.0,57.0,0.033966340124607086,audio
+2.0,58.0,0.03714350610971451,audio
+2.0,59.0,0.03321348503232002,audio
+2.0,67.0,0.030020425096154213,audio
+2.0,68.0,0.034492406994104385,audio
+2.0,69.0,0.02616014890372753,audio
+2.0,77.0,0.03146708384156227,audio
+2.0,78.0,0.04283866658806801,audio
+2.0,79.0,0.02919699251651764,audio
+2.0,87.0,0.03697360306978226,audio
+2.0,88.0,0.03308973088860512,audio
+2.0,89.0,0.02546406351029873,audio
+2.0,97.0,0.03400813788175583,audio
+2.0,98.0,0.02903173305094242,audio
+2.0,99.0,0.028612902387976646,audio
+3.0,7.0,0.03473105654120445,audio
+3.0,8.0,0.0279495008289814,audio
+3.0,9.0,0.027408087626099586,audio
+3.0,17.0,0.038603682070970535,audio
+3.0,18.0,0.036946672946214676,audio
+3.0,19.0,0.02548990212380886,audio
+3.0,27.0,0.03835974261164665,audio
+3.0,28.0,0.04123900458216667,audio
+3.0,29.0,0.030051397159695625,audio
+3.0,37.0,0.03510997071862221,audio
+3.0,38.0,0.03655286878347397,audio
+3.0,39.0,0.024867551401257515,audio
+3.0,47.0,0.03941543027758598,audio
+3.0,48.0,0.035694900900125504,audio
+3.0,49.0,0.04190101847052574,audio
+3.0,57.0,0.033966340124607086,audio
+3.0,58.0,0.03714350610971451,audio
+3.0,59.0,0.03321348503232002,audio
+3.0,67.0,0.030020426958799362,audio
+3.0,68.0,0.034492406994104385,audio
+3.0,69.0,0.02616015076637268,audio
+3.0,77.0,0.03146708384156227,audio
+3.0,78.0,0.04283866658806801,audio
+3.0,79.0,0.02919699251651764,audio
+3.0,87.0,0.03697360306978226,audio
+3.0,88.0,0.03308973088860512,audio
+3.0,89.0,0.02546406351029873,audio
+3.0,97.0,0.03400813788175583,audio
+3.0,98.0,0.02903173491358757,audio
+3.0,99.0,0.028612902387976646,audio

3D-SEMMN/outputs/history.json

@@ -0,0 +1,37 @@
+[
+  {
+    "epoch": 0,
+    "step": 0,
+    "loss_total": 6.665867805480957,
+    "loss_cls": 2.322234630584717,
+    "loss_contrastive": 4.330072402954102,
+    "loss_manifold": 0.007405777927488089,
+    "active_params": 2.0,
+    "energy_proxy": 0.17007668316364288
+  },
+  {
+    "epoch": 0,
+    "step": 10,
+    "loss_total": 5.909451484680176,
+    "loss_cls": 2.2938742637634277,
+    "loss_contrastive": 3.5930726528167725,
+    "loss_manifold": 0.001445187022909522,
+    "active_params": 6.0,
+    "energy_proxy": 0.12415368109941483
+  },
+  {
+    "epoch": 0,
+    "step": 20,
+    "loss_total": 5.335961818695068,
+    "loss_cls": 2.2907979488372803,
+    "loss_contrastive": 3.0338499546051025,
+    "loss_manifold": 0.0009116176515817642,
+    "active_params": 6.0,
+    "energy_proxy": 0.07187842577695847
+  },
+  {
+    "epoch": 0,
+    "val_accuracy": 0.0,
+    "val_cross_modal_retrieval": 0.0
+  }
+]

3D-SEMMN/requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.3
+torchvision>=0.18
+numpy>=1.25
+pandas>=2.2
+PyYAML>=6.0
+matplotlib>=3.8
+plotly>=5.22
+scikit-learn>=1.4
+tqdm>=4.66
+datasets>=2.20

3D-SEMMN/semmn/__init__.py

@@ -0,0 +1,6 @@
+"""3D-SEMMN package."""
+
+from semmn.config import SEMMNConfig
+from semmn.model import SEMMNModel
+
+__all__ = ["SEMMNConfig", "SEMMNModel"]

3D-SEMMN/semmn/columns.py

@@ -0,0 +1,33 @@
+"""Columnar dynamics modules."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ColumnarDynamics(nn.Module):
+    """Dense intra-column dynamics with low-rank factors."""
+
+    def __init__(self, num_columns: int, state_dim: int, rank: int = 8) -> None:
+        super().__init__()
+        self.num_columns = num_columns
+        self.state_dim = state_dim
+        self.rank = rank
+
+        self.u = nn.Parameter(torch.randn(num_columns, state_dim, rank) * 0.05)
+        self.v = nn.Parameter(torch.randn(num_columns, rank, state_dim) * 0.05)
+        self.bias = nn.Parameter(torch.zeros(num_columns, state_dim))
+        self.norm = nn.LayerNorm(state_dim)
+
+    def forward(self, states: torch.Tensor, inter_messages: torch.Tensor) -> torch.Tensor:
+        # W_col = U @ V in factored form to keep parameters compact.
+        w = torch.matmul(self.u, self.v)  # (C, D, D)
+        intra = torch.einsum("bcd,cde->bce", states, w)
+        updated = intra + inter_messages + self.bias.unsqueeze(0)
+        return torch.tanh(self.norm(updated))
+
+
+def pool_columns(states: torch.Tensor) -> torch.Tensor:
+    """Convert `(B, C, D)` into `(B, C)` pooled activity."""
+    return states.mean(dim=-1)

3D-SEMMN/semmn/config.py

@@ -0,0 +1,121 @@
+"""Configuration handling for 3D-SEMMN."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass
+class GridConfig:
+    dims: tuple[int, int, int] = (100, 100, 50)
+    n_columns_xy: tuple[int, int] = (10, 10)
+    vision_y_ratio: float = 0.30
+    audio_y_ratio: float = 0.70
+    hub_y_range: tuple[float, float] = (0.40, 0.60)
+    hub_z_range: tuple[float, float] = (0.30, 0.70)
+
+
+@dataclass
+class ModelConfig:
+    state_dim: int = 32
+    latent_dim: int = 8
+    vision_feature_dim: int = 256
+    audio_feature_dim: int = 256
+    recurrent_steps: int = 6
+    attention_heads: int = 4
+
+
+@dataclass
+class ConnectivityConfig:
+    lambda_decay: float = 5.0
+    inter_density: float = 0.08
+    intra_density_boost: float = 2.0
+    min_weight: float = 0.0
+    max_weight: float = 1.5
+    init_weight_scale: float = 0.15
+    stdp_interval: int = 100
+    stdp_lr: float = 0.01
+    anti_hebbian: float = 0.05
+
+
+@dataclass
+class LossConfig:
+    contrastive_weight: float = 1.0
+    reconstruction_weight: float = 0.5
+    imagination_weight: float = 0.3
+    manifold_weight: float = 0.1
+    spatial_penalty_weight: float = 5e-4
+    temperature: float = 0.1
+    use_spatial_penalty: bool = True
+    use_manifold_loss: bool = True
+
+
+@dataclass
+class TrainingConfig:
+    dataset_source: str = "synthetic"
+    batch_size: int = 64
+    epochs: int = 10
+    lr: float = 1e-3
+    weight_decay: float = 1e-4
+    num_workers: int = 4
+    log_interval: int = 50
+    val_batches: int = 100
+    amp: bool = True
+
+
+@dataclass
+class PathsConfig:
+    output_dir: str = "outputs"
+
+
+@dataclass
+class SEMMNConfig:
+    seed: int = 42
+    grid: GridConfig = field(default_factory=GridConfig)
+    model: ModelConfig = field(default_factory=ModelConfig)
+    connectivity: ConnectivityConfig = field(default_factory=ConnectivityConfig)
+    loss: LossConfig = field(default_factory=LossConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    paths: PathsConfig = field(default_factory=PathsConfig)
+
+    @property
+    def n_columns(self) -> int:
+        return self.grid.n_columns_xy[0] * self.grid.n_columns_xy[1]
+
+    @property
+    def n_neurons(self) -> int:
+        return self.grid.dims[0] * self.grid.dims[1] * self.grid.dims[2]
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> "SEMMNConfig":
+        with Path(path).open("r", encoding="utf-8") as f:
+            raw = yaml.safe_load(f) or {}
+        return cls.from_dict(raw)
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "SEMMNConfig":
+        return cls(
+            seed=raw.get("seed", 42),
+            grid=_merge_dataclass(GridConfig(), raw.get("grid", {})),
+            model=_merge_dataclass(ModelConfig(), raw.get("model", {})),
+            connectivity=_merge_dataclass(
+                ConnectivityConfig(), raw.get("connectivity", {})
+            ),
+            loss=_merge_dataclass(LossConfig(), raw.get("loss", {})),
+            training=_merge_dataclass(TrainingConfig(), raw.get("training", {})),
+            paths=_merge_dataclass(PathsConfig(), raw.get("paths", {})),
+        )
+
+
+def _merge_dataclass(default_obj: Any, overrides: dict[str, Any]) -> Any:
+    current = default_obj.__dict__.copy()
+    for key, value in overrides.items():
+        if isinstance(current.get(key), tuple) and isinstance(value, list):
+            current[key] = tuple(value)
+        else:
+            current[key] = value
+    return type(default_obj)(**current)

3D-SEMMN/semmn/connectivity.py

@@ -0,0 +1,98 @@
+"""Sparse distance-penalized connectivity and STDP rules."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from semmn.config import ConnectivityConfig
+
+
+class SparseInterColumnConnectivity(nn.Module):
+    """Sparse inter-column graph with distance-decay sampling.
+
+    Inspired by seRNN wiring constraints: long-range links carry a larger cost,
+    so we initialize and regularize them with an exponential distance penalty.
+    """
+
+    def __init__(
+        self,
+        column_centers: torch.Tensor,
+        state_dim: int,
+        config: ConnectivityConfig,
+    ) -> None:
+        super().__init__()
+        self.state_dim = state_dim
+        self.config = config
+        self.num_columns = int(column_centers.shape[0])
+        self.register_buffer("column_centers", column_centers.clone())
+
+        row_idx, col_idx, distances, weights = self._build_sparse_edges()
+        self.register_buffer("row_idx", row_idx)
+        self.register_buffer("col_idx", col_idx)
+        self.register_buffer("edge_distances", distances)
+        self.edge_weights = nn.Parameter(weights)
+        self.inter_gain = nn.Parameter(torch.ones(state_dim))
+
+    def _build_sparse_edges(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        centers = self.column_centers
+        c = centers.shape[0]
+        dist = torch.cdist(centers, centers, p=2)
+        decay = torch.exp(-dist / max(self.config.lambda_decay, 1e-6))
+        eye_mask = ~torch.eye(c, dtype=torch.bool, device=centers.device)
+        prob = (decay * self.config.inter_density).clamp(0.0, 1.0)
+        sampled = (torch.rand_like(prob) < prob) & eye_mask
+        row_idx, col_idx = torch.where(sampled)
+
+        if row_idx.numel() == 0:
+            row_idx = torch.arange(0, c - 1, device=centers.device)
+            col_idx = torch.arange(1, c, device=centers.device)
+
+        distances = dist[row_idx, col_idx]
+        weights = torch.randn(row_idx.numel(), device=centers.device) * self.config.init_weight_scale
+        weights = weights.clamp(self.config.min_weight, self.config.max_weight)
+        return row_idx.long(), col_idx.long(), distances.float(), weights.float()
+
+    def sparse_matrix(self) -> torch.Tensor:
+        indices = torch.stack([self.row_idx, self.col_idx], dim=0)
+        return torch.sparse_coo_tensor(
+            indices=indices,
+            values=self.edge_weights,
+            size=(self.num_columns, self.num_columns),
+            device=self.edge_weights.device,
+        ).coalesce()
+
+    def forward(self, states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply sparse inter-column message passing.
+
+        Args:
+            states: `(B, C, D)` column states.
+        Returns:
+            inter_messages: `(B, C, D)` sparse messages.
+            scalar_activity: `(B, C)` pooled activity used for STDP.
+        """
+        scalar_activity = states.mean(dim=-1)  # (B, C)
+        sparse_w = self.sparse_matrix()
+        propagated = torch.sparse.mm(sparse_w, scalar_activity.transpose(0, 1)).transpose(0, 1)
+        inter_messages = propagated.unsqueeze(-1) * self.inter_gain.view(1, 1, -1)
+        return inter_messages, scalar_activity
+
+    @torch.no_grad()
+    def stdp_update(self, pre_activity: torch.Tensor, post_activity: torch.Tensor) -> None:
+        """Hebbian/STDP-like update on sparse edge values."""
+        if pre_activity.numel() == 0:
+            return
+        pre = pre_activity.mean(dim=0)  # (C,)
+        post = post_activity.mean(dim=0)  # (C,)
+        corr = torch.outer(post, pre)
+        edge_corr = corr[self.row_idx, self.col_idx]
+        anti = self.config.anti_hebbian * self.edge_weights.abs()
+        delta = self.config.stdp_lr * (edge_corr - anti)
+        new_w = self.edge_weights + delta
+        self.edge_weights.copy_(new_w.clamp(self.config.min_weight, self.config.max_weight))
+
+    def spatial_wiring_cost(self) -> torch.Tensor:
+        return (self.edge_weights.abs() * self.edge_distances).mean()
+
+    def active_parameters(self) -> int:
+        return int((self.edge_weights.abs() > 1e-8).sum().item())

3D-SEMMN/semmn/dataset.py

@@ -0,0 +1,135 @@
+"""Multimodal AV-MNIST dataset helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torchvision import datasets, transforms
+
+from semmn.config import SEMMNConfig
+
+
+@dataclass
+class Batch:
+    vision: torch.Tensor
+    audio: torch.Tensor
+    label: torch.Tensor
+
+
+class SyntheticAVMNIST(Dataset):
+    """Pair MNIST vision with synthetic spectrogram-like audio."""
+
+    def __init__(self, root: str, train: bool = True) -> None:
+        tfm = transforms.ToTensor()
+        self.mnist = datasets.MNIST(root=root, train=train, download=True, transform=tfm)
+
+    def __len__(self) -> int:
+        return len(self.mnist)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        vision, label = self.mnist[idx]
+        audio = self._label_to_spectrogram(int(label), idx)
+        return vision, audio, torch.tensor(label, dtype=torch.long)
+
+    def _label_to_spectrogram(self, label: int, seed: int) -> torch.Tensor:
+        g = torch.Generator().manual_seed(seed)
+        t = torch.linspace(0, 1, steps=224)
+        base_freq = 110.0 + 22.0 * float(label)
+        sine = (
+            torch.sin(2.0 * torch.pi * base_freq * t)
+            + 0.3 * torch.sin(2.0 * torch.pi * 2.0 * base_freq * t)
+            + 0.2 * torch.randn(224, generator=g)
+        )
+        spec = torch.stft(
+            sine,
+            n_fft=64,
+            hop_length=16,
+            win_length=64,
+            return_complex=True,
+        ).abs()
+        spec = spec.unsqueeze(0)  # (1, F, T)
+        spec = F.interpolate(spec.unsqueeze(0), size=(112, 112), mode="bilinear", align_corners=False)
+        spec = spec.squeeze(0)
+        spec = spec / (spec.max().clamp_min(1e-6))
+        return spec
+
+
+class HuggingFaceAVMNIST(Dataset):
+    """Attempt to read BLOSSOM-framework/AV-MNIST with robust key fallback."""
+
+    def __init__(self, split: str = "train") -> None:
+        from datasets import load_dataset  # lazy import
+
+        self.ds = load_dataset("BLOSSOM-framework/AV-MNIST", split=split)
+
+    def __len__(self) -> int:
+        return len(self.ds)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        row = self.ds[idx]
+        vision = _to_tensor_image(row.get("image") or row.get("vision") or row.get("mnist"))
+        audio = _to_tensor_audio(row.get("audio") or row.get("spectrogram") or row.get("sound"))
+        label = int(row.get("label") or row.get("digit") or row.get("class"))
+        return vision, audio, torch.tensor(label, dtype=torch.long)
+
+
+def _to_tensor_image(value) -> torch.Tensor:
+    if isinstance(value, torch.Tensor):
+        x = value.float()
+    else:
+        x = torch.tensor(value, dtype=torch.float32)
+    if x.ndim == 2:
+        x = x.unsqueeze(0)
+    x = F.interpolate(x.unsqueeze(0), size=(28, 28), mode="bilinear", align_corners=False).squeeze(0)
+    return x.clamp(0.0, 1.0)
+
+
+def _to_tensor_audio(value) -> torch.Tensor:
+    if isinstance(value, dict) and "array" in value:
+        arr = torch.tensor(value["array"], dtype=torch.float32)
+        spec = torch.stft(arr, n_fft=128, hop_length=32, return_complex=True).abs()
+        x = spec.unsqueeze(0)
+    elif isinstance(value, torch.Tensor):
+        x = value.float()
+    else:
+        x = torch.tensor(value, dtype=torch.float32)
+    while x.ndim < 3:
+        x = x.unsqueeze(0)
+    if x.shape[0] != 1:
+        x = x[:1]
+    x = F.interpolate(x.unsqueeze(0), size=(112, 112), mode="bilinear", align_corners=False).squeeze(0)
+    x = x / (x.max().clamp_min(1e-6))
+    return x
+
+
+def build_dataloaders(config: SEMMNConfig) -> tuple[DataLoader, DataLoader]:
+    root = "./data"
+    if config.training.dataset_source == "hf":
+        try:
+            train_ds = HuggingFaceAVMNIST(split="train")
+            val_ds = HuggingFaceAVMNIST(split="test")
+        except Exception:
+            train_ds = SyntheticAVMNIST(root=root, train=True)
+            val_ds = SyntheticAVMNIST(root=root, train=False)
+    else:
+        train_ds = SyntheticAVMNIST(root=root, train=True)
+        val_ds = SyntheticAVMNIST(root=root, train=False)
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=config.training.batch_size,
+        shuffle=True,
+        num_workers=config.training.num_workers,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=config.training.batch_size,
+        shuffle=False,
+        num_workers=max(1, config.training.num_workers // 2),
+        pin_memory=True,
+    )
+    return train_loader, val_loader

3D-SEMMN/semmn/encoders.py

@@ -0,0 +1,66 @@
+"""Vision/audio encoders and lobe injection."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, out_dim: int = 256) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(1, 32, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Flatten(),
+            nn.Linear(64 * 7 * 7, out_dim),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class AudioEncoder1D(nn.Module):
+    """1D Conv encoder over spectrogram time sequences."""
+
+    def __init__(self, out_dim: int = 256, freq_bins: int = 112) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(freq_bins, 128, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(128, 128, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool1d(1),
+            nn.Flatten(),
+            nn.Linear(128, out_dim),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        # Input expected `(B, 1, F, T)`; use `(B, F, T)` for Conv1d.
+        x = spec.squeeze(1)
+        return self.net(x)
+
+
+class LobeInjector(nn.Module):
+    """Project modality features into selected columns."""
+
+    def __init__(self, in_dim: int, num_columns: int, state_dim: int, target_columns: torch.Tensor) -> None:
+        super().__init__()
+        self.num_columns = num_columns
+        self.state_dim = state_dim
+        self.register_buffer("target_columns", target_columns.long())
+        out_dim = max(1, self.target_columns.numel()) * state_dim
+        self.project = nn.Linear(in_dim, out_dim)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        b = features.shape[0]
+        injected = features.new_zeros(b, self.num_columns, self.state_dim)
+        projected = self.project(features).view(b, -1, self.state_dim)
+        injected[:, self.target_columns, :] = projected[:, : self.target_columns.numel(), :]
+        return injected