Model Summary
(directly go to the demo page here)
Background motivation
The development of the latest neural network/deep learning model architectures (State-of-the-Art/SOTA) over the past two decades has seen a trend toward increasingly large-scale model development (in terms of the number of parameters, reaching billions, especially in LLM models), heavy computational burdens/requirement of capable hardware and infrastructure, large energy footprints, and large model sizes, as well as high and unsustainable memory and execution footprints. This may not be immediately apparent to end users, but it actually poses a real environmental threat.
With the current trend of climate change, it is unwise to continue the paradigm and practice of developing and using neural network products that leave a high environmental footprint (energy footprint, memory footprint, inefficient use of resources during execution/inference).
I propose a general neural network model architecture that is more sustainable (GreenAI) and has the potential to be applied to various NN tasks, namely the Network-of-Dendrites (NoD) model architecture.
This model architecture is not specific to a specific domain, but rather represents a core architectural paradigm that is fundamentally different from the model architectures currently being developed. The architecture is fundamentally a computation and routing paradigm (a lean, modular "thinking core") rather than a standalone end-to-end sensor pipeline.
The core paradigm of the NoD model architecture is minimizing the number of parameters by applying the principle of shared parameters, rather than increasing the number to billions.
NoD was developed by emulating the basic idea of how the brain's neurons learn and function biologically. Brain components, modeled as dendrites, go through two main phases:
- the initial phase when the nervous system begins to learn and form new neural networks that become stronger over time after successfully learning from data (myelination); and
- the phase when the nervous system is fully formed and the training and dynamic processing can be stopped and "locked in."
In the context of the NN model architecture, this is achieved by dividing the model architecture into two phases: a. first phase: statistically searching/matching which data can be processed by which archetype/subMLP (gating/routing); after that, the data is fed into a shared-base of parameters before being fed to the selected archetype. This occurs in the first phase during training, where the most appropriate probability for dividing the data into each archetype is still being sought; c. Phase two: the locked-in phase, where the transition from soft probabilistic to deterministic occurs after the optimal routing probability to K other archetypes for processing the data is obtained from Phase one.
Model Provenance & Lineage
- Base Architecture: None (Built completely from scratch).
- Pre-trained Weights: None (Trained with randomly initialized weights).
- Training Dataset: Standard MNIST Handwritten Digit Dataset.
Origin Statement
This model is a custom-designed architecture developed independently. It is not a fork, fine-tuned variant, or distillation of any existing pre-trained model. All weights were trained from scratch exclusively on the attached MNIST dataset.
Usage
Create NoDClassificationLayer class
import tensorflow as tf
from tensorflow.keras import layers, Model, initializers, datasets
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import plot_model
import numpy as np
import export_nod
import monitor
# Reuse the NetworkOfDendritesLayer implementation with Multi-Class Output adjustment
@tf.keras.utils.register_keras_serializable(package='Custom', name='NoDClassificationLayer')
class NoDClassificationLayer(layers.Layer):
def __init__(self, num_inputs, num_classes, num_archetypes, archetype_configs, embedding_dim=16, **kwargs):
super(NoDClassificationLayer, self).__init__(**kwargs)
self.num_inputs = num_inputs
self.num_classes = num_classes
self.num_archetypes = num_archetypes
self.archetype_configs = archetype_configs
self.embedding_dim = embedding_dim
self.is_locked = False
self.temperature = 1.0
self.hard_assignments = None
def build(self, input_shape):
# Discovery parameters
self.input_keys = self.add_weight(
name="input_keys", shape=(self.num_inputs, self.embedding_dim),
initializer=initializers.RandomNormal(stddev=0.1), trainable=True
)
self.archetype_prototypes = self.add_weight(
name="archetype_prototypes", shape=(self.num_archetypes, self.embedding_dim),
initializer=initializers.RandomNormal(stddev=0.1), trainable=True
)
# Bounded Bank of Archetype NPUs
self.archetype_mlps = []
for cfg in self.archetype_configs:
# Output of each mini-MLP now projects to `num_classes` so each pixel
# contributes a non-linear vote to every class score.
# input_shape(1,) is used to ensure the MLPs can process single pixel inputs and stick to 472 params/enforce True 472-params scale in Keras for this NoD arch.
mlp = tf.keras.Sequential([
layers.Dense(cfg["hidden_dim"], activation=cfg["activation"],
kernel_initializer=initializers.VarianceScaling(scale=cfg["init_scale"], mode='fan_in', distribution='normal'), input_shape=(1,)),
layers.Dense(self.num_classes, kernel_initializer=initializers.RandomNormal(stddev=0.1))
])
self.archetype_mlps.append(mlp)
super(NoDClassificationLayer, self).build(input_shape)
def lock_routing(self):
similarity = tf.matmul(self.input_keys, self.archetype_prototypes, transpose_b=True)
self.hard_assignments = tf.argmax(similarity, axis=-1).numpy()
self.is_locked = True
self.input_keys.trainable = False
self.archetype_prototypes.trainable = False
print(f"\n[SYSTEM] Routing Locked. Archetype distribution: {np.bincount(self.hard_assignments)}")
def call(self, inputs):
batch_size = tf.shape(inputs)[0]
if not self.is_locked:
# Phase 1: Soft Dynamic Routing
similarity = tf.matmul(self.input_keys, self.archetype_prototypes, transpose_b=True)
soft_routing = tf.nn.softmax(similarity / self.temperature, axis=-1)
expanded_inputs = tf.expand_dims(inputs, axis=-1) # (Batch, 784, 1)
all_npu_outputs = []
for k in range(self.num_archetypes):
flat_in = tf.reshape(expanded_inputs, [-1, 1])
flat_out = self.archetype_mlps[k](flat_in) # (Batch*784, 10)
npu_out = tf.reshape(flat_out, [batch_size, self.num_inputs, self.num_classes])
all_npu_outputs.append(npu_out)
stacked_outputs = tf.stack(all_npu_outputs, axis=-1) # (Batch, 784, 10, Num_Arch)
routing_expanded = tf.reshape(soft_routing, [1, self.num_inputs, 1, self.num_archetypes])
dendritic_outputs = tf.reduce_sum(stacked_outputs * routing_expanded, axis=-1) # (Batch, 784, 10)
else:
# Phase 2: Fully Vectorized Structural Execution Loop
# self.hard_assignments contains the winning archetype index for each of the 784 inputs (Shape: [784])
# We map each input to its assigned archetype without breaking tensor flow.
expanded_inputs = tf.expand_dims(inputs, axis=-1) # (Batch, 784, 1)
flat_in = tf.reshape(expanded_inputs, [-1, 1]) # (Batch * 784, 1)
# Compute outputs for all archetypes across all inputs first
all_npu_outputs = []
for k in range(self.num_archetypes):
flat_out = self.archetype_mlps[k](flat_in) # (Batch * 784, 10)
npu_out = tf.reshape(flat_out, [batch_size, self.num_inputs, self.num_classes])
all_npu_outputs.append(npu_out)
stacked_outputs = tf.stack(all_npu_outputs, axis=-1) # (Batch, 784, 10, Num_Arch)
# Create a hard one-hot mask from self.hard_assignments: shape (784, Num_Arch)
hard_mask = tf.one_hot(self.hard_assignments, depth=self.num_archetypes, dtype=tf.float32)
# Reshape mask to broadcast correctly: (1, 784, 1, Num_Arch)
routing_expanded = tf.reshape(hard_mask, [1, self.num_inputs, 1, self.num_archetypes])
# Select only the winning archetype's output for each input index deterministically
dendritic_outputs = tf.reduce_sum(stacked_outputs * routing_expanded, axis=-1) # (Batch, 784, 10)
# RMS Normalization over dendritic outputs
rms = tf.math.sqrt(tf.reduce_mean(tf.math.square(dendritic_outputs), axis=1, keepdims=True) + 1e-8)
normalized_outputs = dendritic_outputs / rms
# Macro Neuron Aggregation Sum
macro_sum = tf.reduce_sum(normalized_outputs, axis=1) # (Batch, 10)
return macro_sum
def get_config(self):
"""Serialization support for saving/loading the layer."""
config = super().get_config()
config.update({
"num_inputs": self.num_inputs,
"num_classes": self.num_classes,
"num_archetypes": self.num_archetypes,
"archetype_configs": self.archetype_configs,
"embedding_dim": self.embedding_dim
})
return config
@classmethod
def from_config(cls, config):
"""Deserialization support for loading the layer from a config."""
# config_copy = config.copy()
return cls(**config)
Create the load balaced layer class
import tensorflow as tf
from tensorflow.keras import layers
class BalancedNoDClassificationLayer(layers.Layer):
"""
1D Network-of-Dendrites (NoD) Classification Layer with
Shared-Parameter Core and Load-Balancing Auxiliary Loss.
"""
def __init__(self, num_archetypes=8, archetype_dim=16, balance_weight=0.01, **kwargs):
super(BalancedNoDClassificationLayer, self).__init__(**kwargs)
self.num_archetypes = num_archetypes
self.archetype_dim = archetype_dim
self.balance_weight = balance_weight
def build(self, input_shape):
# 1D Shared-Parameter Core: Shape (num_archetypes, archetype_dim)
self.archetype_core = self.add_weight(
shape=(self.num_archetypes, self.archetype_dim),
initializer='variance_scaling',
trainable=True,
name='nod_1d_shared_core'
)
# Router network to compute soft routing gates across archetypes
self.router_weights = self.add_weight(
shape=(input_shape[-1], self.num_archetypes),
initializer='glorot_uniform',
trainable=True,
name='nod_router_weights'
)
super(BalancedNoDClassificationLayer, self).build(input_shape)
def call(self, inputs):
batch_size = tf.shape(inputs)[0]
# 1. Compute routing probabilities via Softmax
router_logits = tf.matmul(inputs, self.router_weights)
routing_gates = tf.nn.softmax(router_logits, axis=-1) # Shape: (Batch, num_archetypes)
# 2. Load-Balancing Auxiliary Loss (Prevents routing collapse)
mean_gate_per_archetype = tf.reduce_mean(routing_gates, axis=0)
uniform_target = 1.0 / float(self.num_archetypes)
load_balance_loss = self.balance_weight * tf.reduce_sum(
tf.square(mean_gate_per_archetype - uniform_target)
)
self.add_loss(load_balance_loss)
# 3. Weighted combination of the 1D shared archetype core parameters
# routing_gates: (Batch, num_archetypes, 1) x archetype_core: (1, num_archetypes, archetype_dim)
gates_expanded = tf.expand_dims(routing_gates, axis=-1)
core_expanded = tf.expand_dims(self.archetype_core, axis=0)
modulated_archetypes = gates_expanded * core_expanded
output_features = tf.reduce_sum(modulated_archetypes, axis=1) # Shape: (Batch, archetype_dim)
return output_features
def compute_output_shape(self, input_shape):
return (input_shape[0], self.archetype_dim)
def get_config(self):
config = super(BalancedNoDClassificationLayer, self).get_config()
config.update({
"num_archetypes": self.num_archetypes,
"archetype_dim": self.archetype_dim,
"balance_weight": self.balance_weight,
})
return config
@classmethod
def from_config(cls, config):
return cls(**config)
Load and evaluate saved model
def evaluate_saved_model(model_path):
print(f"[INFO] Loading model from '{model_path}'...")
# 1. Load the model safely with compile=False
loaded_model = load_model(
model_path,
custom_objects={"NoDClassificationLayer": NoDClassificationLayer},
compile=False
)
print("[SUCCESS] Model loaded.")
# 2. Load the official MNIST test dataset
print("[INFO] Loading MNIST test dataset...")
(_, _), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
# 3. Apply the exact same preprocessing used during training (Flatten & Normalize)
x_test_processed = x_test.astype(np.float32) / 255.0
x_test_processed = x_test_processed.reshape(-1, 784) # Flatten to (10000, 784)
print(f"[INFO] Running inference on {len(x_test_processed)} test samples...")
# 4. Perform batch prediction
# (If memory is tight, you can batch this, but for 784-dim tiny models, direct prediction is fine)
predictions = loaded_model.predict(x_test_processed, batch_size=1024, verbose=1)
# 5. Extract predicted classes
predicted_labels = np.argmax(predictions, axis=1)
# 6. Calculate True Accuracy
correct_predictions = np.sum(predicted_labels == y_test)
total_samples = len(y_test)
test_accuracy = (correct_predictions / total_samples) * 100.0
print("\n" + "="*40)
print(f" FINAL EVALUATION REPORT")
print("="*40)
print(f" Total Test Samples : {total_samples}")
print(f" Correct Predictions: {correct_predictions}")
print(f" True Test Accuracy : {test_accuracy:.2f}%")
print("="*40)
return test_accuracy
# Test your custom images right away
# batch_test("mnist_test_img")
if __name__ == "__main__":
# Point to your saved model file
model_file = "mnist_nod_model.keras"
evaluate_saved_model(model_file)
Implementation requirements
The model was trained in a standard consumer laptop with no CUDA capable GPU.
Model Characteristics
Illustration
Data flow
Model architecture
Inside the archetype
Model initialization
The model was trained from scratch on MNIST dataset.
Model stats
NoD 1 (1D base shared parameters)
Trainable Parameters: 13.080 File size: 218 Kbytes Architecture: NoD (dendritic non-linear routing and archetype sharing) Nr.of training epoch: 10 epoch True accuracy: 88.18%
Trainable Parameters: 13.080 File size: 218 Kbytes Architecture: NoD (dendritic non-linear routing and archetype sharing) Nr.of training epoch: 20 epoch True accuracy: 90.04%
NoD 2 (2D base shared parameters)
Trainable Parameters: 116,714 parameters File size: 455.91 KB (the actual model .keras file size uploaded is 1.4 MB) Architecture: NoD (dendritic non-linear routing and archetype sharing) Nr.of training epoch: 10 epoch True accuracy: 98.83%
Other details
The model is not pruned nor quantized.
Data Overview
The model trained and evaluated using MNIST dataset and standard method to split train/val/test data.
Evaluation Results
True accuracy on MNIST test set: 88.18% (10 epoch) True accuracy on MNIST test set: 90.04% (20 epoch)
- Downloads last month
- 108
Dataset used to train febrifahmi/NoD
Space using febrifahmi/NoD 1
Evaluation results
- Accuracy on MNIST Test Settest set self-reported0.988