Foto di DeepMind on Unsplash
La moltiplicazione di matrici è un'operazione fondamentale utilizzata in molti sistemi, dalle reti neurali alle routine di calcolo scientifico. Trovare algoritmi efficienti e dimostrabilmente corretti per la moltiplicazione di matrici può avere un enorme impatto nel rendere il calcolo più veloce ed efficiente, ma è un compito molto impegnativo. Lo spazio dei possibili algoritmi è enorme e i metodi tradizionali per scoprire algoritmi, come l'euristica progettata dall'uomo o la ricerca combinatoria, sono spesso subottimali.
DeepMindLa soluzione basata sull'intelligenza artificiale recentemente proposta da per la ricerca automatizzata va ben oltre l'intuizione umana. La soluzione consiste in un agente di apprendimento per rinforzo profondo chiamato AlphaTensor, costruito sopra Alpha Zero. Questo agente è addestrato a giocare a un gioco per giocatore singolo, TensorGame, in cui l'obiettivo è scoprire algoritmi computazionalmente efficienti per la moltiplicazione di matrici.
AlphaTensor è particolarmente bravo a gestire grandi matrici scomponendo grandi moltiplicazioni di matrici in moltiplicazioni più piccole. Inoltre, AlphaTensor può essere utilizzato per ottenere prestazioni all'avanguardia per la moltiplicazione di matrici una volta messo a punto su uno specifico dispositivo hardware.
AlphaTensor ha un grande potenziale per accelerare il deep learning computing. Nel deep learning, molte operazioni dispendiose in termini di tempo possono essere mappate a moltiplicazioni di matrici. Utilizzando AlphaTensor per ottimizzare queste operazioni, le prestazioni complessive dei modelli di deep learning possono essere notevolmente migliorate.
Recentemente, OpenAlphaTensor, il prima implementazione open source di AlphaTensor, che potrebbe rivoluzionare la potenza computazionale dei modelli di deep learning.
Tensore di moltiplicazione di matrici
Per i non esperti nell'ottimizzazione della moltiplicazione di matrici, potrebbe non essere semplice capire come un'operazione come la moltiplicazione di matrici possa essere mappata in un tensore tridimensionale. Cercherò di spiegarlo con parole semplici e con esempi.
Consideriamo il prodotto C = A*B, dove per semplicità sia A che B sono matrici quadrate di dimensione N. L'operazione di moltiplicazione può essere mappata in un tensore 3D di forma (N^2, N^2, N^2). La prima dimensione del tensore rappresenta la matrice appiattita A, la seconda dimensione la matrice appiattita B e la terza dimensione la matrice appiattita C.
Il tensore ha solo valori binari (1 o 0) per ogni voce. Si noti che il tensore rappresenta l'operazione di moltiplicazione, quindi è indipendente dai valori delle matrici A e B.
Ogni voce del tensore corrisponde al coefficiente dell'operazione. Ad esempio, per calcolare C[1,1], è necessario moltiplicare sia A[1,1] che B[1,1]. Pertanto, la voce tensoriale [0,0,0], che corrisponde ad A[1,1], B[1,1] e C[1,1], avrà valore 1. Al contrario, per calcolare C[1,1 ,2,1], A[1] non è necessario. Pertanto, la riga del tensore T[N+0, :, XNUMX] conterrà solo zeri.
L'immagine sotto mostra un esempio di tensore per N=2.
Immagine da DeepMind carta pubblicato nella Natura
Come mostrato in (b) e (c) nella figura sopra, è possibile implementare un algoritmo per calcolare il prodotto utilizzando una decomposizione del tensore 3D. Più specificamente, l'algoritmo seguente può essere utilizzato per convertire una decomposizione tensoriale (le matrici U, V, W) in un algoritmo di moltiplicazione di matrici.
Meta-algoritmo parametrizzato per il calcolo del prodotto di matrice C=AB introdotto in DeepMind carta
Il TensorGame
Il problema di trovare algoritmi efficienti per la moltiplicazione di matrici è estremamente impegnativo perché il numero di possibili algoritmi da considerare è molto più grande del numero di atomi nell'universo, anche per piccoli casi di moltiplicazione di matrici.
DeepMind ha convertito questo problema in un gioco per giocatore singolo e lo ha chiamato TensorGame. In questo gioco, il giocatore sceglie come combinare diverse voci di matrici per moltiplicarle. Viene assegnato un punteggio in base al numero di operazioni necessarie per ottenere il risultato corretto della moltiplicazione. Il gioco termina quando viene raggiunto il tensore zero o quando è stato effettuato il numero massimo di mosse. La fattorizzazione finale viene valutata sulla base di una stima del rango residuo e di determinati criteri di ottimizzazione, come la complessità temporale asintotica o il runtime pratico.
La posizione iniziale nel TensorGame corrisponde al tensore di moltiplicazione della matrice espresso su una base casuale.
In ogni passo t del gioco, il giocatore scrive tre vettori
, che specifica i tensori di rango 1 . Lo stato del gioco viene aggiornato sottraendo i vettori selezionati dal giocatore:
where
è il tensore di moltiplicazione della matrice.Se il gioco termina in p passaggi, significa che il tensore di moltiplicazione della matrice
può essere scomposto in tensori p rango-1 , cioè ha almeno rango p.Il TensorGame può quindi essere interpretato come un algoritmo di scomposizione del rango e AlphaTensor può essere visto come un algoritmo per stimare il rango del tensore.
Architettura AlphaTensor
Finora abbiamo appreso del TensorGame e chiarito come la sua soluzione possa essere vista come un algoritmo di moltiplicazione di matrici. Esploriamo ora i concetti principali di AlphaTensor, l'algoritmo utilizzato per il gioco.
L'architettura AlphaTensor è fondamentalmente un'architettura Transformer encoder-decoder in cui:
- il codificatore prende come input lo stato del gioco , le n azioni precedenti intraprese dal modello (solitamente n=7) e l'indice temporale t dell'azione corrente. Le informazioni sono impilate insieme in un tensore con forma (n+1, N^2, N^2, N^2). Questo tensore viene quindi rimodellato e trasformato (utilizzando tre strati lineari) in un tensore di forma (N^2, N^2, c) dove c è la dimensione interna del modello.
- il decoder genera le azioni n_steps dal vettore embedded fornito dall'encoder in modo autoregressivo. Ogni azione corrisponde a un gettone delle terzine che rappresenta una delle terzine che scompongono il tensore del gioco (cioè riducendone il rango)
Il modello viene addestrato alternando la propagazione all'indietro e la recitazione del modello. La recitazione del modello viene utilizzata per generare dati che vengono quindi utilizzati per addestrare il modello. In pratica, il modello viene addestrato con una combinazione di dati generati sinteticamente e dati generati dal modello durante la recitazione. La fase di recitazione viene eseguita prendendo un tensore 3D corrispondente a un'operazione di matrice e giocando a n_attori su di esso. Ogni attore gioca un gioco su base standard o su base alternativa (il cambio di base viene applicato con una data probabilità). I risultati vengono quindi raccolti e possono essere utilizzati nella fase di addestramento con i dati sintetici.
La fase di recitazione si basa sul Monte Carlo Tree Search (MCTS) di AlphaZero, modificato per supportare ampi spazi d'azione. In breve, prima di scegliere l'azione, vengono esplorati i percorsi n_sims dall'output del modello con un'esplorazione futura massima di 5 passaggi. Le probabilità generate dal modello vengono quindi adeguate tenendo conto dei percorsi generati. Quindi viene scelta l'azione con il/i percorso/i futuro/i più promettente/i per continuare il gioco.
Durante l'addestramento del modello, la ricompensa è in realtà una ricompensa negativa (penalità). Il suo valore assoluto aumenta con ogni passo aggiuntivo richiesto per risolvere il gioco. Se il modello impiega m passi per risolvere un TensorGame, la ricompensa associata al gioco è r=-m. Se il modello non è in grado di risolvere il TensorGame in step max_rank, la ricompensa viene calcolata stimando il rango del tensore rimanente. Il rango è stimato come somma dei ranghi delle matrici che compongono il tensore. La stima è un limite superiore al vero rango del tensore.
Durante la messa a punto del modello, la penalità ricompensata allo stato terminale dovrebbe tenere conto anche della latenza dell’algoritmo prodotto dal modello. La formula della ricompensa diventa rt'=rt+λbt, dove rt è lo schema di ricompensa descritto in precedenza, bt è la ricompensa di riferimento (diversa da zero solo allo stato terminale) e λ è un coefficiente specificato dall'utente.
Accelerazioni (%) degli algoritmi scoperti da AlphaTensor su misura per una GPU e una TPU, estratti dall'articolo di DeepMind. Le accelerazioni sono misurate rispetto alla moltiplicazione di matrice standard (ad es. cuBLAS per la GPU) sullo stesso hardware e confrontate con la Algoritmo del quadrato di Strassen. Fonte: DeepMind.
l'ho rilasciato di recente OpenAlphaTensor, la prima implementazione open source di AlphaTensor. In questa sezione illustrerò l'implementazione. Come abbiamo discusso in precedenza, l'architettura AlphaTensor è abbastanza semplice, basata su un trasformatore standard con un'architettura codificatore-decodificatore. I componenti più interessanti di AlphaTensor sono il primo strato nella parte del codificatore e il modo in cui le azioni vengono campionate.
Iniziamo con il primo livello di codifica.
# x.size = (N, T, S, S, S)
# scalars.size = (N, s)
batch_size = x.shape[0]
S = x.shape[-1]
T = x.shape[1]
x1 = x.permute(0, 2, 3, 4, 1).reshape(batch_size, S, S, S * T)
x2 = x.permute(0, 4, 2, 3, 1).reshape(batch_size, S, S, S * T)
x3 = x.permute(0, 3, 4, 2, 1).reshape(batch_size, S, S, S * T)
input_list = [x1, x2, x3]
for i in range(3): temp = self.linears_1[i](scalars).reshape(batch_size, S, S, 1) input_list[i] = torch.cat([input_list[i], temp], dim=-1) input_list[i] = self.linears_2[i](input_list[i])
x1, x2, x3 = input_list
Nello snippet sopra, mostriamo come il tensore di input viene scomposto in tre tensori, che vengono quindi utilizzati come input di query, chiave e valore del livello trasformatore.
- Attraverso le tre dimensioni del tensore che rappresentano le matrici appiattite (A, B, C), il tensore di input viene appiattito lungo ciascuna dimensione insieme alla dimensione che rappresenta le azioni precedenti. In questo modo, in ogni copia appiattita del tensore di input, la dimensione selezionata è un'aggregazione degli ultimi valori T-1 e del valore effettivo, per tutti i valori S della dimensione selezionata, dove S=N^2. Filosoficamente, è come se, per ogni dimensione, ci concentrassimo su ciò che è accaduto nelle azioni precedenti in quella dimensione.
- Gli scalari sono mappati in tre diversi spazi di dimensione S^2, e poi rimodellati per essere concatenati con i tensori ottenuti al punto precedente. Concettualmente, gli scalari sono mappati su uno spazio di incorporamento di dimensione S^2, quindi le informazioni incorporate vengono suddivise in vettori S e impilate insieme, in modo simile a ciò che accade al testo quando viene tokenizzato.
- I token scalari vengono concatenati con il tensore di input ristrutturato e quindi forniti come input a un livello lineare per mappare le informazioni del focus scalare + cronologia del canale nella dimensione interna del modello.
Questi tre passaggi possono essere interpretati come un modo per fornire al modello sia informazioni sugli scalari (come nel passaggio temporale di TensorGame) sia il focus sulle azioni precedenti per ciascun canale.
Per quanto riguarda il modo in cui vengono prodotte le azioni, è interessante notare che AlphaTensor genera in output la tripletta u, v, w, che mira a ridurre il rango tensoriale. I tre vettori hanno dimensione S e poiché sono concatenati il modello deve produrre un vettore di dimensione 3*S. AlphaTensor è addestrato con un algoritmo RL, quindi tutte le possibili azioni devono essere espresse in termini di probabilità in uno spazio enumerato, cioè il modello produce una probabilità sulle diverse azioni. Ciò significa che ogni vettore nello spazio 3S dovrebbe essere mappato su un'azione diversa. Ciò si traduce in uno spazio azione di dimensione |F|^(3S), dove |F| è il numero di valori diversi che può assumere l'elemento di u, v, w. Di solito, i valori sono limitati a (-2, -1, 0, 1, 2), risultando in una cardinalità di 5 elementi.
Qui arriva una sfida importante: per generare le probabilità di azione per un prodotto matriciale di matrici di dimensione 5 avremmo bisogno di una memoria di 5^75 * 4 byte, il che significherebbe ~ 10^44 GB di memoria. Chiaramente, non possiamo gestire uno spazio d'azione così ampio.
Come risolviamo il problema? Per ridurre l'impronta di memoria delle probabilità di azione possiamo suddividere le terzine in blocchi più piccoli, "tokenizzarli" e trattare i blocchi come token generati nell'architettura del trasformatore, ovvero i token vengono dati come input al decodificatore in un processo autoregressivo modo. Nell'esempio precedente possiamo suddividere le terzine in 15 blocchi, riducendo il consumo di memoria a 15 * 5^(75/15) * 4, ovvero 187.5 KB.
def _eval_forward(self, e: torch.Tensor): bs = e.shape[0] future_g = ( torch.zeros((bs, self.n_samples, self.n_steps)).long().to(e.device) ) ps = torch.ones((bs, self.n_samples)).to(e.device) e = e.unsqueeze(1).repeat(1, self.n_samples, 1, 1) future_g = future_g.view(-1, self.n_steps) ps = ps.view(-1) e = e.view(-1, e.shape[-2], e.shape[-1]) for i in range(self.n_steps): o_s, z_s = self.core(future_g[:, : i + 1], e) future_g[:, i], p_i = sample_from_logits(o_s[:, i]) ps *= p_i future_g = future_g.view(bs, self.n_samples, self.n_steps) ps = ps.view(bs, self.n_samples) return ( future_g, ps, z_s[:, 0].view(bs, self.n_samples, *z_s.shape[2:]).mean(1), )
Sopra mostriamo lo snippet di codice per generare l'azione completa. Nel codice, self.core contiene il livello del decodificatore e il tensore e rappresenta l'output del livello del codificatore. Zero può essere considerato come il token nei modelli NLP e le azioni n_steps che rappresentano i chunk n_steps vengono generate in modo progressivo.
Il modello restituisce tre quantità:
- Le azioni generate
- La probabilità associata all'azione completa
- I logit prodotti per generare la prima azione (il primo blocco) che verrà utilizzato per calcolare il valore del modello.
Vale la pena spendere qualche parola sul parametro n_samples. Il parametro viene utilizzato per la fase di recitazione e consente al modello di generare diverse versioni delle terzine che verranno quindi utilizzate per esplorare lo spazio di azione nell'algoritmo di ricerca dell'albero di Monte Carlo utilizzato nel processo di recitazione. Le diverse azioni n_samples vengono campionate in base alla policy generata dal modello.
Passo di recitazione
La parte più complicata dell'intero algoritmo è probabilmente la fase di recitazione utilizzata per risolvere il TensorGame. L'algoritmo non è spiegato a fondo nell'articolo di AlphaTensor, poiché si basa su diversi articoli precedenti di DeepMind che sono solo citati e dati come noti. Qui ricostruirò tutti i pezzi mancanti e spiegherò passo dopo passo la nostra implementazione.
Possiamo organizzare le fasi di recitazione in tre diverse componenti:
- La ricerca dell'albero di Monte-Carlo
- La simulazione del gioco
- Il calcolo delle politiche migliorato
Analizziamoli uno per uno.
Ricerca di alberi a Monte-Carlo (MCTS)
Monte Carlo Tree Search (MCTS) è una tecnica di intelligenza artificiale ampiamente utilizzata per il gioco, in particolare nei giochi da tavolo e nei videogiochi. L'algoritmo crea un albero di gioco che simula potenziali mosse e risultati e utilizza il campionamento casuale per valutare la ricompensa prevista per ogni mossa. L'algoritmo quindi seleziona in modo iterativo la mossa con la ricompensa prevista più alta e simula i risultati finché non raggiunge uno stato terminale o una condizione di arresto specificata. Le simulazioni vengono utilizzate per stimare la probabilità di vincita per ogni mossa e guidare il processo decisionale. MCTS ha dimostrato di essere efficace in giochi complessi in cui il numero di possibili mosse e risultati è elevato ed è stato utilizzato in sistemi di intelligenza artificiale di gioco di successo, come AlphaGo.
In AlphaTensor viene utilizzata una versione modificata dell'originale MCTS. In particolare, invece di selezionare casualmente l'azione dall'intero spazio azione, l'azione viene selezionata tra un sottoinsieme generato direttamente dal modello (attraverso gli n_campioni presentati prima). La correzione all'aggiornamento dei criteri viene quindi applicata nella fase di calcolo dei criteri migliorati.
Nella nostra implementazione, abbiamo deciso di mantenere tutte le informazioni sull'albero Monte-Carlo in un dizionario avente come chiave la versione hash dello stato di TensorGame e come valori le informazioni associate allo stato stesso. Ogni passo Monte-Carlo parte da un nodo e simula minigiochi n_sim, esplorando il futuro con un orizzonte di 5 mosse. Se il nodo è già stato esplorato in simulazioni precedenti, n_sim viene aggiustato considerando il numero di esplorazioni precedenti. Per ogni nodo il numero di visite è memorizzato nel tensore N_s_a, poiché questo tensore contiene il numero di visite per azione figlia del nodo (tra quelle campionate dal modello).
def monte_carlo_tree_search( model: torch.nn.Module, state: torch.Tensor, n_sim: int, t_time: int, n_steps: int, game_tree: Dict, state_dict: Dict,
): """Runs the monte carlo tree search algorithm. Args: model (torch.nn.Module): The model to use for the simulation. state (torch.Tensor): The initial state. n_sim (int): The number of simulations to run. t_time (int): The current time step. n_steps (int): The maximum number of steps to simulate. game_tree (Dict): The game tree. state_dict (Dict): The dictionary containing the states. """ state_hash = to_hash(extract_present_state(state)) if state_hash in state_dict: with torch.no_grad(): N_s_a = state_dict[state_hash][3] n_sim -= int(N_s_a.sum()) n_sim = max(n_sim, 0) for _ in range(n_sim): simulate_game(model, state, t_time, n_steps, game_tree, state_dict) # return next state possible_states_dict, _, repetitions, N_s_a, q_values, _ = state_dict[ state_hash ] possible_states = _recompose_possible_states(possible_states_dict) next_state_idx = select_future_state( possible_states, q_values, N_s_a, repetitions, return_idx=True ) next_state = possible_states[next_state_idx] return next_state
Il codice sopra mostra la nostra implementazione dell'algoritmo. Per una questione di semplicità del codice, la correzione della politica viene eseguita nella funzione simulate_game.
Simulazione di gioco
La funzione simulate_game è responsabile dell'esplorazione dell'albero composto da nodi che rappresentano un particolare stato del TensorGame. Inoltre esegue il modello ogni volta che viene rilevato un nodo foglia e memorizza tutte le informazioni sul nodo nel dizionario state_dict. Diamo uno sguardo approfondito alla sua implementazione:
@torch.no_grad()
def simulate_game( model, state: torch.Tensor, t_time: int, max_steps: int, game_tree: Dict, states_dict: Dict, horizon: int = 5,
): """Simulates a game from a given state. Args: model: The model to use for the simulation. state (torch.Tensor): The initial state. t_time (int): The current time step. max_steps (int): The maximum number of steps to simulate. game_tree (Dict): The game tree. states_dict (Dict): The states dictionary. horizon (int): The horizon to use for the simulation. """ idx = t_time max_steps = min(max_steps, t_time + horizon) state_hash = to_hash(extract_present_state(state)) trajectory = [] # selection while state_hash in game_tree: ( possible_states_dict, old_idx_to_new_idx, repetition_map, N_s_a, q_values, actions, ) = states_dict[state_hash] possible_states = _recompose_possible_states(possible_states_dict) state_idx = select_future_state( possible_states, q_values, N_s_a, repetition_map, return_idx=True ) trajectory.append((state_hash, state_idx)) # state_hash, action_idx future_state = extract_present_state(possible_states[state_idx]) state = possible_states[state_idx] state_hash = to_hash(future_state) idx += 1 # expansion if idx = max_steps: trajectory.append((state_hash, None)) if not game_is_finished(extract_present_state(state)): state = state.to(model.device) scalars = get_scalars(state, idx).to(state.device) actions, probs, q_values = model(state, scalars) ( possible_states, cloned_idx_to_idx, repetitions, not_dupl_indexes, ) = extract_children_states_from_actions( state, actions, ) not_dupl_actions = actions[:, not_dupl_indexes].to("cpu") not_dupl_q_values = torch.zeros(not_dupl_actions.shape[:-1]).to( "cpu" ) N_s_a = torch.zeros_like(not_dupl_q_values).to("cpu") present_state = extract_present_state(state) states_dict[to_hash(present_state)] = ( _reduce_memory_consumption_before_storing(possible_states), cloned_idx_to_idx, repetitions, N_s_a, not_dupl_q_values, not_dupl_actions, ) game_tree[to_hash(present_state)] = [ to_hash(extract_present_state(fut_state)) for fut_state in possible_states ] leaf_q_value = q_values else: leaf_q_value = -int(torch.linalg.matrix_rank(state).sum()) # backup backward_pass(trajectory, states_dict, leaf_q_value=leaf_q_value)
Ogni simulazione è divisa in tre parti:
- Selezione
- Espansione
- di riserva
Nella parte di selezione la simulazione viene eseguita sui nodi dell'albero già generati e il seguente nodo viene selezionato utilizzando la seguente funzione:
def select_future_state( possible_states: List[torch.Tensor], q_values: torch.Tensor, N_s_a: torch.Tensor, repetitions: Dict[int, list], c_1: float = 1.25, c_2: float = 19652, return_idx: bool = False,
) -> torch.Tensor: """Select the future state maximizing the upper confidence bound."""
# q_values (1, K, 1) pi = torch.tensor( [ len(repetitions[i]) for i in range(len(possible_states)) if i in repetitions ] ).to(q_values.device) ucb = q_values.reshape(-1) + pi * torch.sqrt( torch.sum(N_s_a) / (1 + N_s_a) ) * (c_1 + torch.log((torch.sum(N_s_a) + c_2 + 1) / c_2)) if return_idx: return ucb.argmax() return possible_states[ucb.argmax()]
In pratica, l'azione che massimizza la funzione ucb:
per lo stato specificato è selezionato. Qui Q rappresenta i valori Q generati dal modello e π rappresenta la distribuzione casuale sulle azioni campionate utilizzando la politica del modello. N(s, a) rappresenta il numero di visite del nodo all'azione a dal nodo s.
Una volta che la fase di selezione raggiunge un nodo foglia, se la simulazione non ha raggiunto una condizione terminale (in termini di massima esplorazione, cioè orizzonte futuro, o fine del gioco), il modello viene quindi utilizzato per selezionare n_samples nodi alternativi (saranno nodi foglia nodi nell'iterazione successiva). Questa è chiamata la fase di espansione, poiché nuovi nodi vengono aggiunti all'albero. Quindi, nessun altro nodo viene esplorato nella simulazione corrente, ma la foglia q_value viene inviata al successivo passo della simulazione: il backup.
Il backup è la fase finale di ogni simulazione. Durante il backup, se il nodo foglia era uno stato terminale, viene calcolata la ricompensa finale; in caso contrario, il valore foglia q viene utilizzato come ricompensa stimata. Quindi la ricompensa viene retropropagata sulla traiettoria della simulazione aggiornando sia gli stati q_values che aggiornando il contatore di visite N(s, a). Nello snippet sottostante mostriamo il codice per la propagazione all'indietro della ricompensa.
def backward_pass(trajectory, states_dict, leaf_q_value: torch.Tensor): """Backward pass of the montecarlo algorithm"""
reward = 0 for idx, (state, action_idx) in enumerate(reversed(trajectory)): if action_idx is None: # leaf node reward += leaf_q_value else: ( _, old_idx_to_new_idx, _, N_s_a, q_values, _, ) = states_dict[state] if isinstance(reward, torch.Tensor): reward = reward.to(q_values.device) action_idx = int(action_idx) if action_idx in old_idx_to_new_idx: not_dupl_index = old_idx_to_new_idx[int(action_idx)] else: not_dupl_index = action_idx reward -= 1 q_values[:, not_dupl_index] = ( N_s_a[:, not_dupl_index] * q_values[:, not_dupl_index] + reward ) / (N_s_a[:, not_dupl_index] + 1) N_s_a[:, not_dupl_index] += 1
Calcolo delle politiche migliorato
Una volta che tutte le simulazioni sono state eseguite e l'MCTS offre un'interessante istantanea del prossimo futuro, è il momento di aggiornare la politica associata ai nodi previsti e restituirli, in modo che possano essere utilizzati durante l'addestramento. La politica migliorata, seguendo il metodo descritto in Hubert et al, viene utilizzato per la gestione di ampi spazi di azione. Infatti, per un piccolo spazio di ricerca, è possibile durante MCTS campionare un'azione in modo casuale dallo spazio dell'azione e valutarne l'impatto. Un approccio simile in uno spazio di azione molto più ampio porterebbe a tutte le traiettorie divergenti in percorsi diversi e richiederebbe un numero infinito di traiettorie per ottenere statistiche significative e quindi aggiornare la politica. Poiché qui stiamo usando sample-MCTS per evitare la dispersione, ovvero le azioni n_samples sono campionate in base alla politica del modello e quindi MCTS seleziona solo una delle azioni campionate durante l'esplorazione dell'albero, dobbiamo tenere conto della correzione del campione durante il calcolo la policy aggiornata finale che verrà utilizzata durante il training del modello.
In pratica, la politica migliorata viene calcolata come
where
def compute_improved_policy( state_dict: Dict, states: List[str], model_n_steps: int, model_n_logits: int, N_bar: int,
): """Compute the improved policy given the state_dict, the list of states. The improved policy is computed as (N_s_a / N_s_a.sum())^(1/tau) where tau is (log(N_s_a.sum()) / log(N_bar)) if N_s_a.sum() > N_bar else 1. """ policies = torch.zeros(len(states), model_n_steps, model_n_logits) N_bar = torch.tensor(N_bar) for idx, state in enumerate(states): N_s_a = state_dict[state][3] actions = state_dict[state][5] if N_s_a.sum() > N_bar: tau = (torch.log(N_s_a.sum()) / torch.log(N_bar)).item() else: tau = 1 N_s_a = N_s_a ** (1 / tau) improved_policy = N_s_a / N_s_a.sum() for sample_id in range(actions.shape[1]): action_ids = actions[0, sample_id] for step_id, action_id in enumerate(action_ids): policies[idx, step_id, action_id] += improved_policy[ 0, sample_id ] return policies
Si noti che nella nostra implementazione, dopo aver calcolato la politica dal tensore N_s_a, dobbiamo riportarla al tensore azione originale. Infatti, N_s_a considera solo le azioni campionate dal modello, mentre la policy finale deve contenere probabilità anche per le azioni non esplorate.
Differenze rispetto all'algoritmo di addestramento ChatGPT
AlphaTensor è l'ultimo membro della famiglia AlphaGo/AlphaZero di metodi di intelligenza artificiale di DeepMind. Questi metodi si basano sull’algoritmo Monte Carlo Tree Search (MCTS), che è stato perfezionato e migliorato da DeepMind per affrontare compiti sempre più complessi. Un altro sistema di intelligenza artificiale, ChatGPT di OpenAI, che ha suscitato molto entusiasmo per le sue notevoli prestazioni, è stato addestrato con un approccio diverso, chiamato Reinforcement Learning with Human Feedback (RLHF).
RLHF è una tecnica di messa a punto utilizzata per ottimizzare i modelli linguistici in modo che seguano una serie di istruzioni scritte. Utilizza le preferenze umane come segnale di ricompensa per mettere a punto il modello, allineando così il comportamento del modello linguistico con le preferenze dichiarate di un gruppo specifico di persone, piuttosto che con una nozione più ampia di "valori umani".
Al contrario, MCTS è un algoritmo di ricerca basato su albero utilizzato per determinare le mosse ottimali nei giochi. Simula le mosse potenziali e aggiorna i valori di ciascuna mossa in base ai loro risultati, guidando la selezione della mossa migliore.
RLHF raccoglie dati da dimostrazioni scritte dall'uomo e confronti etichettati dall'uomo tra modelli di intelligenza artificiale e addestra un modello di ricompensa per prevedere le preferenze di un determinato gruppo di persone. Il modello di ricompensa viene quindi utilizzato per mettere a punto i modelli di intelligenza artificiale. MCTS, invece, utilizza simulazioni e valutazioni per determinare la decisione migliore.
Sebbene siano approcci diversi, RLHF e MCTS hanno anche somiglianze. Entrambe le tecniche di intelligenza artificiale utilizzano metodi decisionali e di risoluzione dei problemi ed entrambe utilizzano un approccio per tentativi ed errori per esplorare diverse opzioni e prendere decisioni in base alle informazioni disponibili. Entrambi sono anche processi iterativi che migliorano nel tempo man mano che vengono raccolte più informazioni ed esperienza.
La scelta tra RLHF e MCTS dipende dall'attività da svolgere. RLHF è ideale quando non esiste una metrica chiara per valutare le prestazioni del modello, mentre MCTS si è dimostrato efficace in attività simili a giochi in cui la conoscenza e l'esplorazione del futuro conferiscono al modello un vantaggio significativo.
Ottimizzazione del codice per l'addestramento AlphaTensor
L'implementazione dell'algoritmo di allenamento AlphaTensor richiede di trovare il perfetto compromesso tra velocità di allenamento e consumo di memoria. Come visto nella sezione Modello, la semplice considerazione della tokenizzazione dell'azione può far risparmiare molta memoria, ma una riduzione dello spazio di azione eccessivamente aggressiva può portare sia a un calo della precisione che a prestazioni più lente. Quest'ultimo accade perché tutti i token vengono generati sequenzialmente in modo autoregressivo dal decodificatore del modello. Pertanto, il tempo di inferenza cresce linearmente con il numero di token per azione una volta che il softmax nello spazio azione non è più il collo di bottiglia.
Durante l'impostazione della formazione AlphaTensor, le principali difficoltà sono state riscontrate nell'affrontare il processo di recitazione. Se i tensori non sono memorizzati nel formato corretto, l'MCTS può facilmente causare un aumento incontrollato dell'utilizzo della memoria. D'altra parte, se il numero di tensori memorizzati durante ogni simulazione viene ridotto troppo, l'MCTS può impiegare una quantità infinita di tempo a ricalcolare gli stati richiesti.
Facciamo un esempio della fase di simulazione del gioco, in cui il gioco viene esplorato osservando possibili scenari futuri. Per ogni stato, se non salviamo le azioni generate dal modello e decidiamo di salvare solo il seme casuale utilizzato per campionare le azioni dalla policy, allora ogni volta che esploriamo un nodo dell'albero dovremmo ricalcolare la policy e quindi provare le azioni. Chiaramente abbiamo deciso di archiviare le azioni campionate per risparmiare tempo ed evitare di dover gestire la condivisione del modello tra diversi processi nel caso di parallelizzazione dell'esplorazione MCTS. Tuttavia, il semplice salvataggio delle azioni non è stato sufficiente per ottenere una fase di azione sufficientemente efficiente. Infatti, il tempo necessario per convertire le azioni n_steps nella tripletta (u, v, w), ridurre lo stato tensore del gioco e creare i nuovi tensori 3D dalle azioni n_samples rappresenterebbe facilmente un collo di bottiglia per l'intero addestramento. In secondo luogo, non volevamo memorizzare tutti i possibili stati futuri per ciascuna azione campionata, poiché ciò avrebbe avuto un impatto enorme sulla memoria utilizzata dall'algoritmo. Supponiamo di impostare n_samples=32, n=7 e N=5, e ricordiamo che N è la dimensione del prodotto di matrice quadrata che vogliamo ridurre e n è il numero di azioni precedenti ricordate dal modello. In questa situazione, ogni tensore di stato avrebbe la forma (8, 25, 25, 25), che moltiplicato per 32 darebbe come risultato 3282525254 byte per ogni nodo nel grafico. Ora, considerando che ogni simulazione in fase di espansione genera un nuovo nodo (e n_sim=200), avremmo un consumo di memoria finale di 200328252525*4 = 3.2 GB solo per il primo nodo MCTS. Nello scenario peggiore, durante l'esplorazione dei nodi max_rank attivi (dove max_rank=150), ciò comporterebbe un consumo totale di memoria di 150 * 3.2 GB = 480 GB nella memoria RAM (o memoria GPU se tutti i tensori fossero archiviati nella GPU) . Abbiamo eseguito la formazione sulla nostra workstation con 128 GB di RAM e 48 GB di memoria GPU, quindi abbiamo dovuto ridurre il consumo di memoria.
Non volendo aumentare i tempi di esecuzione, abbiamo adottato un'ottimizzazione che sfrutta la ridondanza nei tensori di stato prodotti. Infatti i tensori hanno in comune n-1 azioni precedenti, che possono quindi essere memorizzate una volta e non ripetute per ogni tensore memorizzato. Ciò si traduce in una riduzione della memoria del 2/7~28%, il che significa che nel peggiore dei casi possono essere archiviati 137GB. A questo punto, semplicemente potando la parte inutilizzata dell'albero (come le traiettorie non selezionate) e memorizzando i tensori nella memoria della CPU, siamo riusciti a evitare qualsiasi errore di memoria durante l'addestramento.
Con OpenAlphaTensor che ora è open source, si aprono diverse strade entusiasmanti per ulteriori sviluppi.
Una progressione naturale è la messa a punto di OpenAlphaTensor sui dispositivi hardware di destinazione. Ciò dovrebbe portare a prestazioni computazionali molto competitive. Pubblicherò di più sulle prestazioni di OpenAlphaTensor su vari hardware su GitHub. Al momento della stesura di questo articolo, OpenAlphaTensor era in fase di formazione.
Un altro importante progresso sarebbe il supporto per la compilazione remota, consentendo agli utenti di creare algoritmi ottimizzati per dispositivi edge. Ciò può essere ottenuto memorizzando il modello OpenAlphaTensor su un server, mentre l'algoritmo di moltiplicazione della matrice viene valutato su hardware diverso.
Potrebbe anche essere importante estendere il supporto a diversi compilatori per calcolare la correzione della ricompensa basata sulla latenza. Diversi compilatori possono portare a diversi algoritmi ottimizzati su un determinato hardware. Ad esempio, il documento DeepMind ha mostrato risultati promettenti utilizzando JAX e il compilatore XLA su TPU e GPU Nvidia. Sarebbe interessante valutarlo utilizzando NCCL su Nvidia o LLVM su CPU.
Infine, l'estensione del modello e dell'algoritmo di addestramento per supportare matrici di dimensioni maggiori rimane una grande sfida aperta. Attualmente, OpenAlphaTensor supporta una dimensione massima della matrice di 5, ma può essere applicata suddividendo le moltiplicazioni di matrici più grandi in gruppi di minuscoli MM con una dimensione inferiore a 5. Questo approccio non è ottimale ed esegue la riduzione direttamente sul tensore grande corrispondente al il MM completo potrebbe teoricamente portare a risultati migliori.
Diego Fiore è il CTO di Nebuly AI, un'azienda impegnata a rendere l'ottimizzazione dell'intelligenza artificiale parte del toolkit di ogni sviluppatore.
- Distribuzione di contenuti basati su SEO e PR. Ricevi amplificazione oggi.
- Platoblockchain. Web3 Metaverse Intelligence. Conoscenza amplificata. Accedi qui.
- Fonte: https://www.kdnuggets.com/2023/03/first-open-source-implementation-deepmind-alphatensor.html?utm_source=rss&utm_medium=rss&utm_campaign=first-open-source-implementation-of-deepminds-alphatensor
- :È
- ][P
- $ SU
- 1
- 3d
- 8
- a
- capace
- Chi siamo
- sopra
- Assoluta
- accelerando
- Secondo
- di conseguenza
- Il mio account
- precisione
- Raggiungere
- raggiunto
- Action
- azioni
- effettivamente
- aggiunto
- aggiuntivo
- Rettificato
- adottato
- avanzare
- Vantaggio
- Dopo shavasana, sedersi in silenzio; saluti;
- Agente
- aggregazione
- aggressivo
- AI
- Sistemi di intelligenza artificiale
- mira
- algoritmo
- Algoritmi
- Tutti
- Consentire
- consente
- da solo
- già
- alternativa
- tra
- quantità
- analizzare
- ed
- Un altro
- applicato
- approccio
- approcci
- architettura
- SONO
- articolo
- artificiale
- intelligenza artificiale
- AS
- addetto
- associato
- At
- Automatizzata
- disponibile
- evitando
- precedente
- di riserva
- basato
- fondamentalmente
- base
- BE
- perché
- diventa
- prima
- essendo
- sotto
- Segno di riferimento
- MIGLIORE
- Meglio
- fra
- Al di là di
- tavola
- Giochi da tavolo
- legato
- più ampia
- BT
- costruire
- costruito
- by
- detto
- Materiale
- non può
- Custodie
- Causare
- ha causato
- certo
- Challenge
- impegnativo
- il cambiamento
- canale
- ChatGPT
- bambino
- scegliere
- la scelta
- scelto
- citato
- pulire campo
- chiaramente
- codice
- raccoglie
- combinare
- impegnata
- Uncommon
- azienda
- rispetto
- competitivo
- complesso
- complessità
- componenti
- composto
- compromesso
- calcolo
- potenza computazionale
- Calcolare
- informatica
- concetti
- concettualmente
- condizione
- fiducia
- Prendere in considerazione
- considerato
- considerando
- ritiene
- consumo
- contiene
- continua
- contrasto
- convertito
- Nucleo
- Corrispondente
- corrisponde
- potuto
- contatore
- CPU
- crea
- Creazione
- criteri
- CTO
- Corrente
- Attualmente
- dati
- trattare
- decide
- deciso
- decisione
- Decision Making
- decisioni
- deep
- apprendimento profondo
- DeepMind
- dipende
- descritta
- Determinare
- Costruttori
- Mercato
- dispositivo
- dispositivi
- DITT
- diverso
- le difficoltà
- Dimensioni
- dimensioni
- direttamente
- scopri
- scoprire
- discusso
- distribuzione
- Diviso
- giù
- Cadere
- durante
- e
- ogni
- In precedenza
- facilmente
- bordo
- Efficace
- efficiente
- o
- elemento
- elementi
- incorporato
- finisce
- migliorata
- enorme
- abbastanza
- iscrizione
- errore
- stima
- stimato
- Etere (ETH)
- valutare
- valutato
- la valutazione
- valutazioni
- Anche
- Ogni
- esempio
- Esempi
- coinvolgenti
- esecuzione
- espansione
- previsto
- esperienza
- Spiegare
- ha spiegato
- gesta
- esplorazione
- esplora
- Esplorazione
- Esplorare
- espresso
- estendere
- estendendo
- estremamente
- abbastanza
- famiglia
- più veloce
- feedback
- pochi
- figura
- finale
- ricerca
- Nome
- galleggiante
- Focus
- seguire
- i seguenti
- Orma
- Nel
- modulo
- formato
- formula
- essere trovato
- da
- pieno
- function
- fondamentale
- ulteriormente
- ulteriori sviluppi
- futuro
- gioco
- Giochi
- generare
- generato
- genera
- la generazione di
- ottenere
- ottenere
- Dare
- dato
- Dare
- scopo
- va
- buono
- GPU
- GPU
- grafico
- grande
- Gruppo
- Gruppo
- cresce
- Crescita
- guida
- cura
- Manovrabilità
- successo
- accade
- Hardware
- dispositivo hardware
- dispositivi hardware
- Avere
- avendo
- qui
- massimo
- orizzonte
- Come
- Tutorial
- Tuttavia
- HTTPS
- Enorme
- umano
- i
- MALATO
- ideale
- IDX
- Immagine
- Impact
- realizzare
- implementazione
- importante
- competenze
- migliorata
- in
- Aumento
- Aumenta
- sempre più
- studente indipendente
- Index
- informazioni
- inizialmente
- ingresso
- invece
- istruzioni
- Intelligence
- interessante
- interno
- introdotto
- intuizione
- IT
- iterazione
- SUO
- stessa
- jpg
- KDnuggets
- mantenere
- Le
- conoscenze
- conosciuto
- Lingua
- grandi
- superiore, se assunto singolarmente.
- Cognome
- Latenza
- con i più recenti
- strato
- galline ovaiole
- portare
- imparato
- apprendimento
- Lista
- Guarda
- cerca
- lotto
- fatto
- Principale
- maggiore
- make
- Fare
- gestire
- gestione
- molti
- carta geografica
- mappatura
- Matrice
- Importanza
- massimo
- significato
- significativo
- si intende
- membro
- Memorie
- metodo
- metodi
- metrico
- mancante
- miscela
- modello
- modelli
- modificato
- modulo
- Scopri di più
- più efficiente
- Inoltre
- maggior parte
- cambiano
- si muove
- moltiplicato
- Naturale
- Natura
- Vicino
- necessaria
- Bisogno
- di applicazione
- negativo.
- reti
- Neurale
- reti neurali
- New
- GENERAZIONE
- nlp
- nodo
- nodi
- non esperto
- Nozione
- numero
- Nvidia
- ottenuto
- of
- Offerte
- on
- ONE
- aprire
- open source
- OpenAI
- operazione
- Operazioni
- ottimale
- ottimizzazione
- OTTIMIZZA
- ottimizzati
- Opzioni
- i
- Altro
- altrimenti
- produzione
- complessivo
- Carta
- documenti
- parametro
- parte
- particolare
- particolarmente
- Ricambi
- Persone
- perfetta
- performance
- esecuzione
- fase
- pezzi
- Platone
- Platone Data Intelligence
- PlatoneDati
- Giocare
- giocatore
- gioco
- punto
- Termini e Condizioni
- politica
- posizione
- possibile
- potenziale
- energia
- Pratico
- pratica
- predire
- previsto
- preferenze
- presentata
- precedente
- probabilità
- probabilmente
- Problema
- processi
- i processi
- produrre
- Prodotto
- Prodotto
- progressione
- progressivo
- promettente
- proposto
- dimostrabile
- comprovata
- pubblicare
- pubblicato
- RAM
- casuale
- ranghi
- piuttosto
- a raggiunto
- raggiunge
- recentemente
- ridurre
- Ridotto
- riducendo
- raffinato
- insegnamento rafforzativo
- rilasciato
- rimanente
- resti
- notevole
- ricorda
- a distanza
- ripetuto
- che rappresenta
- rappresenta
- necessario
- richiede
- responsabile
- limitato
- colpevole
- risultante
- Risultati
- ritorno
- problemi
- rivoluzionare
- Premiare
- RIGA
- rt
- Correre
- s
- stesso
- Risparmi
- risparmio
- scenario
- Scenari
- schema
- Cerca
- Secondo
- Sezione
- seme
- selezionato
- Selezione
- prodotti
- AUTO
- set
- regolazione
- alcuni
- Forma
- compartecipazione
- Corti
- dovrebbero
- mostrare attraverso le sue creazioni
- mostrato
- Spettacoli
- Signal
- significativa
- significativamente
- simile
- somiglianze
- Un'espansione
- semplicità
- semplicemente
- simulazione
- da
- situazione
- Taglia
- Dimensioni
- piccole
- inferiore
- Istantanea
- So
- soluzione
- RISOLVERE
- Soluzione
- alcuni
- Fonte
- lo spazio
- spazi
- specifico
- in particolare
- specificato
- velocità
- spendere
- Spendere
- dividere
- quadrato
- impilato
- Stage
- Standard
- inizia a
- inizio
- Regione / Stato
- state-of-the-art
- ha dichiarato
- stati
- statistica
- step
- Passi
- sosta
- Tornare al suo account
- memorizzati
- negozi
- lineare
- di successo
- tale
- supporto
- supporti
- sintetico
- dati sintetici
- sinteticamente
- sistema
- SISTEMI DI TRATTAMENTO
- su misura
- Fai
- prende
- presa
- Target
- Task
- task
- tecniche
- terminal
- condizioni
- che
- I
- Il futuro
- Il grafo
- le informazioni
- Lo Stato
- loro
- Li
- in tal modo
- perciò
- Strumenti Bowman per analizzare le seguenti finiture:
- Terza
- tre
- tridimensionale
- Attraverso
- tempo
- richiede tempo
- a
- insieme
- token
- tokenizzazione
- token
- Tokens
- pure
- toolkit
- top
- torcia
- Totale
- tradizionale
- Treni
- allenato
- Training
- forma
- traiettoria
- trasformato
- trattare
- vero
- capire
- Universo
- non usato
- Aggiornanento
- aggiornato
- Aggiornamenti
- aggiornamento
- upgrade
- us
- Impiego
- uso
- utenti
- generalmente
- APPREZZIAMO
- Valori
- vario
- versione
- Video
- video games
- Visita
- Visite
- W
- Modo..
- Che
- quale
- while
- ampiamente
- wikipedia
- volere
- vincente
- con
- parole
- stazione di lavoro
- valore
- sarebbe
- scrittura
- scritto
- X
- zefiro
- zero