Photo by DeepMind on Unsplash
Matrix multiplikation er en grundlæggende operation, der bruges i mange systemer, fra neurale netværk til videnskabelige computerrutiner. At finde effektive og beviseligt korrekte algoritmer til matrixmultiplikation kan have en enorm indflydelse på at gøre beregningen hurtigere og mere effektiv, men det er en meget udfordrende opgave. Rummet af mulige algoritmer er enormt, og traditionelle metoder til at opdage algoritmer, såsom menneskedesignede heuristik eller kombinatorisk søgning, er ofte suboptimale.
DeepMind's nyligt foreslåede AI-baserede løsning til automatiseret søgning går langt ud over menneskelig intuition. Løsningen består af en deep reinforcement learning agent kaldet AlphaTensor, bygget ovenpå Alpha Zero. Denne agent er trænet til at spille et enkeltspillerspil, TensorGame, hvor målet er at opdage beregningseffektive algoritmer til matrixmultiplikation.
AlphaTensor er særlig god til at håndtere store matricer ved at dekomponere store matrixmultiplikationer til mindre multiplikationer. Desuden kan AlphaTensor bruges til at opnå den nyeste ydeevne til matrixmultiplikation, når den er finjusteret på en specifik hardwareenhed.
AlphaTensor har et stort potentiale for at accelerere deep learning computing. I deep learning kan mange tidskrævende operationer kortlægges til matrixmultiplikationer. Ved at bruge AlphaTensor til at optimere disse operationer, kan den overordnede ydeevne af deep learning-modeller forbedres væsentligt.
For nylig har OpenAlphaTensor, den første open source-implementering af AlphaTensor, blev udgivet, hvilket kunne revolutionere beregningskraften i deep learning-modeller.
Matrix multiplikationstensor
For ikke-eksperter i matrixmultiplikationsoptimering er det måske ikke ligetil at forstå, hvordan en operation såsom matrixmultiplikation kan kortlægges i en tredimensionel tensor. Jeg vil forsøge at forklare det med enkle ord og med eksempler.
Lad os betragte produktet C = A*B, hvor både A og B for nemheds skyld er kvadratiske matricer af størrelse N. Multiplikationsoperationen kan afbildes i en 3D-tensor af form (N^2, N^2, N^2). Den første tensordimension repræsenterer den fladtrykte matrix A, den anden dimension den fladtrykte matrix B og den tredje dimension den fladtrykte matrix C.
Tensoren har kun binære værdier (enten 1 eller 0) for hver indgang. Bemærk, at tensoren repræsenterer multiplikationsoperationen, så den er uafhængig af værdierne af matricerne A og B.
Hver indtastning af tensoren svarer til operationens koefficient. For at beregne C[1,1] er det for eksempel nødvendigt at gange både A[1,1] og B[1,1]. Derfor vil tensorindgangen [0,0,0], som svarer til A[1,1], B[1,1] og C[1,1], have værdien 1. I modsætning hertil, for at beregne C[1,1 ,2,1], A[1] er ikke nødvendig. Tensorrækken T[N+0, :, XNUMX] vil således kun indeholde nuller.
Billedet nedenfor viser et eksempel på en tensor for N=2.
Billede fra DeepMind's papir offentliggjort i Natur
Som vist i (b) og (c) i figuren ovenfor, er det muligt at implementere en algoritme til at beregne produktet ved hjælp af en dekomponering af 3D-tensoren. Mere specifikt kan algoritmen nedenfor bruges til at konvertere en tensordekomponering (matricerne U, V, W) til en matrixmultiplikationsalgoritme.
Metaalgoritme parametriseret til beregning af matrixproduktet C=AB introduceret i DeepMind's papir
Tensor-spillet
Problemet med at finde effektive algoritmer til matrixmultiplikation er ekstremt udfordrende, fordi antallet af mulige algoritmer at overveje er meget større end antallet af atomer i universet, selv for små tilfælde af matrixmultiplikation.
DeepMind konverterede dette problem til et singleplayer-spil og kaldte det TensorGame. I dette spil vælger spilleren, hvordan man kombinerer forskellige indgange af matricer for at gange dem. En score tildeles baseret på antallet af operationer, der kræves for at opnå det korrekte multiplikationsresultat. Spillet slutter, når nultensoren er nået, eller når det maksimale antal træk er foretaget. Den endelige faktorisering evalueres ud fra en estimering af den resterende rang og visse optimeringskriterier, såsom asymptotisk tidskompleksitet eller praktisk kørselstid.
Startpositionen i TensorGame svarer til Matrix Multiplikation Tensor udtrykt på et tilfældigt grundlag.
I hvert trin t i spillet skriver spilleren tre vektorer ned
, som specificerer rang-1 tensorerne . Spillets tilstand opdateres ved at trække de vektorer, som spilleren har valgt:
hvor
er matrix multiplikationstensoren.Hvis spillet ender i p trin, betyder det, at Matrix Multiplication Tensor
kan dekomponeres til p rank-1 tensorer , dvs. den har mindst rang p.TensorGame kan derefter tolkes som en rang-dekomponeringsalgoritme, og AlphaTensor kan ses som en algoritme til at estimere rangen af tensoren.
AlphaTensor-arkitektur
Indtil videre har vi lært om TensorGame og afklaret, hvordan dets løsning kan ses som en matrixmultiplikationsalgoritme. Lad os nu udforske hovedbegreberne i AlphaTensor, algoritmen, der bruges til spillet.
AlphaTensor-arkitektur er dybest set en encoder-decoder Transformer-arkitektur, hvor:
- indkoderen tager spiltilstanden som input , de n tidligere handlinger udført af modellen (normalt n=7), og tidsindekset t for den aktuelle handling. Information er stablet sammen i en tensor med form (n+1, N^2, N^2, N^2). Denne tensor omformes derefter og transformeres (ved hjælp af tre lineære lag) i en tensor af form (N^2, N^2, c), hvor c er den indre dimension af modellen.
- dekoderen genererer n_steps handlingerne fra den indlejrede vektor givet af indkoderen på en autoregressiv måde. Hver handling svarer til et token af trillingerne repræsenterer en af trillingerne, der nedbryder spiltensoren (dvs. reducerer dens rang)
Modellen trænes ved at skiftevis rygudbredelse og modelagere. Model acting bruges til at generere data, som derefter bruges til at træne modellen. I praksis trænes modellen med en blanding af syntetisk genererede data og data genereret af modellen under skuespil. Skuespiltrinnet udføres ved at tage en 3D-tensor svarende til en matrixoperation og spille n_actors-spil på den. Hver skuespiller spiller et spil enten på standardbasis eller på et alternativt grundlag (ændringen af grundlag anvendes med en given sandsynlighed). Resultaterne indsamles derefter og kan bruges i træningstrinnet med de syntetiske data.
Skuespiltrinnet er baseret på AlphaZero's Monte Carlo Tree Search (MCTS), modificeret til at understøtte store handlingsrum. Kort sagt, før handlingen vælges, udforskes n_sims-stier fra modeloutputtet med en maksimal fremtidig udforskning på 5 trin. De sandsynligheder, der genereres af modellen, justeres derefter under hensyntagen til de genererede stier. Derefter vælges handlingen med den/de mest lovende fremtidige vej(e) for at fortsætte spillet.
Mens man træner modellen, er belønningen faktisk en negativ belønning (straf). Dens absolutte værdi stiger med hvert ekstra trin, der kræves for at løse spillet. Hvis modellen tager m trin for at løse et TensorGame, er belønningen forbundet med spillet r=-m. Hvis modellen ikke er i stand til at løse TensorGame i max_rank-trin, beregnes belønningen ved at estimere rangeringen af den resterende tensor. Rangen estimeres som summen af rækkerne af de matricer, der udgør tensoren. Estimatet er en øvre grænse for tensorens sande rang.
Når du finjusterer modellen, bør strafbelønningen i terminaltilstanden også tage højde for latensen af algoritmen produceret af modellen. Belønningsformlen bliver rt'=rt+λbt, hvor rt er belønningsordningen beskrevet tidligere, bt er benchmarkbelønningen (kun ikke-nul i terminaltilstanden), og λ er en brugerspecificeret koefficient.
Speed-ups (%) af AlphaTensor-opdagede algoritmer skræddersyet til en GPU og en TPU, hentet fra DeepMinds papir. Speed-ups måles i forhold til standard (f.eks. cuBLAS for GPU) matrixmultiplikation på den samme hardware og sammenlignes med Strassen-kvadrat-algoritme. kilde: DeepMind.
jeg for nylig udgivet OpenAlphaTensor, den første open source-implementering af AlphaTensor. I dette afsnit vil jeg gennemgå implementeringen. Som vi diskuterede tidligere, er AlphaTensor-arkitekturen ret ligetil, baseret på en standardtransformer med en encoder-dekoder-arkitektur. De mest interessante komponenter i AlphaTensor er det første lag i encoder-delen og måden, hvorpå handlingerne samples.
Lad os starte med det første kodningslag.
# 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
I uddraget ovenfor viser vi, hvordan inputtensoren dekomponeres i tre tensorer, som derefter bruges som forespørgsels-, nøgle- og værdiinput af transformerlaget.
- På tværs af de tre tensordimensioner, der repræsenterer de fladtrykte matricer (A, B, C), fladgøres inputtensoren langs hver dimension sammen med den dimension, der repræsenterer de foregående handlinger. På denne måde er den valgte dimension i hver fladtrykt kopi af inputtensoren en aggregering af de sidste T-1-værdier og den faktiske værdi for alle S-værdierne af den valgte dimension, hvor S=N^2. Filosofisk er det, som om vi for hver dimension fokuserer på, hvad der skete i de tidligere handlinger i den dimension.
- Skalarerne kortlægges i tre forskellige rum med dimension S^2 og omformes derefter til at blive sammenkædet med tensorerne opnået ved det foregående punkt. Konceptuelt kortlægges skalarerne til et indlejringsrum med dimension S^2, og derefter stykkes den indlejrede information i S-vektorer og stables sammen, svarende til hvad der sker med tekst, når den tokeniseres.
- Skalære tokens sammenkædes med den omstrukturerede inputtensor og gives derefter som input til et lineært lag til kortlægning af skalarer+kanalhistoriens fokusinformation i modellens interne dimension.
Disse tre trin kan tolkes som en måde at give modellen både information om skalarerne (som i TensorGame-tidstrinnet) og fokus på de tidligere handlinger for hver kanal.
Med hensyn til måden handlingerne produceres på, er det interessant at bemærke, at AlphaTensor genererer tripletten u, v, w som output, som har til formål at reducere tensorrangen. De tre vektorer har størrelse S, og da de er sammenkædet, skal modellen producere en vektor på størrelse 3*S. AlphaTensor er trænet med en RL-algoritme, så alle mulige handlinger skal udtrykkes i form af sandsynligheder i et opregnet rum, dvs. modellen producerer en sandsynlighed over de forskellige handlinger. Dette betyder, at hver vektor i 3S-rummet skal kortlægges til en anden handling. Dette resulterer i et handlingsrum med størrelsen |F|^(3S), hvor |F| er antallet af forskellige værdier, som elementet af u, v, w kan tage. Normalt er værdierne begrænset til (-2, -1, 0, 1, 2), hvilket resulterer i en kardinalitet på 5 elementer.
Her kommer en stor udfordring: For at generere handlingssandsynlighederne for et matrixprodukt af matricer af størrelse 5 ville vi have brug for en hukommelse på 5^75 * 4 bytes, hvilket ville betyde ~10^44 GB hukommelse. Det er klart, at vi ikke kan styre så stort et handlingsrum.
Hvordan løser vi problemet? For at reducere handlingssandsynlighedens hukommelse kan vi opdele trillingerne i mindre bidder, "tokenisere" dem og behandle bidderne som genererede tokens i transformatorarkitekturen, dvs. tokens gives som input til dekoderen i en autoregressiv vej. I eksemplet ovenfor kan vi opdele trillingerne i 15 bidder, hvilket reducerer hukommelsesforbruget til 15 * 5^(75/15) * 4, dvs. 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), )
Ovenfor viser vi kodestykket til generering af den fulde handling. I koden indeholder self.core dekoderlaget, og tensoren e repræsenterer outputtet fra encoderlaget. Nul kan betragtes som token i NLP-modeller og n_steps-handlingerne, der repræsenterer n_steps-bidderne, genereres på en progressiv måde.
Modellen returnerer tre mængder:
- De genererede handlinger
- Sandsynligheden forbundet med den fulde handling
- De logits, der er produceret til at generere den første handling (den første del), der vil blive brugt til at beregne modelværdien.
Det er værd at bruge et par ord på parameteren n_samples. Parameteren bruges til skuespiltrinnet, og det giver modellen mulighed for at generere forskellige versioner af trillingerne, som derefter vil blive brugt til at udforske handlingsrummet i Monte Carlo Tree Search-algoritmen, der bruges i skuespilprocessen. De n_samples forskellige handlinger samples i henhold til den politik, der genereres af modellen.
Handler Trin
Den mest besværlige del af hele algoritmen er sandsynligvis det handletrin, der bruges til at løse TensorGame. Algoritmen er ikke dybt forklaret i AlphaTensor-papiret, da den er baseret på flere DeepMinds tidligere artikler, som netop er citeret og givet som kendt. Her vil jeg rekonstruere alle de manglende brikker og trin for trin forklare vores implementering.
Vi kan organisere skuespiltrinene i tre forskellige komponenter:
- Monte-Carlo træsøgning
- Spillets simulering
- Den forbedrede politikberegning
Lad os analysere dem én efter én.
Monte-Carlo Tree Search (MCTS)
Monte Carlo Tree Search (MCTS) er en meget brugt kunstig intelligens-teknik til spil, især i brætspil og videospil. Algoritmen opretter et spiltræ, der simulerer potentielle træk og resultater og bruger tilfældige stikprøver til at evaluere den forventede belønning for hvert træk. Algoritmen vælger derefter iterativt træk med den højeste forventede belønning og simulerer resultater, indtil den når en terminal tilstand eller en specificeret standsningstilstand. Simuleringerne bruges til at estimere sandsynligheden for at vinde for hvert træk og guide beslutningsprocessen. MCTS har vist sig at være effektiv i komplekse spil, hvor antallet af mulige træk og udfald er stort, og det er blevet brugt i vellykkede game-playing AI-systemer, såsom AlphaGo.
I AlphaTensor bruges en modificeret version af den originale MCTS. Især i stedet for tilfældigt at vælge handlingen fra hele handlingsrummet, vælges handlingen blandt en delmængde, der genereres direkte af modellen (gennem n_samples præsenteret før). Korrektionen til politikopgraderingen anvendes derefter i trinnet forbedret politikberegning.
I vores implementering besluttede vi at opbevare alle oplysninger om Monte-Carlo-træet i en ordbog med hash-versionen af TensorGame-tilstanden som nøgle og som værdier den information, der er forbundet med selve staten. Hvert Monte-Carlo-trin starter fra en node og simulerer n_sim-minispil og udforsker fremtiden med en horisont på 5 træk. Hvis noden allerede er blevet udforsket i tidligere simuleringer, justeres n_sim i betragtning af antallet af tidligere udforskninger. For hver knude lagres antallet af besøg i N_s_a-tensoren, da denne tensor indeholder antallet af besøg pr. knudeunderordnet handling (blandt dem, der er samplet af modellen).
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
Ovenstående kode viser vores implementering af algoritmen. For et spørgsmål om kodes enkelthed udføres politikkorrektionen i simulate_game-funktionen.
Spil Simulation
Funktionen simulate_game er ansvarlig for at udforske træet, der består af noder, der repræsenterer en bestemt tilstand af TensorGame. Den kører også modellen, hver gang der stødes på en bladknude, og den gemmer al nodeinformation i state_dict-ordbogen. Lad os give et dybt kig på implementeringen:
@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)
Hver simulering er opdelt i tre dele:
- Udvælgelse
- Ekspansion
- backup
I udvælgelsesdelen køres simuleringen på de allerede genererede træ-noder, og følgende node vælges ved hjælp af følgende funktion:
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()]
I praksis er handlingen, der maksimerer ucb-funktionen:
for den givne tilstand er valgt. Her repræsenterer Q Q-værdierne genereret af modellen, og π repræsenterer den tilfældige fordeling over de handlinger, der er samplet ved hjælp af modelpolitikken. N(s, a) repræsenterer antallet af besøg af noden til handling a fra node s.
Når først udvælgelsesfasen når en bladknude, hvis simuleringen ikke har nået en terminal tilstand (med hensyn til enten maksimal udforskning, dvs. fremtidig horisont, eller spilafslutning), bruges modellen derefter til at vælge n_samples alternative noder (de vil være blade) noder i den efterfølgende iteration). Dette kaldes ekspansionsfasen, da der tilføjes nye noder til træet. Derefter udforskes ingen yderligere node i den aktuelle simulering, men leaf q_value sendes til følgende simuleringstrin: backup.
Sikkerhedskopiering er den sidste fase af hver simulering. Under backup, hvis bladknuden var en terminal tilstand, beregnes den endelige belønning; ellers bruges blad q værdien som en estimeret belønning. Derefter forplantes belønningen tilbage på simuleringsbanen, idet både tilstandene q_values opdateres og besøgstælleren N(s, a) opdateres. I uddraget nedenfor viser vi koden for belønningens tilbage-udbredelse.
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
Forbedret politikberegning
Når alle simuleringer er blevet kørt, og MCTS giver et interessant øjebliksbillede af den nærmeste fremtid, er det tid til at opdatere politikken forbundet med de forudsagte noder og returnere dem, så de kan bruges under træning. Den forbedrede politik, efter metoden beskrevet i Hubert et al, bruges til at administrere store handlingsrum. Faktisk, for lille søgeplads, er det muligt under MCTS at sample en handling tilfældigt fra handlingsrummet og evaluere dens virkning. En lignende tilgang i et meget større handlingsrum ville føre til, at alle baner divergerer ad forskellige veje, og det ville kræve en uendelig mængde af baner for at få meningsfulde statistikker og derefter opdatere politikken. Da vi her bruger sample-MCTS til at undgå spredningen, dvs. n_samples-handlinger samples i overensstemmelse med modelpolitikken, og så vælger MCTS bare en af samplede handlinger, mens vi udforsker træet, skal vi tage højde for sample-korrektionen, når vi beregner den endelige opdaterede politik, der vil blive brugt under træning af modellen.
I praksis opgøres den forbedrede politik som
hvor
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
Bemærk, at i vores implementering efter at have beregnet politikken fra N_s_a-tensoren, skal vi kortlægge den tilbage til den oprindelige handlingstensor. Faktisk betragter N_s_a blot handlingerne samplet af modellen, mens den endelige politik skal indeholde sandsynligheder også for de ikke-udforskede handlinger.
Forskelle med hensyn til ChatGPT træningsalgoritme
AlphaTensor er det seneste medlem af AlphaGo/AlphaZero-familien af kunstig intelligens-metoder fra DeepMind. Disse metoder er baseret på Monte Carlo Tree Search (MCTS) algoritmen, som er blevet forfinet og forbedret af DeepMind til at tackle stadig mere komplekse opgaver. Et andet kunstig intelligens-system, OpenAIs ChatGPT, som har skabt en masse buzz for sin bemærkelsesværdige ydeevne, blev trænet med en anden tilgang, kaldet Reinforcement Learning with Human Feedback (RLHF).
RLHF er en finjusteringsteknik, der bruges til at tune sprogmodeller til at følge et sæt skriftlige instruktioner. Den bruger menneskelige præferencer som et belønningssignal til at finjustere modellen og derved afstemme sprogmodellens adfærd med de erklærede præferencer for en specifik gruppe mennesker, snarere end en bredere forestilling om 'menneskelige værdier'.
I modsætning hertil er MCTS en træbaseret søgealgoritme, der bruges til at bestemme de optimale træk i spil. Den simulerer potentielle træk og opdaterer værdierne for hvert træk baseret på deres resultater, og styrer valget af det bedste træk.
RLHF indsamler data fra menneskeskrevne demonstrationer og menneskemærkede sammenligninger mellem AI-modeller og træner en belønningsmodel til at forudsige præferencerne for en given gruppe mennesker. Belønningsmodellen bruges derefter til at finjustere AI-modellerne. MCTS, på den anden side, bruger simuleringer og evalueringer til at bestemme den bedste beslutning.
Selvom de er forskellige tilgange, har RLHF og MCTS også ligheder. Begge teknikker til kunstig intelligens bruger beslutningstagning og problemløsningsmetoder, og begge bruger en trial-and-error-tilgang til at udforske forskellige muligheder og træffe beslutninger baseret på tilgængelig information. Begge er også iterative processer, der forbedres over tid, efterhånden som der indsamles mere information og erfaring.
Valget mellem RLHF og MCTS afhænger af opgaven. RLHF er ideel, når der ikke er en klar metrik til at evaluere modellens ydeevne, mens MCTS har vist sig effektiv i spillignende opgaver, hvor viden og udforskning af fremtiden giver modellen en væsentlig fordel.
Kodeoptimering til AlphaTensor træning
Implementering af AlphaTensor træningsalgoritmen kræver at finde det perfekte kompromis mellem træningshastighed og hukommelsesforbrug. Som det ses i Model-sektionen, kan blot overvejelse af handlingstokenisering spare en masse hukommelse, men en alt for aggressiv handlingspladsreduktion kan føre til både fald i nøjagtighed og langsommere ydeevne. Sidstnævnte sker, fordi alle tokens genereres sekventielt på en autoregressiv måde af modeldekoderen. Derfor vokser inferenstiden lineært med antallet af tokens pr. handling, når softmax på handlingsrummet ikke længere er flaskehalsen.
Ved opsætning af AlphaTensor-træning fandt man de største vanskeligheder med at håndtere skuespilprocessen. Hvis tensorerne ikke er gemt i det korrekte format, kan MCTS let forårsage ukontrolleret vækst i hukommelsesforbruget. På den anden side, hvis antallet af tensorer, der er lagret under hver simulering, reduceres for meget, kan MCTS bruge uendelig meget tid på at genberegne de nødvendige tilstande.
Lad os tage et eksempel på spilsimuleringstrinnet, hvor spillet udforskes ved at se på mulige fremtidige scenarier. For hver stat, hvis vi ikke gemmer handlingerne genereret af modellen, og vi beslutter kun at gemme det tilfældige frø, der blev brugt til at prøve handlingerne fra politikken, så skal vi, hver gang vi udforsker en træknude, genberegne politikken og Prøv derefter handlingerne. Det er klart, at vi besluttede at gemme de samplede handlinger for at spare tid og for at undgå at skulle styre modeldeling mellem forskellige processer i tilfælde af MCTS-udforskningsparallelisering. Men blot at gemme handlingerne var ikke nok til at få et tilstrækkeligt effektivt handletrin. Faktisk ville tiden til at konvertere n_steps-handlingerne til (u, v, w) tripletten, reducere spillets tensortilstand og skabe de nye 3D-tensorer fra n_samples-handlingerne let være en flaskehals for hele træningen. For det andet ønskede vi ikke at gemme alle mulige fremtidige tilstande for hver samplet handling, da dette ville have en enorm indflydelse på den hukommelse, der bruges af algoritmen. Antag, at vi sætter n_samples=32, n=7 og N=5, og lad os huske, at N er størrelsen af det kvadratiske matrixprodukt, vi ønsker at reducere, og n er antallet af tidligere handlinger, som modellen husker. I denne situation ville hver tilstandstensor have formen (8, 25, 25, 25), som ganget med 32 ville resultere i 3282525254 bytes for hver node i grafen. I betragtning af, at hver simulering i udvidelsesfasen genererer en ny node (og n_sim=200), ville vi have et endeligt hukommelsesforbrug på 200328252525*4 = 3.2GB for den første MCTS-knude alene. I det værste tilfælde ville dette, mens man udforsker fungerende max_rank noder (hvor max_rank=150), resultere i et samlet hukommelsesforbrug på 150 * 3.2 GB = 480 GB i RAM-hukommelse (eller GPU-hukommelse, hvis alle tensorer var gemt på GPU'en) . Vi kørte træningen på vores arbejdsstation med 128 GB RAM og 48 GB GPU-hukommelse, så vi var nødt til at reducere hukommelsesforbruget.
Da vi ikke ønskede at øge eksekveringstiden, vedtog vi en optimering, der udnytter redundansen i de producerede statstensorer. Faktisk har tensorerne n-1 tidligere handlinger til fælles, som så kan gemmes én gang og ikke gentages for hver lagret tensor. Dette resulterer i en hukommelsesreduktion på 2/7~28%, hvilket betyder, at der i værste fald kan lagres 137 GB. På dette tidspunkt, ved blot at beskære den ubrugte del af træet (såsom de uvalgte baner) og gemme tensorerne i CPU-hukommelsen, var vi i stand til at undgå enhver hukommelsesfejl under træning.
Da OpenAlphaTensor nu er open source, åbner der sig flere spændende muligheder for videre udvikling.
En naturlig progression er finjusteringen af OpenAlphaTensor på målhardwareenheder. Dette forventes at føre til meget konkurrencedygtige beregningspræstationer. Jeg vil publicere mere om ydeevnen af OpenAlphaTensor på diverse hardware på GitHub. På tidspunktet for skrivning af denne artikel var OpenAlphaTensor under uddannelse.
Et andet vigtigt fremskridt ville være understøttelsen af fjernkompilering, hvilket giver brugerne mulighed for at bygge algoritmer, der er optimeret til edge-enheder. Dette kan opnås ved at gemme OpenAlphaTensor-modellen på en server, mens matrixmultiplikationsalgoritmen evalueres på forskellig hardware.
Det kan også være vigtigt at udvide understøttelsen af forskellige compilere til at beregne den latensbaserede belønningskorrektion. Forskellige compilere kan føre til forskellige optimerede algoritmer på en given hardware. For eksempel viste DeepMind-papiret lovende resultater ved brug af JAX og XLA-kompileren på TPU og Nvidia GPU'er. Det ville være interessant at evaluere dette ved hjælp af NCCL på Nvidia eller LLVM på CPU'er.
Endelig er det stadig en stor åben udfordring at udvide modellen og træningsalgoritmen til at understøtte større matrixstørrelser. I øjeblikket understøtter OpenAlphaTensor en maksimal matrixstørrelse på 5, men den kan anvendes ved at opdele større matrixmultiplikationer i grupper af bittesmå MM'er med en størrelse mindre end 5. Denne fremgangsmåde er suboptimal og udfører reduktionen direkte på den store tensor svarende til fuld MM kunne teoretisk set føre til bedre resultater.
Diego Fiori er CTO for Nebuly AI, en virksomhed, der er forpligtet til at gøre AI-optimering til en del af enhver udviklers værktøjssæt.
- SEO Powered Content & PR Distribution. Bliv forstærket i dag.
- Platoblokkæde. Web3 Metaverse Intelligence. Viden forstærket. Adgang her.
- Kilde: 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
- :er
- ][s
- $OP
- 1
- 3d
- 8
- a
- I stand
- Om
- over
- absolutte
- accelererende
- Ifølge
- derfor
- Konto
- nøjagtighed
- opnå
- opnået
- Handling
- aktioner
- faktisk
- tilføjet
- Yderligere
- Justeret
- vedtaget
- fremme
- Fordel
- Efter
- Agent
- aggregering
- aggressive
- AI
- AI-systemer
- målsætninger
- algoritme
- algoritmer
- Alle
- tillade
- tillader
- alene
- allerede
- alternativ
- blandt
- beløb
- analysere
- ,
- En anden
- anvendt
- tilgang
- tilgange
- arkitektur
- ER
- artikel
- kunstig
- kunstig intelligens
- AS
- tildelt
- forbundet
- At
- Automatiseret
- til rådighed
- undgå
- tilbage
- backup
- baseret
- I bund og grund
- grundlag
- BE
- fordi
- bliver
- før
- være
- jf. nedenstående
- benchmark
- BEDSTE
- Bedre
- mellem
- Beyond
- board
- Brætspil
- bundet
- bredere
- BT
- bygge
- bygget
- by
- kaldet
- CAN
- kan ikke
- tilfælde
- Årsag
- forårsagede
- vis
- udfordre
- udfordrende
- lave om
- Kanal
- ChatGPT
- barn
- valg
- vælge
- valgt
- citeret
- klar
- tydeligt
- kode
- indsamler
- kombinerer
- engageret
- Fælles
- selskab
- sammenlignet
- konkurrencedygtig
- komplekse
- kompleksitet
- komponenter
- sammensat
- kompromis
- beregning
- computerkraft
- Compute
- computing
- begreber
- Begrebsmæssigt
- betingelse
- tillid
- Overvej
- betragtes
- Overvejer
- anser
- forbrug
- indeholder
- fortsæt
- kontrast
- konverteret
- Core
- Tilsvarende
- svarer
- kunne
- Counter
- CPU
- skaber
- Oprettelse af
- kriterier
- CTO
- Nuværende
- For øjeblikket
- data
- beskæftiger
- beslutte
- besluttede
- beslutning
- Beslutningstagning
- afgørelser
- dyb
- dyb læring
- DeepMind
- afhænger
- beskrevet
- Bestem
- Udvikler
- Udvikling
- enhed
- Enheder
- DICT
- forskellige
- vanskeligheder
- Dimension
- størrelse
- direkte
- opdage
- opdage
- drøftet
- fordeling
- Divided
- ned
- Drop
- i løbet af
- e
- hver
- tidligere
- nemt
- Edge
- Effektiv
- effektiv
- enten
- element
- elementer
- indlejret
- ender
- forbedret
- enorm
- nok
- indrejse
- fejl
- skøn
- anslået
- Ether (ETH)
- evaluere
- evalueret
- evaluere
- evalueringer
- Endog
- Hver
- eksempel
- eksempler
- spændende
- udførelse
- udvidelse
- forventet
- erfaring
- Forklar
- forklarede
- exploits
- udforskning
- udforske
- udforsket
- Udforskning
- udtrykt
- udvide
- strækker
- ekstremt
- retfærdigt
- familie
- hurtigere
- tilbagemeldinger
- få
- Figur
- endelige
- finde
- Fornavn
- Flyde
- Fokus
- følger
- efter
- Fodspor
- Til
- formular
- format
- Formula
- fundet
- fra
- fuld
- funktion
- fundamental
- yderligere
- videre udvikling
- fremtiden
- spil
- Spil
- generere
- genereret
- genererer
- generere
- få
- få
- Giv
- given
- Give
- mål
- Goes
- godt
- GPU
- GPU'er
- graf
- stor
- gruppe
- Gruppens
- Vokser
- Vækst
- vejlede
- hånd
- Håndtering
- skete
- sker
- Hardware
- hardwareenhed
- hardwareenheder
- Have
- have
- link.
- højeste
- horisont
- Hvordan
- How To
- Men
- HTTPS
- kæmpe
- menneskelig
- i
- SYG
- ideal
- IDX
- billede
- KIMOs Succeshistorier
- gennemføre
- implementering
- vigtigt
- Forbedre
- forbedret
- in
- Forøg
- Stigninger
- stigende
- uafhængig
- indeks
- oplysninger
- initial
- indgang
- i stedet
- anvisninger
- Intelligens
- interessant
- interne
- introduceret
- intuition
- IT
- iteration
- ITS
- selv
- jpg
- KDnuggets
- Holde
- Nøgle
- viden
- kendt
- Sprog
- stor
- større
- Efternavn
- Latency
- seneste
- lag
- lag
- føre
- lærte
- læring
- Liste
- Se
- leder
- Lot
- lavet
- Main
- større
- lave
- Making
- administrere
- styring
- mange
- kort
- kortlægning
- Matrix
- Matter
- maksimal
- betyder
- meningsfuld
- midler
- medlem
- Hukommelse
- metode
- metoder
- metrisk
- mangler
- blanding
- model
- modeller
- modificeret
- modul
- mere
- mere effektiv
- Desuden
- mest
- bevæge sig
- bevæger sig
- ganget
- Natural
- Natur
- I nærheden af
- nødvendig
- Behov
- behov
- negativ
- net
- Neural
- neurale netværk
- Ny
- næste
- NLP
- node
- noder
- ikke-eksperter
- Begreb
- nummer
- Nvidia
- opnået
- of
- Tilbud
- on
- ONE
- åbent
- open source
- OpenAI
- drift
- Produktion
- optimal
- optimering
- Optimer
- optimeret
- Indstillinger
- original
- Andet
- Ellers
- output
- samlet
- Papir
- papirer
- parameter
- del
- særlig
- især
- dele
- Mennesker
- perfekt
- ydeevne
- udfører
- fase
- stykker
- plato
- Platon Data Intelligence
- PlatoData
- Leg
- spiller
- spiller
- Punkt
- politikker
- politik
- position
- mulig
- potentiale
- magt
- Praktisk
- praksis
- forudsige
- forudsagde
- præferencer
- forelagt
- tidligere
- sandsynlighed
- sandsynligvis
- Problem
- behandle
- Processer
- producere
- produceret
- Produkt
- progression
- progressiv
- lovende
- foreslog
- beviseligt
- gennemprøvet
- offentliggøre
- offentliggjort
- RAM
- tilfældig
- rækker
- hellere
- nået
- når
- for nylig
- reducere
- Reduceret
- reducere
- raffinerede
- forstærkning læring
- frigivet
- resterende
- resterne
- bemærkelsesværdig
- huske
- fjern
- gentaget
- repræsenterer
- repræsenterer
- påkrævet
- Kræver
- ansvarlige
- begrænset
- resultere
- resulterer
- Resultater
- afkast
- afkast
- revolutionere
- Beløn
- RÆKKE
- rt
- Kør
- s
- samme
- Gem
- besparelse
- scenarie
- scenarier
- Ordningen
- Søg
- Anden
- Sektion
- frø
- valgt
- udvælgelse
- valg
- SELV
- sæt
- indstilling
- flere
- Shape
- deling
- Kort
- bør
- Vis
- vist
- Shows
- Signal
- signifikant
- betydeligt
- lignende
- ligheder
- Simpelt
- enkelhed
- ganske enkelt
- simulation
- siden
- Situationen
- Størrelse
- størrelser
- lille
- mindre
- Snapshot
- So
- løsninger
- SOLVE
- Løsning
- nogle
- Kilde
- Space
- rum
- specifikke
- specifikt
- specificeret
- hastighed
- tilbringe
- udgifterne
- delt
- firkant
- stablet
- Stage
- standard
- starte
- starter
- Tilstand
- state-of-the-art
- erklærede
- Stater
- statistik
- Trin
- Steps
- standsning
- butik
- opbevaret
- forhandler
- ligetil
- vellykket
- sådan
- support
- Understøtter
- syntetisk
- syntetiske data
- syntetisk
- systemet
- Systemer
- skræddersyet
- Tag
- tager
- tager
- mål
- Opgaver
- opgaver
- teknikker
- terminal
- vilkår
- at
- Fremtiden
- Grafen
- oplysninger
- Staten
- deres
- Them
- derved
- derfor
- Disse
- Tredje
- tre
- tredimensionale
- Gennem
- tid
- tidskrævende
- til
- sammen
- token
- tokenization
- tokeniseret
- Tokens
- også
- toolkit
- top
- fakkel
- I alt
- traditionelle
- Tog
- uddannet
- Kurser
- tog
- bane
- omdannet
- behandle
- sand
- forstå
- Universe
- ubrugt
- Opdatering
- opdateret
- opdateringer
- opdatering
- opgradering
- us
- Brug
- brug
- brugere
- sædvanligvis
- værdi
- Værdier
- forskellige
- udgave
- video
- videospil
- Besøg
- Besøg
- W
- Vej..
- Hvad
- som
- mens
- bredt
- Wikipedia
- vilje
- vindende
- med
- ord
- arbejdsstation
- værd
- ville
- skrivning
- skriftlig
- X
- zephyrnet
- nul