Fotografie de DeepMind on Unsplash
Înmulțirea matricelor este o operație fundamentală utilizată în multe sisteme, de la rețele neuronale la rutine de calcul științific. Găsirea unor algoritmi eficienți și corecti pentru înmulțirea matricelor poate avea un impact uriaș asupra procesului de calcul mai rapid și mai eficient, dar este o sarcină foarte dificilă. Spațiul de algoritmi posibili este enorm, iar metodele tradiționale de descoperire a algoritmilor, cum ar fi euristica proiectată de om sau căutarea combinatorie, sunt adesea suboptime.
DeepMind's recently proposed AI-based solution for automated search goes far beyond human intuition. The solution consists of a deep reinforcement learning agent called AlphaTensor, built on top of alphazero. Acest agent este antrenat să joace un joc pentru un singur jucător, TensorGame, în care scopul este să descopere algoritmi eficienți din punct de vedere computațional pentru multiplicarea matricei.
AlphaTensor este deosebit de bun la manipularea matricilor mari prin descompunerea înmulțirilor mari de matrice în înmulțiri mai mici. Mai mult, AlphaTensor poate fi folosit pentru a obține performanțe de ultimă generație pentru multiplicarea matricei, odată reglat fin pe un anumit dispozitiv hardware.
AlphaTensor are un mare potențial pentru accelerarea calculului de învățare profundă. În învățarea profundă, multe operațiuni consumatoare de timp pot fi mapate la înmulțiri de matrice. Prin utilizarea AlphaTensor pentru a optimiza aceste operațiuni, performanța generală a modelelor de învățare profundă poate fi îmbunătățită semnificativ.
Recent, OpenAlphaTensor, the prima implementare open source a AlphaTensor, a fost lansat, care ar putea revoluționa puterea de calcul a modelelor de învățare profundă.
Tensorul de multiplicare a matricei
Pentru cei care nu sunt experți în optimizarea înmulțirii matricelor, s-ar putea să nu fie simplu să înțeleagă cum o operație, cum ar fi multiplicarea matricei, poate fi mapată într-un tensor tridimensional. Voi încerca să explic în cuvinte simple și cu exemple.
Să considerăm produsul C = A*B, unde pentru simplitate atât A cât și B sunt matrici pătrate de dimensiunea N. Operația de înmulțire poate fi mapată într-un tensor 3D de formă (N^2, N^2, N^2). Prima dimensiune tensorală reprezintă matricea aplatizată A, a doua dimensiune matricea aplatizată B și a treia dimensiune matricea aplatizată C.
Tensorul are doar valori binare (fie 1, fie 0) pentru fiecare intrare. Rețineți că tensorul reprezintă operația de înmulțire, deci este independent de valorile matricelor A și B.
Fiecare intrare a tensorului corespunde coeficientului operației. De exemplu, pentru a calcula C[1,1], este necesar să se înmulțească atât A[1,1], cât și B[1,1]. Prin urmare, intrarea tensorului [0,0,0], care corespunde lui A[1,1], B[1,1] și C[1,1], va avea valoarea 1. În schimb, pentru a calcula C[1,1 ,2,1], A[1] nu este necesar. Astfel, rândul tensorului T[N+0, :, XNUMX] va conține doar zerouri.
Imaginea de mai jos arată un exemplu de tensor pentru N=2.
Image from DeepMind's hârtie publicat în Natură
După cum se arată în (b) și (c) în figura de mai sus, este posibil să se implementeze un algoritm pentru calcularea produsului folosind o descompunere a tensorului 3D. Mai precis, algoritmul de mai jos poate fi folosit pentru a converti o descompunere tensorală (matricele U, V, W) într-un algoritm de multiplicare a matricei.
Meta-algorithm parameterized for computing the matrix product C=AB introduced in DeepMind's hârtie
TensorGame
Problema găsirii unor algoritmi eficienți pentru multiplicarea matricei este extrem de provocatoare, deoarece numărul de algoritmi posibili de luat în considerare este mult mai mare decât numărul de atomi din univers, chiar și pentru cazurile mici de multiplicare a matricei.
DeepMind a transformat această problemă într-un joc pentru un singur jucător și l-a numit TensorGame. În acest joc, jucătorul alege cum să combine diferite intrări de matrice pentru a le multiplica. Se atribuie un punctaj în funcție de numărul de operații necesare pentru a obține rezultatul corect al înmulțirii. Jocul se termină când se atinge tensorul zero sau când s-a făcut numărul maxim de mișcări. Factorizarea finală este evaluată pe baza unei estimări a rangului rezidual și a anumitor criterii de optimizare, cum ar fi complexitatea timpului asimptotic sau timpul de rulare practic.
Poziția inițială în TensorGame corespunde Tensorului de multiplicare a matricei exprimat în mod aleatoriu.
În fiecare pas t al jocului, jucătorul notează trei vectori
, care specifică tensorii de rang-1 . Starea jocului este actualizată prin scăderea vectorilor selectați de jucător:
Unde
este tensorul de multiplicare a matricei.Dacă jocul se termină în p pași, aceasta înseamnă că Tensorul de multiplicare a matricei
pot fi descompuse în tensori p rang-1 , adică are cel puțin rangul p.TensorGame poate fi apoi interpretat ca un algoritm de descompunere a rangului, iar AlphaTensor poate fi văzut ca un algoritm pentru estimarea rangului tensorului.
Arhitectura AlphaTensor
Până acum am aflat despre TensorGame și am clarificat cum soluția sa poate fi văzută ca un algoritm de multiplicare a matricei. Să explorăm acum principalele concepte ale AlphaTensor, algoritmul folosit pentru joc.
Arhitectura AlphaTensor este practic o arhitectură Transformator codificator-decodor în care:
- codificatorul ia ca intrare starea jocului , cele n acțiuni anterioare întreprinse de model (de obicei n=7) și indicele de timp t al acțiunii curente. Informațiile sunt stivuite împreună într-un tensor cu formă (n+1, N^2, N^2, N^2). Acest tensor este apoi remodelat și transformat (folosind trei straturi liniare) într-un tensor de formă (N^2, N^2, c) unde c este dimensiunea interioară a modelului.
- decodorul generează acțiunile n_steps din vectorul încorporat dat de codificator într-un mod auto-regresiv. Fiecare acțiune corespunde unui simbol al tripleților reprezentând unul dintre tripletele care descompun tensorul jocului (adică reducerea rangului acestuia)
Modelul este antrenat prin alternarea back-propagation și model acting. Acționarea modelului este utilizată pentru a genera date care sunt apoi folosite pentru a antrena modelul. În practică, modelul este antrenat cu un amestec de date generate sintetic și date generate de model în timpul acțiunii. Pasul de acțiune se face prin luarea unui tensor 3D corespunzător unei operații matrice și jucând jocuri n_actors pe acesta. Fiecare actor joacă un joc fie pe o bază standard, fie pe o bază alternativă (schimbarea bazei se aplică cu o probabilitate dată). Rezultatele sunt apoi colectate și pot fi utilizate în etapa de antrenament cu datele sintetice.
The acting step is based on AlphaZero's Monte Carlo Tree Search (MCTS), modified to support large action spaces. In short, before choosing the action, n_sims paths are explored from the model output with a maximum future exploration of 5 steps. The probabilities generated by the model are then adjusted taking into account the generated paths. Then the action with the most promising future path(s) is chosen to continue the game.
În timpul antrenării modelului, recompensa este de fapt o recompensă negativă (penalizare). Valoarea sa absolută crește cu fiecare pas suplimentar necesar pentru rezolvarea jocului. Dacă modelul face m pași pentru a rezolva un TensorGame, recompensa asociată jocului este r=-m. Dacă modelul nu este capabil să rezolve TensorGame în pași max_rank, recompensa este calculată prin estimarea rangului tensorului rămas. Rangul este estimat ca suma rangurilor matricelor care compun tensorul. Estimarea este o limită superioară a rangului adevărat al tensorului.
When fine-tuning the model, the penalty reward at the terminal state should also take into account the latency of the algorithm produced by the model. The reward formula becomes rt'=rt+λbt, where rt is the reward scheme described earlier, bt is the benchmark reward (non-zero only at the terminal state), and λ este un coeficient specificat de utilizator.
Accelerări (%) ale algoritmilor descoperiți de AlphaTensor, adaptați pentru un GPU și un TPU, extrași din lucrarea DeepMind. Accelerările sunt măsurate în raport cu multiplicarea matricei standard (de exemplu cuBLAS pentru GPU) pe același hardware și în comparație cu Algoritmul Strassen-pătrat. sursa: DeepMind.
Am eliberat recent OpenAlphaTensor, prima implementare open source a AlphaTensor. În această secțiune voi parcurge implementarea. După cum am discutat mai devreme, arhitectura AlphaTensor este destul de simplă, bazată pe un transformator standard cu o arhitectură codificator-decodor. Cele mai interesante componente ale AlphaTensor sunt primul strat din partea codificatorului și modul în care sunt eșantionate acțiunile.
Să începem cu primul strat de codificare.
# 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
În fragmentul de mai sus, arătăm cum tensorul de intrare este descompus în trei tensori, care sunt apoi utilizați ca intrări de interogare, cheie și valoare ale stratului transformator.
- Pe cele trei dimensiuni tensoare care reprezintă matricele aplatizate (A, B, C), tensorul de intrare este aplatizat de-a lungul fiecărei dimensiuni împreună cu dimensiunea care reprezintă acțiunile anterioare. În acest fel, în fiecare copie aplatizată a tensorului de intrare, dimensiunea selectată este o agregare a ultimelor valori T-1 și a valorii reale, pentru toate valorile S ale dimensiunii selectate, unde S=N^2. Din punct de vedere filosofic, este ca și cum, pentru fiecare dimensiune, ne concentrăm asupra a ceea ce s-a întâmplat în acțiunile anterioare din acea dimensiune.
- Scalarii sunt mapați în trei spații diferite de dimensiunea S^2 și apoi remodelați pentru a fi concatenați cu tensorii obținuți la punctul anterior. Din punct de vedere conceptual, scalarii sunt mapați la un spațiu de încorporare de dimensiunea S^2, iar apoi informațiile încorporate sunt împărțite în S vectori și stivuite împreună, similar cu ceea ce se întâmplă cu textul atunci când este tokenizată.
- Tokenurile scalare sunt concatenate cu tensorul de intrare restructurat și apoi date ca intrare unui strat liniar pentru maparea informațiilor focalizate scalari+istorie canal în dimensiunea internă a modelului.
Acești trei pași pot fi interpretați ca o modalitate de a oferi modelului atât informații despre scalari (ca în pasul de timp TensorGame), cât și focalizarea pe acțiunile anterioare pentru fiecare canal.
În ceea ce privește modul în care sunt produse acțiunile, este interesant de observat că AlphaTensor generează ca ieșire tripletul u, v, w, care urmărește reducerea rangului tensorului. Cei trei vectori au dimensiunea S și, deoarece sunt concatenați, modelul trebuie să producă un vector de dimensiunea 3*S. AlphaTensor este antrenat cu un algoritm RL, deci toate acțiunile posibile trebuie exprimate în termeni de probabilități într-un spațiu enumerat, adică modelul produce o probabilitate asupra diferitelor acțiuni. Aceasta înseamnă că fiecare vector din spațiul 3S ar trebui mapat la o acțiune diferită. Rezultă un spațiu de acțiune de dimensiunea |F|^(3S), unde |F| este numărul de valori diferite pe care le poate lua elementul u, v, w. De obicei, valorile sunt limitate la (-2, -1, 0, 1, 2), rezultând o cardinalitate de 5 elemente.
Iată o provocare majoră: pentru a genera probabilitățile de acțiune pentru un produs matriceal de matrici de dimensiunea 5, am avea nevoie de o memorie de 5^75 * 4 octeți, ceea ce ar însemna ~10^44 GB de memorie. În mod clar, nu putem gestiona un spațiu de acțiune atât de mare.
Cum rezolvăm problema? Pentru a reduce amprenta de memorie a probabilităților de acțiune, putem împărți tripleții în bucăți mai mici, le putem „tokeniza” și trata bucățile ca jetoane generate în arhitectura transformatorului, adică jetoanele sunt date ca intrare la decodor într-un mod auto-regresiv. cale. În exemplul de mai sus putem împărți tripleții în 15 bucăți, reducând consumul de memorie la 15 * 5^(75/15) * 4, adică 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), )
Mai sus arătăm fragmentul de cod pentru generarea acțiunii complete. În cod, self.core conține stratul de decodor, iar tensorul e reprezintă ieșirea stratului de codificator. Zero poate fi considerat drept tokenul în modelele NLP și acțiunile n_steps care reprezintă bucățile n_steps sunt generate într-un mod progresiv.
Modelul returnează trei cantități:
- Acțiunile generate
- Probabilitatea asociată acțiunii complete
- Logiti-urile produse pentru generarea primei acțiuni (prima bucată) care vor fi utilizate pentru calcularea valorii modelului.
Merită să cheltuiți câteva cuvinte pe parametrul n_samples. Parametrul este utilizat pentru pasul de acțiune și permite modelului să genereze diferite versiuni ale tripleților care vor fi apoi utilizate pentru explorarea spațiului de acțiune în algoritmul de căutare a arborelui Monte Carlo utilizat în procesul de acțiune. Diferitele acțiuni n_samples sunt eșantionate conform politicii generate de model.
Pasul de actorie
Cea mai dificilă parte a întregului algoritm este probabil pasul de acțiune folosit pentru rezolvarea TensorGame. Algoritmul nu este explicat profund în lucrarea AlphaTensor, deoarece se bazează pe mai multe lucrări anterioare ale DeepMind, care tocmai sunt citate și date ca fiind cunoscute. Aici, voi reconstrui toate piesele lipsă și voi explica pas cu pas implementarea noastră.
Putem organiza etapele de acțiune în trei componente diferite:
- Căutarea arborelui Monte-Carlo
- Simularea jocului
- Calcularea politicii îmbunătățite
Să le analizăm unul câte unul.
Căutarea arborilor din Monte-Carlo (MCTS)
Monte Carlo Tree Search (MCTS) este o tehnică de inteligență artificială utilizată pe scară largă pentru jocuri, în special în jocurile de societate și jocurile video. Algoritmul creează un arbore de joc care simulează potențialele mișcări și rezultate și utilizează eșantionarea aleatorie pentru a evalua recompensa așteptată pentru fiecare mișcare. Algoritmul selectează apoi în mod iterativ mișcarea cu cea mai mare recompensă așteptată și simulează rezultatele până când atinge o stare terminală sau o condiție de oprire specificată. Simulările sunt folosite pentru a estima probabilitatea de câștig pentru fiecare mișcare și pentru a ghida procesul de luare a deciziilor. MCTS s-a dovedit a fi eficient în jocurile complexe în care numărul de mișcări și rezultate posibile este mare și a fost folosit în sisteme AI de succes, cum ar fi AlphaGo.
În AlphaTensor este utilizată o versiune modificată a MCTS original. În special, în loc să se selecteze aleatoriu acțiunea din întreg spațiul de acțiune, acțiunea este selectată dintr-un subset generat direct de model (prin n_samples prezentate anterior). Corecția la actualizarea politicii este apoi aplicată în pasul de calcul al politicii îmbunătățite.
În implementarea noastră, am decis să păstrăm toate informațiile despre arborele Monte-Carlo într-un dicționar având ca cheie versiunea hash a stării TensorGame și ca valori informațiile asociate stării în sine. Fiecare pas Monte-Carlo pornește de la un nod și simulează mini-jocuri n_sim, explorând viitorul cu un orizont de 5 mișcări. Dacă nodul a fost deja explorat în simulările anterioare, n_sim este ajustat luând în considerare numărul de explorări anterioare. Pentru fiecare nod, numărul de vizite este stocat în tensorul N_s_a, deoarece acest tensor conține numărul de vizite pe acțiune copil de nod (dintre cele eșantionate de model).
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
Codul de mai sus arată implementarea noastră a algoritmului. Din motive de simplitate a codului, corectarea politicii este efectuată în funcția simulate_game.
Simulare de joc
Funcția simulate_game este responsabilă pentru explorarea arborelui compus din noduri reprezentând o anumită stare a TensorGame. De asemenea, rulează modelul ori de câte ori este întâlnit un nod frunză și stochează toate informațiile despre nod în dicționarul state_dict. Să aruncăm o privire profundă asupra implementării sale:
@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)
Fiecare simulare este împărțită în trei părți:
- Selecţie
- Expansiune
- Backup
În partea de selecție, simularea este rulată pe nodurile arborescente deja generate, iar următorul nod este selectat folosind următoarea funcție:
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()]
În practică, acțiunea de maximizare a funcției ucb:
pentru starea dată este selectată. Aici Q reprezintă valorile Q generate de model și π reprezintă distribuția aleatoare a acțiunilor eșantionate folosind politica modelului. N(s, a) reprezintă numărul de vizite ale nodului la acțiunea a de la nodul s.
Odată ce faza de selecție ajunge la un nod frunză, dacă simularea nu a atins o condiție terminală (fie în ceea ce privește explorarea maximă, adică orizontul viitor, fie sfârșitul jocului), modelul este apoi utilizat pentru selectarea n_samples noduri alternative (acestea vor fi frunze). noduri în iterația succesivă). Aceasta se numește faza de expansiune, deoarece noi noduri sunt adăugate în arbore. Apoi, nu mai este explorat niciun nod în simularea curentă, dar valoarea q_value a frunzei este trimisă la următorul pas de simulare: backup.
Backup-ul este etapa finală a fiecărei simulări. În timpul copiei de rezervă, dacă nodul frunză a fost o stare terminală, recompensa finală este calculată; în caz contrar, valoarea q frunzei este folosită ca recompensă estimată. Apoi recompensa este propagată înapoi pe traiectoria simulării actualizând ambele stări q_values și actualizând contorul de vizite N(s, a). În fragmentul de mai jos arătăm codul pentru propagarea înapoi a recompensei.
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
Calcul îmbunătățit al politicilor
Odată ce toate simulările au fost executate și MCTS oferă un instantaneu interesant al viitorului apropiat, este timpul să actualizăm politica asociată cu nodurile prezise și să le returnăm, astfel încât să poată fi utilizate în timpul antrenamentului. Politica îmbunătățită, urmând metoda descrisă în Hubert şi colab, este folosit pentru gestionarea spațiilor mari de acțiune. De fapt, pentru spațiul de căutare mic, în timpul MCTS este posibil să eșantionați o acțiune în mod aleatoriu din spațiul de acțiune și să evaluați impactul acesteia. O abordare similară într-un spațiu de acțiune mult mai mare ar duce la divergerea tuturor traiectoriilor pe căi diferite și ar avea nevoie de o cantitate infinită de traiectorii pentru obținerea de statistici semnificative și apoi actualizarea politicii. Deoarece aici folosim sample-MCTS pentru a evita dispersia, adică acțiunile n_samples sunt eșantionate în conformitate cu politica modelului și apoi MCTS selectează doar una dintre acțiunile eșantionate în timp ce explorăm arborele, trebuie să luăm în considerare corecția eșantionului atunci când calculăm politica finală actualizată care va fi utilizată în timpul antrenării modelului.
În practică, politica îmbunătățită este calculată ca
Unde
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
Rețineți că în implementarea noastră, după ce am calculat politica de la tensorul N_s_a, trebuie să o mapam înapoi la tensorul de acțiune original. De fapt, N_s_a ia în considerare doar acțiunile eșantionate de model, în timp ce politica finală trebuie să conțină probabilități și pentru acțiunile neexplorate.
Diferențe în ceea ce privește algoritmul de antrenament ChatGPT
AlphaTensor is the latest member of the AlphaGo/AlphaZero family of artificial intelligence methods by DeepMind. These methods are based on the Monte Carlo Tree Search (MCTS) algorithm, which has been refined and enhanced by DeepMind to tackle increasingly complex tasks. Another AI system, OpenAI's ChatGPT, which has caused a lot of buzz for its remarkable performance, was trained with a different approach, called Reinforcement Learning with Human Feedback (RLHF).
RLHF este o tehnică de reglare fină utilizată pentru a regla modelele de limbaj pentru a urma un set de instrucțiuni scrise. Folosește preferințele umane ca semnal de recompensă pentru a ajusta modelul, aliniind astfel comportamentul modelului lingvistic cu preferințele declarate ale unui anumit grup de oameni, mai degrabă decât o noțiune mai largă de „valori umane”.
În schimb, MCTS este un algoritm de căutare bazat pe arbore, folosit pentru a determina mișcările optime în jocuri. Simulează potențialele mișcări și actualizează valorile fiecărei mișcări pe baza rezultatelor acestora, ghidând selecția celei mai bune mișcări.
RLHF colectează date din demonstrații scrise de oameni și comparații etichetate de oameni între modelele AI și antrenează un model de recompensă pentru a prezice preferințele unui anumit grup de oameni. Modelul de recompensă este apoi utilizat pentru a ajusta modelele AI. MCTS, pe de altă parte, utilizează simulări și evaluări pentru a determina cea mai bună decizie.
Deși sunt abordări diferite, RLHF și MCTS au, de asemenea, asemănări. Ambele tehnici de inteligență artificială folosesc metode de luare a deciziilor și de rezolvare a problemelor și ambele folosesc o abordare prin încercare și eroare pentru a explora diferite opțiuni și a lua decizii pe baza informațiilor disponibile. Ambele sunt, de asemenea, procese iterative care se îmbunătățesc în timp pe măsură ce se adună mai multe informații și experiență.
Alegerea între RLHF și MCTS depinde de sarcina în cauză. RLHF este ideal atunci când nu există o metrică clară pentru evaluarea performanței modelului, în timp ce MCTS s-a dovedit eficient în sarcini asemănătoare jocurilor în care cunoștințele și explorarea viitorului oferă modelului un avantaj semnificativ.
Optimizarea codului pentru antrenamentul AlphaTensor
Implementarea algoritmului de antrenament AlphaTensor necesită găsirea compromisului perfect între viteza de antrenament și consumul de memorie. După cum se vede în secțiunea Model, pur și simplu luarea în considerare a tokenizării acțiunii poate economisi multă memorie, dar o reducere prea agresivă a spațiului de acțiune poate duce atât la scăderea preciziei, cât și la o performanță mai lentă. Acesta din urmă se întâmplă deoarece toate jetoanele sunt generate secvenţial într-un mod autoregresiv de către decodorul modelului. Prin urmare, timpul de inferență crește liniar cu numărul de jetoane pe acțiune, odată ce softmax-ul de pe spațiul de acțiune nu mai este blocajul.
La configurarea antrenamentului AlphaTensor, principalele dificultăți au fost găsite în abordarea procesului de actorie. Dacă tensorii nu sunt stocați în formatul corect, MCTS poate provoca cu ușurință creșterea necontrolată a utilizării memoriei. Pe de altă parte, dacă numărul de tensori stocați în timpul fiecărei simulări este redus prea mult, MCTS poate petrece o perioadă infinită de timp recalculând stările necesare.
Let's take an example of the game simulation step, where the game is explored by looking at possible future scenarios. For each state, if we don't save the actions generated by the model and we decide to save only the random seed used to sample the actions from the policy, then each time we explore a tree node we would have to recompute the policy and then sample the actions. Clearly, we decided to store the sampled actions to save time and to avoid having to manage model sharing between different processes in the case of MCTS exploration parallelization. However, just saving the actions was not enough to get a sufficiently efficient acting step. In fact, the time for converting the n_steps actions into the (u, v, w) triplet, reducing the game tensor state and creating the new3D tensors from the n_samples actions would easily be a bottleneck for the whole training. Secondly, we didn't want to store all possible future states for each sampled action, as this would have a huge impact on the memory used by the algorithm. Suppose we set n_samples=32, n=7 and N=5, and let's remember that N is the size of the square matrix product we want to reduce and n is the number of previous actions remembered by the model. In this situation, each state tensor would have the form (8, 25, 25, 25), which multiplied by 32 would result in 3282525254 octeți pentru fiecare nod din grafic. Acum, având în vedere că fiecare simulare din faza de expansiune generează un nou nod (și n_sim=200), am avea un consum final de memorie de 200328252525*4 = 3.2 GB numai pentru primul nod MCTS. În cel mai rău caz, în timp ce se explorează nodurile max_rank active (unde max_rank=150), acest lucru ar duce la un consum total de memorie de 150 * 3.2 GB = 480 GB în memoria RAM (sau memoria GPU dacă toți tensorii au fost stocați pe GPU) . Am desfășurat antrenamentul pe stația noastră de lucru cu 128 GB RAM și 48 GB memorie GPU, așa că a trebuit să reducem consumul de memorie.
Since we didn't want to increase the execution time, we adopted an optimization that exploits the redundancy in the state tensors produced. In fact, the tensors have n-1 previous actions in common, which can then be stored once and not repeated for each stored tensor. This results in a memory reduction of 2/7~28%, meaning that in the worst-case 137GB can be stored. At this point, by simply pruning the unused part of the tree (such as the unselected trajectories) and storing the tensors in CPU memory, we were able to avoid any memory error during training.
Având în vedere că OpenAlphaTensor este acum open source, se deschid câteva căi interesante pentru dezvoltarea ulterioară.
O progresie naturală este reglarea fină a OpenAlphaTensor pe dispozitivele hardware țintă. Acest lucru este de așteptat să conducă la o performanță de calcul foarte competitivă. Voi publica mai multe despre performanța OpenAlphaTensor pe diverse hardware pe GitHub. La momentul scrierii acestui articol, OpenAlphaTensor era în curs de formare.
Un alt progres important ar fi suportul pentru compilarea de la distanță, permițând utilizatorilor să construiască algoritmi optimizați pentru dispozitivele edge. Acest lucru poate fi realizat prin stocarea modelului OpenAlphaTensor pe un server, în timp ce algoritmul de multiplicare a matricei este evaluat pe hardware diferit.
De asemenea, ar putea fi importantă extinderea suportului pentru diferiți compilatori pentru a calcula corecția recompensei bazată pe latență. Diferiți compilatori pot duce la diferiți algoritmi optimizați pe un anumit hardware. De exemplu, lucrarea DeepMind a arătat rezultate promițătoare folosind JAX și compilatorul XLA pe TPU și GPU-uri Nvidia. Ar fi interesant să evaluăm acest lucru folosind NCCL pe Nvidia sau LLVM pe procesoare.
În cele din urmă, extinderea modelului și a algoritmului de antrenament pentru a suporta dimensiuni mai mari ale matricei rămâne o provocare majoră deschisă. În prezent, OpenAlphaTensor acceptă o dimensiune maximă a matricei de 5, dar poate fi aplicată prin împărțirea înmulțirilor de matrice mai mari în grupuri de MM-uri mici cu o dimensiune mai mică de 5. Această abordare este suboptimă și efectuând reducerea direct pe tensorul mare corespunzător MM complet ar putea duce teoretic la rezultate mai bune.
Diego Fiori is the CTO of Nebuly AI, a company committed to making AI optimization part of every developer's toolkit.
- Distribuție de conținut bazat pe SEO și PR. Amplifică-te astăzi.
- Platoblockchain. Web3 Metaverse Intelligence. Cunoștințe amplificate. Accesați Aici.
- Sursa: 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
- :este
- ][p
- $UP
- 1
- 3d
- 8
- a
- Capabil
- Despre Noi
- mai sus
- Absolut
- accelerarea
- Conform
- în consecință
- Cont
- precizie
- Obține
- realizat
- Acțiune
- acțiuni
- de fapt
- adăugat
- Suplimentar
- Ajustat
- adoptată
- avansa
- Avantaj
- După
- Agent
- agregare
- agresiv
- AI
- Sisteme AI
- isi propune
- Algoritmul
- algoritmi
- TOATE
- Permiterea
- permite
- singur
- deja
- alternativă
- printre
- sumă
- analiza
- și
- O alta
- aplicat
- abordare
- abordari
- arhitectură
- SUNT
- articol
- artificial
- inteligență artificială
- AS
- alocate
- asociate
- At
- Automata
- disponibil
- evitarea
- înapoi
- Backup
- bazat
- Pe scurt
- bază
- BE
- deoarece
- devine
- înainte
- fiind
- de mai jos
- Benchmark
- CEL MAI BUN
- Mai bine
- între
- Dincolo de
- bord
- Consiliul de Jocuri
- legat
- mai larg
- BT
- construi
- construit
- by
- denumit
- CAN
- nu poti
- caz
- Provoca
- cauzată
- sigur
- contesta
- provocare
- Schimbare
- Canal
- Chat GPT
- copil
- alegere
- alegere
- ales
- citată
- clar
- clar
- cod
- colecte
- combina
- comise
- Comun
- companie
- comparație
- competitiv
- complex
- complexitate
- componente
- compuse
- compromis
- calcul
- puterea de calcul
- Calcula
- tehnica de calcul
- Concepte
- conceptual
- condiție
- încredere
- Lua în considerare
- luate în considerare
- luand in considerare
- consideră
- consum
- conține
- continua
- contrast
- convertit
- Nucleu
- Corespunzător
- corespunde
- ar putea
- Contracara
- Procesor
- creează
- Crearea
- Criteriile de
- CTO
- Curent
- În prezent
- de date
- abuzive
- decide
- hotărât
- decizie
- Luarea deciziilor
- Deciziile
- adânc
- învățare profundă
- DeepMind
- depinde de
- descris
- Determina
- Dezvoltator
- Dezvoltare
- dispozitiv
- Dispozitive
- DICT
- diferit
- dificultăți
- Dimensiune
- Dimensiuni
- direct
- descoperi
- descoperirea
- discutat
- distribuire
- împărțit
- jos
- Picătură
- în timpul
- e
- fiecare
- Mai devreme
- cu ușurință
- Margine
- Eficace
- eficient
- oricare
- element
- element
- încorporat
- se încheie
- sporită
- enorm
- suficient de
- intrare
- eroare
- estima
- estimativ
- Eter (ETH)
- evalua
- evaluat
- evaluarea
- evaluări
- Chiar
- Fiecare
- exemplu
- exemple
- captivant
- execuție
- expansiune
- de aşteptat
- experienţă
- Explica
- a explicat
- exploit
- explorare
- explora
- explorat
- Explorarea
- și-a exprimat
- extinde
- extindere
- extrem
- destul de
- familie
- mai repede
- feedback-ul
- puțini
- Figura
- final
- descoperire
- First
- pluti
- Concentra
- urma
- următor
- urmă
- Pentru
- formă
- format
- formulă
- găsit
- din
- Complet
- funcţie
- fundamental
- mai mult
- dezvoltare ulterioară
- viitor
- joc
- Jocuri
- genera
- generată
- generează
- generator
- obține
- obtinerea
- Da
- dat
- Oferirea
- scop
- Merge
- bine
- GPU
- unități de procesare grafică
- grafic
- mare
- grup
- Grupului
- creste
- Creștere
- ghida
- mână
- Manipularea
- sa întâmplat
- se întâmplă
- Piese metalice
- dispozitiv hardware
- dispozitive hardware
- Avea
- având în
- aici
- cea mai mare
- orizont
- Cum
- Cum Pentru a
- Totuși
- HTTPS
- mare
- uman
- i
- BOLNAV
- ideal
- idx
- imagine
- Impactul
- punerea în aplicare a
- implementarea
- important
- îmbunătăţi
- îmbunătățit
- in
- Crește
- Creșteri
- tot mai mult
- independent
- index
- informații
- inițială
- intrare
- in schimb
- instrucțiuni
- Inteligență
- interesant
- intern
- introdus
- intuiţie
- IT
- repetare
- ESTE
- în sine
- jpg
- KDnuggets
- A pastra
- Cheie
- cunoştinţe
- cunoscut
- limbă
- mare
- mai mare
- Nume
- Latență
- Ultimele
- strat
- straturi
- conduce
- învățat
- învăţare
- Listă
- Uite
- cautati
- Lot
- făcut
- Principal
- major
- face
- Efectuarea
- administra
- de conducere
- multe
- Hartă
- cartografiere
- Matrice
- materie
- maxim
- sens
- semnificativ
- mijloace
- membru
- Memorie
- metodă
- Metode
- metric
- dispărut
- amestec
- model
- Modele
- modificată
- modul
- mai mult
- mai eficient
- În plus
- cele mai multe
- muta
- mişcă
- înmulțit
- Natural
- Natură
- În apropiere
- necesar
- Nevoie
- necesar
- negativ
- rețele
- neural
- rețele neuronale
- Nou
- următor
- nlp
- nod
- noduri
- neexperti
- noțiune
- număr
- Nvidia
- obținut
- of
- promoții
- on
- ONE
- deschide
- open-source
- OpenAI
- operaţie
- Operațiuni
- optimă
- optimizare
- Optimizați
- optimizate
- Opţiuni
- original
- Altele
- in caz contrar
- producție
- global
- Hârtie
- lucrări
- parametru
- parte
- special
- în special
- piese
- oameni
- Perfect
- performanță
- efectuarea
- fază
- piese
- Plato
- Informații despre date Platon
- PlatoData
- Joaca
- player
- joc
- Punct
- Politicile
- Politica
- poziţie
- posibil
- potenţial
- putere
- Practic
- practică
- prezice
- a prezis
- preferinţele
- prezentat
- precedent
- probabilitate
- probabil
- Problemă
- proces
- procese
- produce
- Produs
- Produs
- progresie
- progresiv
- promițător
- propus
- dovedit
- dovedit
- publica
- publicat
- RAM
- aleator
- rândurile
- mai degraba
- atins
- aTINGE
- recent
- reduce
- Redus
- reducerea
- rafinat
- Consolidarea învățării
- eliberat
- rămas
- rămășițe
- remarcabil
- minte
- la distanta
- repetat
- reprezentând
- reprezintă
- necesar
- Necesită
- responsabil
- limitat
- rezultat
- rezultând
- REZULTATE
- reveni
- Returnează
- revoluţiona
- Răsplăti
- RÂND
- rt
- Alerga
- s
- acelaşi
- Economisiți
- economisire
- scenariu
- scenarii
- schemă
- Caută
- Al doilea
- Secțiune
- sămânţă
- selectate
- selectarea
- selecţie
- SELF
- set
- instalare
- câteva
- Modela
- partajarea
- Pantaloni scurți
- să
- Arăta
- indicat
- Emisiuni
- Semnal
- semnificativ
- semnificativ
- asemănător
- asemănări
- simplu
- simplitate
- pur şi simplu
- simulare
- întrucât
- situație
- Mărimea
- dimensiuni
- mic
- mai mici
- Instantaneu
- So
- soluţie
- REZOLVAREA
- Rezolvarea
- unele
- Sursă
- Spaţiu
- spații
- specific
- specific
- specificată
- viteză
- petrece
- Cheltuire
- împărţi
- pătrat
- stivuite
- Etapă
- standard
- Începe
- începe
- Stat
- de ultimă oră
- stabilit
- Statele
- statistică
- Pas
- paşi
- oprire
- stoca
- stocate
- magazine
- simplu
- de succes
- astfel de
- a sustine
- Sprijină
- sintetic
- date sintetice
- sintetic
- sistem
- sisteme
- adaptate
- Lua
- ia
- luare
- Ţintă
- Sarcină
- sarcini
- tehnici de
- Terminal
- termeni
- acea
- Viitorul
- Graficul
- informațiile
- Statul
- lor
- Lor
- astfel
- prin urmare
- Acestea
- Al treilea
- trei
- tri-dimensională
- Prin
- timp
- consumă timp
- la
- împreună
- semn
- tokenizarea
- cuvinte pot
- indicativele
- de asemenea
- Toolkit
- top
- lanternă
- Total
- tradiţional
- Tren
- dresat
- Pregătire
- trenuri
- traiectorie
- transformat
- trata
- adevărat
- înţelege
- Univers
- nefolosit
- Actualizează
- actualizat
- actualizări
- actualizarea
- upgrade-ul
- us
- Folosire
- utilizare
- utilizatorii
- obișnuit
- valoare
- Valori
- diverse
- versiune
- Video
- jocuri video
- Vizita
- Vizite
- W
- Cale..
- Ce
- care
- în timp ce
- pe larg
- Wikipedia
- voi
- câștigător
- cu
- cuvinte
- Statie de lucru
- valoare
- ar
- scris
- scris
- X
- zephyrnet
- zero