Photo by DeepMind on Unsplash
Matrisemultiplikasjon er en grunnleggende operasjon som brukes i mange systemer, fra nevrale nettverk til vitenskapelige datarutiner. Å finne effektive og beviselig riktige algoritmer for matrisemultiplikasjon kan ha stor innvirkning på å gjøre beregningen raskere og mer effektiv, men er en svært utfordrende oppgave. Rommet av mulige algoritmer er enormt, og tradisjonelle metoder for å oppdage algoritmer, som menneskeskapte heuristikk eller kombinatorisk søk, er ofte suboptimale.
DeepMindsin nylig foreslåtte AI-baserte løsning for automatisert søk går langt utover menneskelig intuisjon. Løsningen består av en dyp forsterkningslæringsagent kalt AlphaTensor, bygget på toppen av alphazero. Denne agenten er opplært til å spille et enkeltspillerspill, TensorGame, hvor målet er å oppdage beregningseffektive algoritmer for matrisemultiplikasjon.
AlphaTensor er spesielt flink til å håndtere store matriser ved å dekomponere store matrisemultiplikasjoner til mindre multiplikasjoner. Dessuten kan AlphaTensor brukes til å oppnå toppmoderne ytelse for matrisemultiplikasjon når den er finjustert på en spesifikk maskinvareenhet.
AlphaTensor har et stort potensial for å akselerere dyp læringsdatabehandling. I dyp læring kan mange tidkrevende operasjoner kartlegges til matrisemultiplikasjoner. Ved å bruke AlphaTensor for å optimalisere disse operasjonene, kan den generelle ytelsen til dyplæringsmodeller forbedres betydelig.
Nylig har OpenAlphaTensor, den første åpen kildekode-implementering av AlphaTensor, ble utgitt, noe som kan revolusjonere beregningskraften til dyplæringsmodeller.
Matrisemultiplikasjonstensor
For ikke-eksperter innen matrisemultiplikasjonsoptimalisering er det kanskje ikke enkelt å forstå hvordan en operasjon som matrisemultiplikasjon kan kartlegges i en tredimensjonal tensor. Jeg skal prøve å forklare det med enkle ord og med eksempler.
La oss vurdere produktet C = A*B, hvor for enkelhets skyld både A og B er kvadratiske matriser av størrelse N. Multiplikasjonsoperasjonen kan kartlegges i en 3D-tensor av form (N^2, N^2, N^2). Den første tensordimensjonen representerer den flatede matrisen A, den andre dimensjonen den flatede matrisen B og den tredje dimensjonen den flatede matrisen C.
Tensoren har bare binære verdier (enten 1 eller 0) for hver oppføring. Merk at tensoren representerer multiplikasjonsoperasjonen, så den er uavhengig av verdiene til matrisene A og B.
Hver inntasting av tensoren tilsvarer koeffisienten for operasjonen. For eksempel, for å beregne C[1,1], er det nødvendig å multiplisere både A[1,1] og B[1,1]. Derfor vil tensorinngangen [0,0,0], som tilsvarer A[1,1], B[1,1] og C[1,1], ha verdi 1. I motsetning til dette, for å beregne C[1,1 ,2,1], A[1] er ikke nødvendig. Tensorraden T[N+0, :, XNUMX] vil således kun inneholde nuller.
Bildet nedenfor viser et eksempel på en tensor for N=2.
Bilde fra DeepMind's papir publisert i Natur
Som vist i (b) og (c) i figuren ovenfor, er det mulig å implementere en algoritme for å beregne produktet ved å bruke en dekomponering av 3D-tensoren. Mer spesifikt kan algoritmen nedenfor brukes til å konvertere en tensordekomponering (matrisene U, V, W) til en matrisemultiplikasjonsalgoritme.
Metaalgoritme parametrisert for å beregne matriseproduktet C=AB introdusert i DeepMinds papir
Tensor-spillet
Problemet med å finne effektive algoritmer for matrisemultiplikasjon er ekstremt utfordrende fordi antallet mulige algoritmer å vurdere er mye større enn antall atomer i universet, selv for små tilfeller av matrisemultiplikasjon.
DeepMind konverterte dette problemet til et enkeltspillerspill, og kalte det TensorGame. I dette spillet velger spilleren hvordan de skal kombinere ulike oppføringer av matriser for å multiplisere dem. En poengsum tildeles basert på antall operasjoner som kreves for å oppnå riktig multiplikasjonsresultat. Spillet avsluttes når nulltensoren er nådd eller når maksimalt antall trekk er gjort. Den endelige faktoriseringen blir evaluert basert på en estimering av gjenværende rangering og visse optimaliseringskriterier, for eksempel asymptotisk tidskompleksitet eller praktisk kjøretid.
Startposisjonen i TensorGame tilsvarer Matrix Multiplication Tensor uttrykt på tilfeldig basis.
I hvert trinn t i spillet skriver spilleren ned tre vektorer
, som spesifiserer rang-1 tensorer . Tilstanden til spillet oppdateres ved å trekke fra vektorene valgt av spilleren:
hvor
er matrisemultiplikasjonstensoren.Hvis spillet ender i p trinn, betyr dette at Matrix Multiplication Tensor
kan dekomponeres til p rank-1 tensorer , dvs. den har minst rang p.TensorGame kan da tolkes som en rang-dekomponeringsalgoritme og AlphaTensor kan sees på som en algoritme for å estimere rangeringen til tensoren.
AlphaTensor-arkitektur
Så langt har vi lært om TensorGame og avklart hvordan løsningen kan sees på som en matrisemultiplikasjonsalgoritme. La oss nå utforske hovedkonseptene til AlphaTensor, algoritmen som brukes for spillet.
AlphaTensor-arkitektur er i utgangspunktet en koder-dekoder-transformatorarkitektur der:
- koderen tar spillets tilstand som input , de n tidligere handlingene som ble utført av modellen (vanligvis n=7), og tidsindeksen t for gjeldende handling. Informasjon er stablet sammen i en tensor med form (n+1, N^2, N^2, N^2). Denne tensoren blir deretter omformet og transformert (ved hjelp av tre lineære lag) i en tensor av form (N^2, N^2, c) der c er den indre dimensjonen til modellen.
- dekoderen genererer n_steps-handlingene fra den innebygde vektoren gitt av koderen på en autoregressiv måte. Hver handling tilsvarer en token av trillingene som representerer en av trillingene som bryter ned spilltensoren (dvs. reduserer rangeringen)
Modellen trenes opp ved å veksle mellom ryggutbredelse og modellskuespill. Model acting brukes til å generere data som deretter brukes til å trene modellen. I praksis trenes modellen med en blanding av syntetisk genererte data og data generert av modellen under skuespill. Skuespilltrinnet gjøres ved å ta en 3D-tensor som tilsvarer en matriseoperasjon og spille n_actors-spill på den. Hver skuespiller spiller et spill enten på standardbasis eller på alternativ basis (endringen av grunnlaget brukes med en gitt sannsynlighet). Resultatene samles så inn og kan brukes i treningstrinnet med de syntetiske dataene.
Skuespillertrinnet er basert på AlphaZeros Monte Carlo Tree Search (MCTS), modifisert for å støtte store handlingsrom. Kort sagt, før du velger handling, utforskes n_sims-stier fra modellutgangen med en maksimal fremtidig utforskning på 5 trinn. Sannsynlighetene som genereres av modellen justeres deretter under hensyntagen til de genererte banene. Deretter velges handlingen med den/de mest lovende fremtidige veien(e) for å fortsette spillet.
Mens du trener modellen, er belønningen faktisk en negativ belønning (straff). Dens absolutte verdi øker for hvert ekstra trinn som kreves for å løse spillet. Hvis modellen tar m skritt for å løse et TensorGame, er belønningen knyttet til spillet r=-m. Hvis modellen ikke er i stand til å løse TensorGame i max_rank-trinn, beregnes belønningen ved å estimere rangeringen til den gjenværende tensoren. Rangeringen er estimert som summen av rekkene til matrisene som utgjør tensoren. Estimatet er en øvre grense for den sanne rangeringen av tensoren.
Når du finjusterer modellen, bør straffebelønningen ved terminaltilstanden også ta hensyn til latensen til algoritmen produsert av modellen. Belønningsformelen blir rt'=rt+λbt, der rt er belønningsordningen beskrevet tidligere, bt er referansebelønningen (ikke null bare ved terminaltilstanden), og λ er en brukerspesifisert koeffisient.
Speed-ups (%) av AlphaTensor-oppdagede algoritmer skreddersydd for en GPU og en TPU, hentet fra DeepMinds papir. Speed-ups måles i forhold til standard (f.eks. cuBLAS for GPU) matrisemultiplikasjon på samme maskinvare og sammenlignet med Strassen-kvadratalgoritme. Kilde: DeepMind.
Jeg slapp nylig OpenAlphaTensor, den første åpen kildekode-implementeringen av AlphaTensor. I denne delen vil jeg gå gjennom implementeringen. Som vi diskuterte tidligere, er AlphaTensor-arkitekturen ganske enkel, basert på en standard transformator med en koder-dekoder-arkitektur. De mest interessante komponentene til AlphaTensor er det første laget i koderdelen og måten handlingene samples på.
La oss starte med det første kodingslaget.
# 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 utdraget ovenfor viser vi hvordan inngangstensoren dekomponeres i tre tensorer, som deretter brukes som spørrings-, nøkkel- og verdiinndata for transformatorlaget.
- På tvers av de tre tensordimensjonene som representerer de flate matrisene (A, B, C), er inngangstensoren flatet langs hver dimensjon sammen med dimensjonen som representerer de foregående handlingene. På denne måten, i hver flatet-kopi av inngangstensoren, er den valgte dimensjonen en aggregering av de siste T-1-verdiene og den faktiske verdien, for alle S-verdiene til den valgte dimensjonen, der S=N^2. Filosofisk sett er det som om vi for hver dimensjon fokuserer på det som skjedde i de tidligere handlingene i den dimensjonen.
- Skalarene er kartlagt i tre forskjellige rom med dimensjon S^2, og deretter omformet for å bli sammenkoblet med tensorene oppnådd ved forrige punkt. Konseptuelt blir skalarene kartlagt til et innebygd rom med dimensjon S^2, og deretter blir den innebygde informasjonen delt inn i S-vektorer og stablet sammen, på samme måte som det som skjer med tekst når den blir tokenisert.
- Skalare tokens er sammenkoblet med den restrukturerte inputtensoren og deretter gitt som input til et lineært lag for å kartlegge fokusinformasjonen for skalarer+kanalhistorie i den interne dimensjonen til modellen.
Disse tre trinnene kan tolkes som en måte å gi modellen både informasjon om skalarene (som i TensorGame-tidstrinnet) og fokus på de tidligere handlingene for hver kanal.
Når det gjelder måten handlingene produseres på, er det interessant å merke seg at AlphaTensor genererer tripletten u, v, w som utgang, som har som mål å redusere tensorrangen. De tre vektorene har størrelse S og siden de er sammenkoblet, må modellen produsere en vektor med størrelse 3*S. AlphaTensor er trent med en RL-algoritme, så alle mulige handlinger må uttrykkes i form av sannsynligheter i et oppregnet rom, dvs. at modellen produserer en sannsynlighet over de forskjellige handlingene. Dette betyr at hver vektor i 3S-rommet bør kartlegges til en annen handling. Dette resulterer i et handlingsrom med størrelse |F|^(3S), hvor |F| er antallet forskjellige verdier som elementet av u, v, w kan ta. Vanligvis er verdiene begrenset til (-2, -1, 0, 1, 2), noe som resulterer i en kardinalitet på 5 elementer.
Her kommer en stor utfordring: for å generere handlingssannsynlighetene for et matriseprodukt av matriser av størrelse 5, trenger vi et minne på 5^75 * 4 byte, som vil bety ~10^44 GB minne. Det er klart at vi ikke klarer et så stort handlingsrom.
Hvordan løser vi problemet? For å redusere minnefotavtrykket til handlingssannsynlighetene kan vi dele trillingene i mindre biter, "tokenisere" dem, og behandle bitene som genererte tokens i transformatorarkitekturen, dvs. tokens gis som input til dekoderen i en autoregressiv vei. I eksemplet ovenfor kan vi dele trillingene i 15 biter, og redusere minneforbruket 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 kodebiten for å generere hele handlingen. I koden inneholder self.core dekoderlaget og tensoren e representerer utdata fra koderlaget. Null kan betraktes som token i NLP-modeller og n_steps-handlingene som representerer n_steps-bitene genereres på en progressiv måte.
Modellen returnerer tre mengder:
- De genererte handlingene
- Sannsynligheten knyttet til hele handlingen
- Logittene produsert for å generere den første handlingen (den første delen) som vil bli brukt til å beregne modellverdien.
Det er verdt å bruke noen få ord på parameteren n_samples. Parameteren brukes for handlingstrinnet og lar modellen generere forskjellige versjoner av trillingene som deretter vil bli brukt til å utforske handlingsrommet i Monte Carlo Tree Search-algoritmen som brukes i skuespillerprosessen. De forskjellige handlingene n_samples samples i henhold til policyen som genereres av modellen.
Fungerende trinn
Den mest vanskelige delen av hele algoritmen er sannsynligvis Acting-trinnet som brukes for å løse TensorGame. Algoritmen er ikke dypt forklart i AlphaTensor-artikkelen, siden den er basert på flere DeepMinds tidligere artikler som bare er sitert og gitt som kjent. Her vil jeg rekonstruere alle de manglende delene og forklare trinn for trinn implementeringen vår.
Vi kan organisere skuespilletrinnene i tre forskjellige komponenter:
- Monte-Carlo Tree Search
- Spillsimuleringen
- Den forbedrede policyberegningen
La oss analysere dem en etter en.
Monte-Carlo Tree Search (MCTS)
Monte Carlo Tree Search (MCTS) er en mye brukt kunstig intelligens-teknikk for spill, spesielt i brettspill og videospill. Algoritmen lager et spilltre som simulerer potensielle trekk og utfall og bruker tilfeldig prøvetaking for å evaluere forventet belønning for hvert trekk. Algoritmen velger deretter iterativt trekket med den høyeste forventede belønningen og simulerer utfall til den når en terminal tilstand eller en spesifisert stopptilstand. Simuleringene brukes til å estimere sannsynligheten for å vinne for hvert trekk og veilede beslutningsprosessen. MCTS har vist seg å være effektiv i komplekse spill hvor antallet mulige trekk og utfall er stort, og det har blitt brukt i vellykkede AI-systemer som AlphaGo.
I AlphaTensor brukes en modifisert versjon av den originale MCTS. Spesielt, i stedet for å velge handlingen tilfeldig fra hele handlingsrommet, velges handlingen blant et delsett generert direkte av modellen (gjennom n_samplene presentert før). Korrigeringen av policyoppgraderingen blir deretter brukt i trinnet for forbedret policyberegning.
I implementeringen vår bestemte vi oss for å beholde all informasjon om Monte-Carlo-treet i en ordbok som har hash-versjonen av TensorGame-tilstanden som nøkkel, og som verdier informasjonen knyttet til selve staten. Hvert Monte-Carlo-trinn starter fra en node og simulerer n_sim-minispill, og utforsker fremtiden med en horisont på 5 trekk. Hvis noden allerede er utforsket i tidligere simuleringer, justeres n_sim med tanke på antall tidligere undersøkelser. For hver node lagres antall besøk i N_s_a-tensoren, siden denne tensoren inneholder antall besøk per node-underordnet handling (blant de samplet av 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
Koden ovenfor viser vår implementering av algoritmen. For enkelhets skyld utføres policykorreksjonen i simulate_game-funksjonen.
Spill Simulering
Simulate_game-funksjonen er ansvarlig for å utforske treet som består av noder som representerer en bestemt tilstand av TensorGame. Den kjører også modellen hver gang en bladnode støtes på, og den lagrer all nodeinformasjon i state_dict-ordboken. La oss gi en dyp titt 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 delt inn i tre deler:
- utvalg
- Ekspansjon
- Backup
I utvalgsdelen kjøres simuleringen på de allerede genererte trenodene, og følgende node velges ved hjelp av følgende funksjon:
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, handlingen som maksimerer ucb-funksjonen:
for den gitte tilstanden er valgt. Her representerer Q Q-verdiene generert av modellen og π representerer den tilfeldige fordelingen over handlingene samplet ved bruk av modellpolicyen. N(s, a) representerer antall besøk av noden til handling a fra node s.
Når seleksjonsfasen når en bladnode, hvis simuleringen ikke har nådd en terminal tilstand (i form av enten maksimal utforskning, dvs. fremtidig horisont, eller spillslutt), brukes modellen for å velge n_samples alternative noder (de vil være blad noder i den påfølgende iterasjonen). Dette kalles utvidelsesfasen, siden nye noder legges til treet. Deretter utforskes ingen ytterligere node i gjeldende simulering, men leaf q_value sendes til følgende simuleringstrinn: backup.
Sikkerhetskopiering er den siste fasen av hver simulering. Under backup, hvis bladnoden var en terminal tilstand, beregnes den endelige belønningen; ellers brukes blad q-verdien som en estimert belønning. Deretter forplantes belønningen tilbake på simuleringsbanen som oppdaterer både tilstandene q_values og oppdaterer besøkstelleren N(s, a). I utdraget nedenfor viser vi koden for belønningens tilbakeformidling.
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 policyberegning
Når alle simuleringene er kjørt og MCTS gir et interessant øyeblikksbilde av nær fremtid, er det på tide å oppdatere policyen knyttet til de forutsagte nodene og returnere dem, slik at de kan brukes under trening. Den forbedrede policyen, etter metoden beskrevet i Hubert et al, brukes til å administrere store handlingsrom. Faktisk, for lite søkerom, er det mulig under MCTS å prøve en handling tilfeldig fra handlingsrommet og evaluere virkningen. En lignende tilnærming i et mye større handlingsrom vil føre til at alle baner divergerer i forskjellige baner, og det vil trenge en uendelig mengde baner for å få meningsfull statistikk og deretter oppdatere policyen. Siden vi her bruker sample-MCTS for å unngå spredningen, dvs. at n_samples-handlinger samples i henhold til modellpolicyen og så velger MCTS bare en av samplede handlinger mens vi utforsker treet, må vi ta hensyn til sample-korrigeringen når vi beregner den endelige oppdaterte policyen som vil bli brukt under opplæring av modellen.
I praksis er den forbedrede politikken beregnet 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
Merk at i implementeringen vår etter å ha beregnet policyen fra N_s_a-tensoren, må vi kartlegge den tilbake til den opprinnelige handlingstensoren. Faktisk vurderer N_s_a bare handlingene samplet av modellen, mens den endelige policyen må inneholde sannsynligheter også for de ikke-utforskede handlingene.
Forskjeller med hensyn til ChatGPT treningsalgoritme
AlphaTensor er det siste medlemmet av AlphaGo/AlphaZero-familien av kunstig intelligens-metoder fra DeepMind. Disse metodene er basert på Monte Carlo Tree Search (MCTS) algoritmen, som har blitt raffinert og forbedret av DeepMind for å takle stadig mer komplekse oppgaver. Et annet AI-system, OpenAIs ChatGPT, som har skapt mye buzz for sin bemerkelsesverdige ytelse, ble trent med en annen tilnærming, kalt Reinforcement Learning with Human Feedback (RLHF).
RLHF er en finjusteringsteknikk som brukes til å stille inn språkmodeller for å følge et sett med skriftlige instruksjoner. Den bruker menneskelige preferanser som et belønningssignal for å finjustere modellen, og dermed justere oppførselen til språkmodellen med de uttalte preferansene til en spesifikk gruppe mennesker, i stedet for en bredere forestilling om 'menneskelige verdier'.
Derimot er MCTS en trebasert søkealgoritme som brukes til å bestemme de optimale bevegelsene i spill. Den simulerer potensielle trekk og oppdaterer verdiene for hvert trekk basert på resultatene, og veileder valget av det beste trekket.
RLHF samler inn data fra menneskeskrevne demonstrasjoner og menneskemerkede sammenligninger mellom AI-modeller, og trener opp en belønningsmodell for å forutsi preferansene til en gitt gruppe mennesker. Belønningsmodellen brukes deretter til å finjustere AI-modellene. MCTS, derimot, bruker simuleringer og evalueringer for å bestemme den beste beslutningen.
Selv om de er forskjellige tilnærminger, har RLHF og MCTS også likheter. Begge teknikkene for kunstig intelligens bruker beslutnings- og problemløsningsmetoder, og begge bruker en prøving-og-feil-tilnærming for å utforske ulike alternativer og ta avgjørelser basert på tilgjengelig informasjon. Begge er også iterative prosesser som forbedres over tid etter hvert som mer informasjon og erfaring samles inn.
Valget mellom RLHF og MCTS avhenger av oppgaven. RLHF er ideell når det ikke finnes en klar metrikk for å evaluere modellens ytelse, mens MCTS har vist seg effektiv i spilllignende oppgaver hvor kunnskap og utforskning av fremtiden gir modellen en betydelig fordel.
Kodeoptimalisering for AlphaTensor-trening
Implementering av AlphaTensor-treningsalgoritmen krever å finne det perfekte kompromisset mellom treningshastighet og minneforbruk. Som sett i modelldelen, kan det å bare vurdere handlingstokeniseringen spare mye minne, men en altfor aggressiv reduksjon av handlingsrom kan føre til både fall i nøyaktighet og tregere ytelse. Det siste skjer fordi alle tokens genereres sekvensielt på en autoregressiv måte av modelldekoderen. Derfor vokser slutningstiden lineært med antall tokens per handling når softmax på handlingsområdet ikke lenger er flaskehalsen.
Ved oppsett av AlphaTensor-trening ble hovedvanskene funnet i å håndtere skuespillerprosessen. Hvis tensorene ikke er lagret i riktig format, kan MCTS lett forårsake ukontrollert vekst i minnebruken. På den annen side, hvis antall tensorer som er lagret under hver simulering reduseres for mye, kan MCTS bruke uendelig mye tid på å beregne de nødvendige tilstandene på nytt.
La oss ta et eksempel på spillsimuleringstrinnet, der spillet utforskes ved å se på mulige fremtidige scenarier. For hver stat, hvis vi ikke lagrer handlingene generert av modellen og vi bestemmer oss for å lagre bare det tilfeldige frøet som ble brukt til å prøve handlingene fra policyen, må vi beregne policyen på nytt hver gang vi utforsker en trenode. deretter prøve handlingene. Det er klart at vi bestemte oss for å lagre de samplede handlingene for å spare tid og for å unngå å måtte administrere modelldeling mellom ulike prosesser i tilfelle MCTS-utforskningsparallellisering. Det var imidlertid ikke nok å lagre handlingene for å få et tilstrekkelig effektivt handlingstrinn. Faktisk vil tiden for å konvertere n_steps-handlingene til (u, v, w) tripletten, redusere spillets tensortilstand og lage de nye 3D-tensorene fra n_samples-handlingene lett være en flaskehals for hele treningen. For det andre ønsket vi ikke å lagre alle mulige fremtidige tilstander for hver samplet handling, da dette ville ha en enorm innvirkning på minnet som brukes av algoritmen. Anta at vi setter n_samples=32, n=7 og N=5, og la oss huske at N er størrelsen på det kvadratiske matriseproduktet vi ønsker å redusere og n er antallet tidligere handlinger som er husket av modellen. I denne situasjonen vil hver tilstandstensor ha formen (8, 25, 25, 25), som multiplisert med 32 vil resultere i 3282525254 byte for hver node i grafen. Nå, med tanke på at hver simulering i utvidelsesfasen genererer en ny node (og n_sim=200), vil vi ha et endelig minneforbruk på 200328252525*4 = 3.2 GB for den første MCTS-noden alene. I verste fall, mens du utforsker fungerende max_rank-noder (der max_rank=150), vil dette resultere i et totalt minneforbruk på 150 * 3.2 GB = 480 GB i RAM-minne (eller GPU-minne hvis alle tensorer var lagret på GPUen) . Vi kjørte treningen på arbeidsstasjonen vår med 128 GB RAM og 48 GB GPU-minne, så vi måtte redusere minneforbruket.
Siden vi ikke ønsket å øke gjennomføringstiden, tok vi i bruk en optimalisering som utnytter redundansen i de statlige tensorene som produseres. Tensorene har faktisk n-1 tidligere handlinger til felles, som deretter kan lagres en gang og ikke gjentas for hver lagret tensor. Dette resulterer i en minnereduksjon på 2/7~28%, noe som betyr at i verste fall kan 137 GB lagres. På dette tidspunktet, ved ganske enkelt å beskjære den ubrukte delen av treet (som de uvalgte banene) og lagre tensorene i CPU-minnet, var vi i stand til å unngå minnefeil under trening.
Når OpenAlphaTensor nå er åpen kildekode, åpner det seg flere spennende veier for videre utvikling.
En naturlig progresjon er finjusteringen av OpenAlphaTensor på målmaskinvareenheter. Dette forventes å føre til svært konkurransedyktige beregningsytelser. Jeg vil publisere mer om ytelsen til OpenAlphaTensor på diverse maskinvare på GitHub. På tidspunktet for skriving av denne artikkelen var OpenAlphaTensor under opplæring.
Et annet viktig fremskritt ville være støtte for ekstern kompilering, slik at brukere kan bygge algoritmer optimalisert for edge-enheter. Dette kan oppnås ved å lagre OpenAlphaTensor-modellen på en server, mens matrisemultiplikasjonsalgoritmen evalueres på forskjellig maskinvare.
Det kan også være viktig å utvide støtten til forskjellige kompilatorer for å beregne den latensbaserte belønningskorrigeringen. Ulike kompilatorer kan føre til forskjellige optimaliserte algoritmer på en gitt maskinvare. For eksempel viste DeepMind-papiret lovende resultater ved å bruke JAX og XLA-kompilatoren på TPU og Nvidia GPUer. Det ville vært interessant å evaluere dette ved å bruke NCCL på Nvidia eller LLVM på CPUer.
Til slutt er det fortsatt en stor åpen utfordring å utvide modellen og treningsalgoritmen til å støtte større matrisestørrelser. For øyeblikket støtter OpenAlphaTensor en maksimal matrisestørrelse på 5, men den kan brukes ved å dele større matrisemultiplikasjoner i grupper av bittesmå MM-er med en størrelse mindre enn 5. Denne tilnærmingen er suboptimal, og utfører reduksjonen direkte på den store tensoren som tilsvarer full MM vil teoretisk kunne føre til bedre resultater.
Diego Fiori er CTO for Nebuly AI, et selskap som er forpliktet til å gjøre AI-optimalisering til en del av hver utvikleres verktøysett.
- SEO-drevet innhold og PR-distribusjon. Bli forsterket i dag.
- Platoblokkkjede. Web3 Metaverse Intelligence. Kunnskap forsterket. Tilgang 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
- $OPP
- 1
- 3d
- 8
- a
- I stand
- Om oss
- ovenfor
- Absolute
- akselerer
- Ifølge
- tilsvar
- Logg inn
- nøyaktighet
- Oppnå
- oppnådd
- Handling
- handlinger
- faktisk
- la til
- Ytterligere
- justert
- vedtatt
- avansere
- Fordel
- Etter
- Agent
- aggregering
- aggressiv
- AI
- AI-systemer
- mål
- algoritme
- algoritmer
- Alle
- tillate
- tillater
- alene
- allerede
- alternativ
- blant
- beløp
- analysere
- og
- En annen
- anvendt
- tilnærming
- tilnærminger
- arkitektur
- ER
- Artikkel
- kunstig
- kunstig intelligens
- AS
- tildelt
- assosiert
- At
- Automatisert
- tilgjengelig
- unngå
- tilbake
- Backup
- basert
- I utgangspunktet
- basis
- BE
- fordi
- blir
- før du
- være
- under
- benchmark
- BEST
- Bedre
- mellom
- Beyond
- borde
- Board Games
- bundet
- bredere
- BT
- bygge
- bygget
- by
- som heter
- CAN
- kan ikke
- saken
- Årsak
- forårsaket
- viss
- utfordre
- utfordrende
- endring
- Kanal
- ChatGPT
- barn
- valg
- velge
- valgt ut
- sitert
- fjerne
- klart
- kode
- innsamler
- kombinere
- forpliktet
- Felles
- Selskapet
- sammenlignet
- konkurranse
- komplekse
- kompleksitet
- komponenter
- komponert
- kompromiss
- beregningen
- regnekraft
- Beregn
- databehandling
- konsepter
- Konseptuelt
- tilstand
- selvtillit
- Vurder
- ansett
- vurderer
- anser
- forbruk
- inneholder
- fortsette
- kontrast
- konvertert
- Kjerne
- Tilsvarende
- tilsvarer
- kunne
- Motvirke
- prosessor
- skaper
- Opprette
- kriterier
- CTO
- Gjeldende
- I dag
- dato
- håndtering
- bestemme
- besluttet
- avgjørelse
- Beslutningstaking
- avgjørelser
- dyp
- dyp læring
- DeepMind
- avhenger
- beskrevet
- Bestem
- Utvikler
- Utvikling
- enhet
- Enheter
- DIKT
- forskjellig
- vanskeligheter
- Dimensjon
- dimensjoner
- direkte
- oppdage
- oppdage
- diskutert
- distribusjon
- Divided
- ned
- Drop
- under
- e
- hver enkelt
- Tidligere
- lett
- Edge
- Effektiv
- effektiv
- enten
- element
- elementer
- innebygd
- slutter
- forbedret
- enorm
- nok
- entry
- feil
- anslag
- anslått
- Eter (ETH)
- evaluere
- evaluert
- evaluere
- evalueringer
- Selv
- Hver
- eksempel
- eksempler
- spennende
- gjennomføring
- utvidelse
- forventet
- erfaring
- Forklar
- forklarte
- exploits
- leting
- utforske
- utforsket
- Utforske
- uttrykte
- utvide
- strekker
- ekstremt
- ganske
- familie
- raskere
- tilbakemelding
- Noen få
- Figur
- slutt~~POS=TRUNC
- finne
- Først
- Flyte
- Fokus
- følge
- etter
- Fotspor
- Til
- skjema
- format
- formel
- funnet
- fra
- fullt
- funksjon
- fundamental
- videre
- videre utvikling
- framtid
- spill
- Games
- generere
- generert
- genererer
- genererer
- få
- få
- Gi
- gitt
- Giving
- mål
- Går
- god
- GPU
- GPU
- graf
- flott
- Gruppe
- Gruppens
- Vokser
- Vekst
- veilede
- hånd
- Håndtering
- skjedde
- skjer
- maskinvare
- maskinvareenhet
- maskinvareenheter
- Ha
- å ha
- her.
- høyest
- horisont
- Hvordan
- Hvordan
- Men
- HTTPS
- stort
- menneskelig
- i
- JEG VIL
- ideell
- iDX
- bilde
- Påvirkning
- iverksette
- gjennomføring
- viktig
- forbedre
- forbedret
- in
- Øke
- øker
- stadig
- uavhengig
- indeks
- informasjon
- innledende
- inngang
- i stedet
- instruksjoner
- Intelligens
- interessant
- intern
- introdusert
- intuisjon
- IT
- køyring
- DET ER
- selv
- jpg
- KDnuggets
- Hold
- nøkkel
- kunnskap
- kjent
- Språk
- stor
- større
- Siste
- Ventetid
- siste
- lag
- lag
- føre
- lært
- læring
- Liste
- Se
- ser
- Lot
- laget
- Hoved
- større
- gjøre
- Making
- administrer
- administrerende
- mange
- kart
- kartlegging
- Matrix
- Saken
- maksimal
- betyr
- meningsfylt
- midler
- medlem
- Minne
- metode
- metoder
- metrisk
- mangler
- blanding
- modell
- modeller
- modifisert
- moduler
- mer
- mer effektivt
- Videre
- mest
- flytte
- trekk
- multiplisert
- Naturlig
- Natur
- Nær
- nødvendig
- Trenger
- nødvendig
- negativ
- nettverk
- neural
- nevrale nettverk
- Ny
- neste
- nlp
- node
- noder
- ikke-eksperter
- Forestilling
- Antall
- Nvidia
- innhentet
- of
- Tilbud
- on
- ONE
- åpen
- åpen kildekode
- OpenAI
- drift
- Drift
- optimal
- optimalisering
- Optimalisere
- optimalisert
- alternativer
- original
- Annen
- ellers
- produksjon
- samlet
- Papir
- papirer
- parameter
- del
- Spesielt
- spesielt
- deler
- porsjoner
- perfekt
- ytelse
- utfører
- fase
- stykker
- plato
- Platon Data Intelligence
- PlatonData
- Spille
- spiller
- spiller
- Point
- Politikk
- politikk
- posisjon
- mulig
- potensiell
- makt
- Praktisk
- praksis
- forutsi
- spådd
- preferanser
- presentert
- forrige
- sannsynlighet
- sannsynligvis
- Problem
- prosess
- Prosesser
- produsere
- produsert
- Produkt
- progresjon
- progressiv
- lovende
- foreslått
- beviselig
- utprøvd
- publisere
- publisert
- RAM
- tilfeldig
- rekkene
- heller
- nådd
- Når
- nylig
- redusere
- Redusert
- redusere
- raffinert
- forsterkning læring
- utgitt
- gjenværende
- forblir
- bemerkelsesverdig
- husker
- fjernkontroll
- gjentatt
- representerer
- representerer
- påkrevd
- Krever
- ansvarlig
- begrenset
- resultere
- resulterende
- Resultater
- retur
- avkastning
- Revolusjonere
- Belønn
- RAD
- rt
- Kjør
- s
- samme
- Spar
- besparende
- scenario
- scenarier
- ordningen
- Søk
- Sekund
- Seksjon
- seed
- valgt
- velge
- utvalg
- SELV
- sett
- innstilling
- flere
- Form
- deling
- Kort
- bør
- Vis
- vist
- Viser
- Signal
- signifikant
- betydelig
- lignende
- likheter
- Enkelt
- enkelhet
- ganske enkelt
- simulering
- siden
- situasjon
- Størrelse
- størrelser
- liten
- mindre
- Snapshot
- So
- løsning
- LØSE
- løse
- noen
- kilde
- Rom
- mellomrom
- spesifikk
- spesielt
- spesifisert
- fart
- bruke
- utgifter
- splittet
- kvadrat
- stablet
- Scene
- Standard
- Begynn
- starter
- Tilstand
- state-of-the-art
- uttalte
- Stater
- statistikk
- Trinn
- Steps
- stoppe
- oppbevare
- lagret
- butikker
- rett fram
- vellykket
- slik
- støtte
- Støtter
- syntetisk
- syntetiske data
- syntetisk
- system
- Systemer
- skreddersydd
- Ta
- tar
- ta
- Target
- Oppgave
- oppgaver
- teknikker
- terminal
- vilkår
- Det
- De
- Fremtiden
- Grafen
- informasjonen
- Staten
- deres
- Dem
- derved
- derfor
- Disse
- Tredje
- tre
- tredimensjonal
- Gjennom
- tid
- tidkrevende
- til
- sammen
- token
- tokenization
- symbolbaserte
- tokens
- også
- verktøykasse
- topp
- lommelykt
- Totalt
- tradisjonelle
- Tog
- trent
- Kurs
- Togene
- bane
- forvandlet
- behandle
- sant
- forstå
- Universe
- ubrukt
- Oppdater
- oppdatert
- oppdateringer
- oppdatering
- oppgradering
- us
- bruk
- bruke
- Brukere
- vanligvis
- verdi
- Verdier
- ulike
- versjon
- video
- videospill
- Besøk
- Besøk
- W
- Vei..
- Hva
- hvilken
- mens
- allment
- Wikipedia
- vil
- vinne
- med
- ord
- arbeidsstasjon
- verdt
- ville
- skriving
- skrevet
- X
- zephyrnet
- null