Genesis: Un'Architettura di Reti Neurali Ispirata alle Connessioni del Cervello Umano

July 5, 2024

In questo articolo, vi guiderò attraverso uno dei miei ultimi progetti, Genesis, un'architettura di reti neurali progettata per imitare le connessioni neurali umane. Questo progetto è stato sia una sfida che un'esperienza gratificante, poiché mi ha richiesto di approfondire le complessità delle strutture neurali e delle loro funzioni. L'obiettivo era creare un modello che potesse offrire calcoli più efficienti e plausibili biologicamente.

GenesisConfig: Lo Scheletro della Configurazione

Uno dei primi passi nella creazione di Genesis è stata la definizione di una classe di configurazione robusta. La classe GenesisConfig consente configurazioni flessibili e dettagliate. Include parametri per le dimensioni della rete, il numero di sezioni, neuroni, neuroni attentivi, dendritici, sinapsi e altro. Questa configurabilità è cruciale poiché mi permette di sperimentare con diverse strutture di rete e il loro impatto sulle prestazioni.

@dataclass
class GenesisConfig:
    n_dim: int
    n_sections: int
    n_neurons: int
    n_attentive_neurons: int
    n_attentive_neuron_heads: int
    n_dendritics: int
    n_synapses: int
    n_pos_size: int
    n_neighbors: int
    n_seed: int
    n_input: str
    n_output: str
    n_vocab_size: int
    n_cross_attention: bool
    n_outputs: int

La Classe Genesis: Combinare Sezioni e Neuroni

La classe Genesis è dove avviene la magia. Combina varie sezioni, ciascuna contenente neuroni che possono interagire e propagare segnali. L'architettura supporta più tipi di input e output, rendendola versatile per diversi compiti. Ad esempio, è possibile configurarla per il modeling causale del linguaggio o per la classificazione di sequenze.

Inizializzazione

Durante l'inizializzazione, la rete configura il suo input, le sezioni, le divisioni per gli input e output delle sezioni e lo strato di output finale basato sulla configurazione fornita.

class Genesis(nn.Module):
    def __init__(self, config: GenesisConfig, tokenizer: AutoTokenizer = None):
        super(Genesis, self).__init__()
        self.config = config
        self.tokenizer = tokenizer
        self.n_dim = self.config.n_dim
        self.n_inner_dim = self.n_dim // self.config.n_sections if self.config.n_section == "split" else self.n_dim
        self.input = self._init_input()
        self.sections = self._init_sections()
        self.a_splits, self.c_splits, self.combiner = self._init_splits()
        self.output = self._init_output()

Passaggio Forward

Il passaggio forward gestisce il flusso dei dati attraverso la rete, gestendo l'attivazione dei neuroni e l'interazione tra le sezioni. Se configurato per l'attenzione incrociata, elabora gli input secondari di conseguenza.

def forward(self, inputs, attention_mask=None, y=None, labels=None):
    total_neuron_activation_count = 0
    hidden_states = self._process_input(inputs)
    hidden_states = self._process_sections(hidden_states, attention_mask, y, total_neuron_activation_count)
    outputs, loss = self._process_output(hidden_states, labels)
    return GenesisOutput(loss=loss, outputs=outputs, neuron_activation_count=total_neuron_activation_count)

Sezioni e Neuroni: I Componenti Principali

Classe Section

Ogni Sezione contiene un insieme di neuroni. I neuroni possono essere attentivi o non attentivi, con strutture interne e funzioni diverse. La classe section include anche metodi per i passaggi forward e la gestione dell'attivazione dei neuroni.

class Section(nn.Module):
    def __init__(self, n_dim, n_neurons, n_attentive_neurons, n_attentive_neuron_heads, n_dendritics, n_synapses, n_pos_size, n_neighbors):
        super(Section, self).__init__()
        self.neurons = self._init_neurons()
        self.indexer = Indexer(self.n_dim, self.neurons, self.n_attentive_neurons)

    def forward(self, x, attention_mask=None, y=None):
        neuron = self.indexer(x)
        x = neuron.forward(x, lambda: self.increment_neuron_count(), attention_mask, y)
        self.reset_neurons()
        return x, self.neuron_activation_count

Classe Neuron

La classe Neuron è il cuore di Genesis. Gestisce l'elaborazione dendritica, l'interazione somatica, la propagazione del segnale assonico e la trasmissione sinaptica. I neuroni possono anche stabilire connessioni con i loro vicini, simulando la natura interconnessa dei neuroni biologici.

class Neuron(nn.Module):
    def __init__(self, n_dim, attentive_neuron, n_attentive_neuron_heads, n_dendritics, n_synapses, n_pos_size, n_neighbors):
        super(Neuron, self).__init__()
        self.dendritics = nn.ModuleList([Dendritic(self.n_dim) for _ in range(self.n_dendritics)])
        self.soma = Soma(self.n_dim)
        self.axon = Axon(self.n_dim, self.attentive_neuron, self.n_attentive_neuron_heads)
        self.synapses = nn.ModuleList([Synapse(self.n_dim) for _ in range(self.n_synapses)])
        self.ln = nn.LayerNorm(self.n_dim, bias=False)
        self.neighbors = []

    def forward(self, x, increment_neuron_count_func, attention_mask=None, y=None):
        if self.activated:
            return x
        self.activated = True
        increment_neuron_count_func()
        x = self.process(x, attention_mask, y)
        x = self.propagate(x)
        x = self.connect(x, increment_neuron_count_func)
        return x

Strati di Input e Output

### Reading e Bridge Le classi Reading e Bridge gestiscono diversi tipi di elaborazione di input e output. La classe Reading è utilizzata per l'incorporamento dell'input testuale, mentre la classe Bridge può essere utilizzata per l'elaborazione intermedia o come strato di output.

class Reading(nn.Module):
    def __init__(self, n_embed_size, n_vocab_size, n_pos_size):
        super(Reading, self).__init__()
        self.wte = nn.Embedding(self.n_vocab_size, self.n_embed_size)
        self.wpe = nn.Embedding(self.n_pos_size, self.n_embed_size)
        self.fc_out = nn.Linear(self.n_embed_size, self.n_embed_size)
        self.ln = nn.LayerNorm(self.n_embed_size)
        self.act = nn.SiLU()

    def forward(self, input_ids):
        position_ids = torch.arange(0, input_ids.shape[-1], dtype=torch.long, device=input_ids.device)
        inputs_embeds = self.wte(input_ids)
        position_embeds = self.wpe(position_ids.unsqueeze(0))
        hidden_states = inputs_embeds + position_embeds
        hidden_states = self.fc_out(hidden_states)
        hidden_states = self.ln(hidden_states)
        hidden_states = self.act(hidden_states)
        return hidden_states

CausalLM e SequenceClassifier

Queste classi definiscono l'elaborazione finale dell'output, che sia per il modeling del linguaggio o per la classificazione delle sequenze.

class CausalLM(nn.Module):
    def __init__(self, n_dim, n_vocab_size):
        super(CausalLM, self).__init__()
        self.lm_head = nn.Linear(self.n_dim, self.n_vocab_size)

    def forward(self, x):
        return self.lm_head(x)

class SequenceClassifier(nn.Module):
    def __init__(self, n_dim, n_outputs):
        super(SequenceClassifier, self).__init__()
        self.fc_out = nn.Linear(self.n_dim, self.n_outputs)

    def forward(self, x):
        x = torch.mean(x, dim=1)
        x = self.fc_out(x)
        x = torch.softmax(x, dim=-1) if self.n_outputs > 1 else torch.sigmoid(x)
        return x

Conclusione

Genesis rappresenta un passo avanti nell'architettura delle reti neurali, traendo ispirazione dal cervello umano. Il suo design modulare, la configurabilità e i calcoli plausibili biologicamente lo rendono un approccio promettente per vari compiti di AI. Sviluppare questo modello è stato un viaggio affascinante, combinando la ricerca teorica con le sfide pratiche dell'implementazione. Non vedo l'ora di ottimizzare ulteriormente e ampliare le sue capacità.

Questo progetto è open-source e accetto volentieri collaborazioni da chiunque sia interessato a esplorare le intersezioni tra intelligenza artificiale e biologica. Spingiamo insieme i confini di ciò che è possibile!