Foto door DeepMind on Unsplash
Matrixvermenigvuldiging is een fundamentele bewerking die in veel systemen wordt gebruikt, van neurale netwerken tot wetenschappelijke computerroutines. Het vinden van efficiënte en aantoonbaar correcte algoritmen voor matrixvermenigvuldiging kan een enorme impact hebben op het sneller en efficiënter maken van berekeningen, maar het is een zeer uitdagende taak. De ruimte van mogelijke algoritmen is enorm, en traditionele methoden voor het ontdekken van algoritmen, zoals door mensen ontworpen heuristieken of combinatorisch zoeken, zijn vaak niet optimaal.
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 alfanul. Deze agent is getraind om een spel voor één speler te spelen, TensorGame, waarbij het doel is om rekenkundig efficiënte algoritmen voor matrixvermenigvuldiging te ontdekken.
AlphaTensor is bijzonder goed in het omgaan met grote matrices door grote matrixvermenigvuldigingen op te splitsen in kleinere vermenigvuldigingen. Bovendien kan AlphaTensor worden gebruikt om state-of-the-art prestaties voor matrixvermenigvuldiging te bereiken, eenmaal nauwkeurig afgestemd op een specifiek hardwareapparaat.
AlphaTensor heeft een groot potentieel voor het versnellen van deep learning computing. Bij deep learning kunnen veel tijdrovende bewerkingen worden toegewezen aan matrixvermenigvuldigingen. Door AlphaTensor te gebruiken om deze bewerkingen te optimaliseren, kunnen de algehele prestaties van deep learning-modellen aanzienlijk worden verbeterd.
Onlangs heeft OpenAlphaTensor, de eerste open source-implementatie van AlphaTensor, werd uitgebracht, wat een revolutie teweeg zou kunnen brengen in de rekenkracht van deep learning-modellen.
Matrixvermenigvuldigingstensor
Voor niet-experts in optimalisatie van matrixvermenigvuldiging is het misschien niet eenvoudig te begrijpen hoe een bewerking zoals matrixvermenigvuldiging in een driedimensionale tensor kan worden afgebeeld. Ik zal proberen het in eenvoudige woorden en met voorbeelden uit te leggen.
Laten we het product C = A*B bekijken, waarbij voor de eenvoud zowel A als B vierkante matrices van grootte N zijn. De vermenigvuldigingsbewerking kan in kaart worden gebracht in een 3D-vormtensor (N^2, N^2, N^2). De eerste tensordimensie vertegenwoordigt de afgeplatte matrix A, de tweede dimensie de afgeplatte matrix B en de derde dimensie de afgeplatte matrix C.
De tensor heeft alleen binaire waarden (1 of 0) voor elk item. Merk op dat de tensor de vermenigvuldigingsoperatie vertegenwoordigt, dus het is onafhankelijk van de waarden van de matrices A en B.
Elke invoer van de tensor komt overeen met de coëfficiënt van de bewerking. Om bijvoorbeeld C[1,1] te berekenen, is het nodig om zowel A[1,1] als B[1,1] te vermenigvuldigen. Daarom zal de tensorinvoer [0,0,0], die overeenkomt met A[1,1], B[1,1] en C[1,1], waarde 1 hebben. Om C[1,1 te berekenen ,2,1], A[1] is niet nodig. De tensorrij T[N+0, :, XNUMX] zal dus alleen nullen bevatten.
De onderstaande afbeelding toont een voorbeeld van een tensor voor N=2.
Image from DeepMind's papier gepubliceerd NATUUR
Zoals weergegeven in (b) en (c) in de bovenstaande afbeelding, is het mogelijk om een algoritme te implementeren voor het berekenen van het product met behulp van een ontleding van de 3D-tensor. Meer specifiek kan onderstaand algoritme gebruikt worden om een tensorontleding (de matrices U, V, W) om te zetten in een matrixvermenigvuldigingsalgoritme.
Meta-algorithm parameterized for computing the matrix product C=AB introduced in DeepMind's papier
Het Tensorspel
Het probleem van het vinden van efficiënte algoritmen voor matrixvermenigvuldiging is buitengewoon uitdagend, omdat het aantal mogelijke algoritmen dat overwogen moet worden veel groter is dan het aantal atomen in het universum, zelfs voor kleine gevallen van matrixvermenigvuldiging.
DeepMind zette dit probleem om in een spel voor één speler en noemde het de TensorGame. In dit spel kiest de speler hoe hij verschillende ingangen van matrices combineert om ze te vermenigvuldigen. Er wordt een score toegekend op basis van het aantal bewerkingen dat nodig is om het juiste vermenigvuldigingsresultaat te bereiken. Het spel eindigt wanneer de tensor nul is bereikt of wanneer het maximale aantal zetten is gedaan. De uiteindelijke factorisatie wordt geëvalueerd op basis van een schatting van de resterende rang en bepaalde optimalisatiecriteria, zoals asymptotische tijdcomplexiteit of praktische looptijd.
De beginpositie in de TensorGame komt overeen met de matrixvermenigvuldigingstensor die op willekeurige basis wordt uitgedrukt.
In elke stap t van het spel schrijft de speler drie vectoren op
, die de tensoren van rang 1 specificeert . De status van het spel wordt bijgewerkt door de door de speler geselecteerde vectoren af te trekken:
WAAR
is de matrixvermenigvuldigingstensor.Als het spel eindigt in p stappen, betekent dit dat de Matrix Vermenigvuldigingstensor
kan worden ontleed in p rang-1 tensoren , dwz het heeft ten minste rang p.De TensorGame kan dan worden geïnterpreteerd als een rangontledingsalgoritme en AlphaTensor kan worden gezien als een algoritme voor het schatten van de rangorde van de tensor.
AlphaTensor-architectuur
Tot nu toe hebben we geleerd over de TensorGame en verduidelijkt hoe de oplossing kan worden gezien als een algoritme voor matrixvermenigvuldiging. Laten we nu eens kijken naar de belangrijkste concepten van AlphaTensor, het algoritme dat voor het spel wordt gebruikt.
AlphaTensor-architectuur is in feite een encoder-decoder Transformer-architectuur waarbij:
- de encoder neemt als invoer de spelstatus , de n eerdere acties die door het model zijn ondernomen (meestal n=7) en de tijdsindex t van de huidige actie. Informatie wordt op elkaar gestapeld in een tensor met vorm (n+1, N^2, N^2, N^2). Deze tensor wordt vervolgens hervormd en getransformeerd (met behulp van drie lineaire lagen) in een vorm tensor (N^2, N^2, c) waarbij c de binnenste dimensie van het model is.
- de decoder genereert de n_steps-acties op basis van de ingebedde vector die door de encoder wordt gegeven op een autoregressieve manier. Elke actie komt overeen met een fiche van de drieling die een van de drielingen vertegenwoordigt die de speltensor ontbinden (dwz de rangorde verminderen)
Het model wordt getraind door back-propagation en model acting af te wisselen. Modelacteren wordt gebruikt om gegevens te genereren die vervolgens worden gebruikt om het model te trainen. In de praktijk wordt het model getraind met een mix van synthetisch gegenereerde data en data die het model tijdens het acteren genereert. De acteerstap wordt gedaan door een 3D-tensor te nemen die overeenkomt met een matrixbewerking en daarop n_actors-spellen te spelen. Elke acteur speelt een spel op standaardbasis of op alternatieve basis (de verandering van basis wordt toegepast met een bepaalde waarschijnlijkheid). De resultaten worden vervolgens verzameld en kunnen worden gebruikt in de trainingsstap met de synthetische gegevens.
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.
Tijdens het trainen van het model is de beloning eigenlijk een negatieve beloning (penalty). De absolute waarde neemt toe met elke extra stap die nodig is om het spel op te lossen. Als het model m stappen zet om een TensorGame op te lossen, is de beloning die bij het spel hoort r=-m. Als het model de TensorGame niet in stappen van max_rank kan oplossen, wordt de beloning berekend door de rang van de resterende tensor te schatten. De rangorde wordt geschat als de som van de rangordes van de matrices die de tensor vormen. De schatting is een bovengrens van de ware rang van de tensor.
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 λ is een door de gebruiker gespecificeerde coëfficiënt.
Versnellingen (%) van door AlphaTensor ontdekte algoritmen op maat gemaakt voor een GPU en een TPU, ontleend aan DeepMind's paper. Versnellingen worden gemeten ten opzichte van standaard (bijv. cuBLAS voor de GPU) matrixvermenigvuldiging op dezelfde hardware en vergeleken met de Strassen-kwadraat algoritme. Bron: DeepMind.
Ik heb onlangs uitgebracht OpenAlphaTensor, de eerste open source-implementatie van AlphaTensor. In deze paragraaf zal ik de implementatie doornemen. Zoals we eerder hebben besproken, is de AlphaTensor-architectuur redelijk eenvoudig, gebaseerd op een standaardtransformator met een encoder-decoderarchitectuur. De meest interessante componenten van AlphaTensor zijn de eerste laag in het encodergedeelte en de manier waarop de acties worden gesampled.
Laten we beginnen met de eerste coderingslaag.
# 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
In het bovenstaande fragment laten we zien hoe de invoertensor wordt ontleed in drie tensoren, die vervolgens worden gebruikt als query-, sleutel- en waarde-invoer van de transformatorlaag.
- Over de drie tensordimensies die de afgeplatte matrices (A, B, C) vertegenwoordigen, wordt de invoertensor langs elke dimensie afgeplat samen met de dimensie die de vorige acties vertegenwoordigt. Op deze manier is in elke afgeplatte kopie van de invoertensor de geselecteerde dimensie een aggregatie van de laatste T-1-waarden en de werkelijke waarde, voor alle S-waarden van de geselecteerde dimensie, waarbij S=N^2. Filosofisch gezien is het alsof we voor elke dimensie focussen op wat er gebeurde in de voorgaande acties in die dimensie.
- De scalairen worden afgebeeld in drie verschillende ruimtes van dimensie S ^ 2, en vervolgens hervormd om te worden samengevoegd met de tensoren die op het vorige punt zijn verkregen. Conceptueel worden de scalairen toegewezen aan een inbeddingsruimte van dimensie S ^ 2, en vervolgens wordt de ingebedde informatie opgedeeld in S-vectoren en op elkaar gestapeld, vergelijkbaar met wat er gebeurt met tekst wanneer tokenized.
- Scalaire tokens worden aaneengeschakeld met de geherstructureerde invoertensor en vervolgens als invoer gegeven aan een lineaire laag voor het in kaart brengen van de focusinformatie scalairen + kanaalgeschiedenis in de interne dimensie van het model.
Deze drie stappen kunnen worden geïnterpreteerd als een manier om het model zowel informatie over de scalaire waarden (zoals in de TensorGame-tijdstap) als de focus op de voorgaande acties voor elk kanaal te geven.
Met betrekking tot de manier waarop de acties worden geproduceerd, is het interessant om op te merken dat AlphaTensor als uitvoer het triplet u, v, w genereert, dat tot doel heeft de tensorrang te verminderen. De drie vectoren hebben grootte S en aangezien ze aaneengeschakeld zijn, moet het model een vector van grootte 3*S produceren. AlphaTensor is getraind met een RL-algoritme, dus alle mogelijke acties moeten worden uitgedrukt in termen van kansen in een opgesomde ruimte, dwz het model produceert een kans over de verschillende acties. Dit betekent dat elke vector in de 3S-ruimte moet worden toegewezen aan een andere actie. Dit resulteert in een actieruimte van grootte |F|^(3S), waarbij |F| is het aantal verschillende waarden dat het element u, v, w kan aannemen. Gewoonlijk zijn de waarden beperkt tot (-2, -1, 0, 1, 2), wat resulteert in een kardinaliteit van 5 elementen.
Hier komt een grote uitdaging: om de actiekansen te genereren voor een matrixproduct van matrices van grootte 5, hebben we een geheugen nodig van 5 ^ 75 * 4 bytes, wat neerkomt op ~ 10 ^ 44 GB geheugen. Het is duidelijk dat we zo'n grote actieruimte niet aankunnen.
Hoe lossen we het probleem op? Om de geheugenvoetafdruk van de actiekansen te verkleinen, kunnen we de triplets splitsen in kleinere stukken, ze "tokeniseren" en de stukken behandelen als gegenereerde tokens in de transformatorarchitectuur, dwz de tokens worden gegeven als input aan de decoder in een auto-regressieve manier. In het bovenstaande voorbeeld kunnen we de triplets opsplitsen in 15 stukken, waardoor het geheugenverbruik wordt teruggebracht tot 15 * 5 ^ (75/15) * 4, ofwel 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), )
Hierboven tonen we het codefragment voor het genereren van de volledige actie. In de code bevat self.core de decoderlaag en vertegenwoordigt de tensor e de uitvoer van de encoderlaag. Nul kan worden beschouwd als de token in NLP-modellen en de n_steps-acties die de n_steps-chunks vertegenwoordigen, worden op een progressieve manier gegenereerd.
Het model retourneert drie grootheden:
- De gegenereerde acties
- De waarschijnlijkheid die is gekoppeld aan de volledige actie
- De logs die zijn geproduceerd voor het genereren van de eerste actie (het eerste deel) die zal worden gebruikt voor het berekenen van de modelwaarde.
Het is de moeite waard om een paar woorden te besteden aan de parameter n_samples. De parameter wordt gebruikt voor de acteerstap en stelt het model in staat om verschillende versies van de tripletten te genereren die vervolgens worden gebruikt voor het verkennen van de actieruimte in het Monte Carlo Tree Search-algoritme dat wordt gebruikt in het acteerproces. De verschillende acties van n_samples worden bemonsterd volgens het beleid dat door het model wordt gegenereerd.
Handelende stap
Het meest lastige deel van het hele algoritme is waarschijnlijk de Acting-stap die wordt gebruikt voor het oplossen van de TensorGame. Het algoritme wordt niet diep uitgelegd in het AlphaTensor-artikel, omdat het gebaseerd is op verschillende eerdere artikelen van DeepMind die zojuist zijn geciteerd en als bekend worden gegeven. Hier reconstrueer ik alle ontbrekende stukjes en leg ik stap voor stap onze implementatie uit.
We kunnen de acteerstappen in drie verschillende componenten indelen:
- De boomzoektocht in Monte Carlo
- De spelsimulatie
- De verbeterde beleidsberekening
Laten we ze één voor één analyseren.
Boom zoeken in Monte Carlo (MCTS)
Monte Carlo Tree Search (MCTS) is een veelgebruikte kunstmatige-intelligentietechniek voor het spelen van games, met name in bordspellen en videogames. Het algoritme creëert een spelboom die mogelijke zetten en uitkomsten simuleert en gebruikt willekeurige steekproeven om de verwachte beloning voor elke zet te evalueren. Het algoritme selecteert vervolgens iteratief de zet met de hoogste verwachte beloning en simuleert de resultaten totdat deze een terminale status of een gespecificeerde stopconditie bereikt. De simulaties worden gebruikt om de winkans voor elke zet in te schatten en het besluitvormingsproces te begeleiden. Het is aangetoond dat MCTS effectief is in complexe spellen waarbij het aantal mogelijke zetten en uitkomsten groot is, en het is gebruikt in succesvolle AI-systemen voor het spelen van games, zoals AlphaGo.
In AlphaTensor wordt een aangepaste versie van het originele MCTS gebruikt. In het bijzonder, in plaats van de actie willekeurig te selecteren uit de hele actieruimte, wordt de actie geselecteerd uit een subset die rechtstreeks door het model wordt gegenereerd (via de eerder gepresenteerde n_samples). De correctie op de beleidsupgrade wordt vervolgens toegepast in de stap Verbeterd beleid.
In onze implementatie hebben we besloten om alle informatie over de Monte-Carlo-boom in een woordenboek te bewaren met als sleutel de hash-versie van de TensorGame-status en als waarden de informatie die aan de staat zelf is gekoppeld. Elke Monte-Carlo-stap begint bij een knooppunt en simuleert n_sim-minigames, waarbij de toekomst wordt verkend met een horizon van 5 zetten. Als het knooppunt al is verkend in eerdere simulaties, wordt n_sim aangepast rekening houdend met het aantal eerdere verkenningen. Voor elk knooppunt wordt het aantal bezoeken opgeslagen in de N_s_a-tensor, aangezien deze tensor het aantal bezoeken per onderliggende actie van het knooppunt bevat (onder de door het model bemonsterde).
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
De bovenstaande code toont onze implementatie van het algoritme. Voor de eenvoud van de code wordt de beleidscorrectie uitgevoerd in de functie simuleren_spel.
Spelsimulatie
De functie simulation_game is verantwoordelijk voor het verkennen van de boom die is samengesteld uit knooppunten die een bepaalde toestand van de TensorGame vertegenwoordigen. Het voert het model ook uit wanneer een bladknooppunt wordt aangetroffen en het slaat alle knooppuntinformatie op in het state_dict-woordenboek. Laten we de implementatie eens nader bekijken:
@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)
Elke simulatie bestaat uit drie delen:
- Selectie
- Uitbreiding
- backup
In het selectiegedeelte wordt de simulatie uitgevoerd op de reeds gegenereerde boomknooppunten en wordt het volgende knooppunt geselecteerd met behulp van de volgende functie:
def select_future_state( possible_states: List[torch.Tensor], q_values: torch.Tensor, N_s_a: torch.Tensor, repetitions: Dict[int, list], c_1: float = 1.25, c_2: float = 19652, return_idx: bool = False,
) -> torch.Tensor: """Select the future state maximizing the upper confidence bound."""
# q_values (1, K, 1) pi = torch.tensor( [ len(repetitions[i]) for i in range(len(possible_states)) if i in repetitions ] ).to(q_values.device) ucb = q_values.reshape(-1) + pi * torch.sqrt( torch.sum(N_s_a) / (1 + N_s_a) ) * (c_1 + torch.log((torch.sum(N_s_a) + c_2 + 1) / c_2)) if return_idx: return ucb.argmax() return possible_states[ucb.argmax()]
In de praktijk is de actie die de ucb-functie maximaliseert:
voor de gegeven status is geselecteerd. Hierin staat Q voor de door het model gegenereerde Q-waarden en staat π voor de willekeurige verdeling over de acties die zijn bemonsterd met behulp van het modelbeleid. N(s, a) vertegenwoordigt het aantal bezoeken van het knooppunt aan actie a van knooppunt s.
Zodra de selectiefase een bladknooppunt bereikt en de simulatie geen eindvoorwaarde heeft bereikt (in termen van maximale verkenning, dwz toekomstige horizon of einde van het spel), wordt het model vervolgens gebruikt voor het selecteren van n_samples alternatieve knooppunten (ze zullen bladknooppunten zijn). knooppunten in de opeenvolgende iteratie). Dit wordt de uitbreidingsfase genoemd, omdat er nieuwe knooppunten aan de boom worden toegevoegd. Vervolgens wordt er in de huidige simulatie geen verder knooppunt verkend, maar wordt het blad q_value naar de volgende simulatiestap gestuurd: de back-up.
Back-up is de laatste fase van elke simulatie. Als tijdens de back-up het bladknooppunt een terminale toestand was, wordt de uiteindelijke beloning berekend; anders wordt de blad q-waarde gebruikt als een geschatte beloning. Vervolgens wordt de beloning teruggepropageerd op het simulatietraject, waarbij zowel de toestanden q_values worden bijgewerkt als de bezoekteller N (s, a) wordt bijgewerkt. In het onderstaande fragment laten we de code zien voor de back-propagation van de beloning.
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
Verbeterde beleidsberekening
Zodra alle simulaties zijn uitgevoerd en de MCTS een interessante momentopname biedt van de nabije toekomst, is het tijd om het beleid dat is gekoppeld aan de voorspelde knooppunten bij te werken en terug te sturen, zodat ze tijdens de training kunnen worden gebruikt. Het verbeterde beleid, volgens de werkwijze beschreven in Hubertus et al, wordt gebruikt voor het beheer van grote actieruimten. Voor een kleine zoekruimte is het tijdens MCTS zelfs mogelijk om willekeurig een actie uit de actieruimte te nemen en de impact ervan te evalueren. Een vergelijkbare aanpak in een veel grotere actieruimte zou ertoe leiden dat alle trajecten uiteenlopen in verschillende paden en er zou een oneindig aantal trajecten nodig zijn om zinvolle statistieken te krijgen en vervolgens het beleid bij te werken. Aangezien we hier sample-MCTS gebruiken om de spreiding te vermijden, d.w.z. n_samples-acties worden gesampled overeenkomstig het modelbeleid en vervolgens selecteert MCTS slechts een van de gesamplede acties tijdens het verkennen van de boom, we moeten rekening houden met de sample-correctie bij het berekenen het laatste bijgewerkte beleid dat zal worden gebruikt tijdens het trainen van het model.
In de praktijk wordt het verbeterde beleid berekend als
WAAR
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 op dat we in onze implementatie, nadat we het beleid van de N_s_a-tensor hebben berekend, het terug moeten koppelen aan de oorspronkelijke actie-tensor. In feite houdt N_s_a alleen rekening met de acties die door het model worden bemonsterd, terwijl het uiteindelijke beleid ook waarschijnlijkheden moet bevatten voor de niet-verkende acties.
Verschillen met betrekking tot het ChatGPT-trainingsalgoritme
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 is een fijnafstemmingstechniek die wordt gebruikt om taalmodellen af te stemmen op een reeks schriftelijke instructies. Het gebruikt menselijke voorkeuren als beloningssignaal om het model te verfijnen, waardoor het gedrag van het taalmodel wordt afgestemd op de verklaarde voorkeuren van een specifieke groep mensen, in plaats van op een breder begrip van 'menselijke waarden'.
MCTS daarentegen is een op bomen gebaseerd zoekalgoritme dat wordt gebruikt om de optimale zetten in games te bepalen. Het simuleert mogelijke zetten en werkt de waarden van elke zet bij op basis van hun resultaten, en begeleidt de selectie van de beste zet.
RLHF verzamelt gegevens van door mensen geschreven demonstraties en door mensen gelabelde vergelijkingen tussen AI-modellen, en traint een beloningsmodel om de voorkeuren van een bepaalde groep mensen te voorspellen. Het beloningsmodel wordt vervolgens gebruikt om de AI-modellen te verfijnen. MCTS daarentegen gebruikt simulaties en evaluaties om de beste beslissing te bepalen.
Hoewel het verschillende benaderingen zijn, hebben RLHF en MCTS ook overeenkomsten. Beide kunstmatige-intelligentietechnieken maken gebruik van methoden voor besluitvorming en probleemoplossing, en beide gebruiken een trial-and-error-benadering om verschillende opties te verkennen en beslissingen te nemen op basis van beschikbare informatie. Beide zijn ook iteratieve processen die in de loop van de tijd verbeteren naarmate er meer informatie en ervaring wordt verzameld.
De keuze tussen RLHF en MCTS is afhankelijk van de uit te voeren taak. RLHF is ideaal wanneer er geen duidelijke maatstaf is voor het evalueren van de modelprestaties, terwijl MCTS effectief is gebleken bij spelachtige taken waarbij kennis en verkenning van de toekomst het model een aanzienlijk voordeel geven.
Code-optimalisatie voor AlphaTensor-training
Het implementeren van het AlphaTensor-trainingsalgoritme vereist het vinden van het perfecte compromis tussen trainingssnelheid en geheugenverbruik. Zoals te zien is in de sectie Model, kan het simpelweg overwegen van de actietokenisatie veel geheugen besparen, maar een al te agressieve vermindering van de actieruimte kan leiden tot zowel verminderde nauwkeurigheid als tragere prestaties. Dit laatste gebeurt omdat alle tokens achtereenvolgens op een autoregressieve manier worden gegenereerd door de modeldecoder. Daarom groeit de inferentietijd lineair met het aantal fiches per actie zodra de softmax op het actieveld niet meer het knelpunt is.
Bij het opzetten van de AlphaTensor-training werden de grootste moeilijkheden gevonden in het omgaan met het acteerproces. Als de tensoren niet in het juiste formaat worden opgeslagen, kan de MCTS gemakkelijk een ongecontroleerde groei van het geheugengebruik veroorzaken. Aan de andere kant, als het aantal opgeslagen tensoren tijdens elke simulatie te veel wordt verminderd, kan de MCTS een oneindige hoeveelheid tijd besteden aan het opnieuw berekenen van de vereiste toestanden.
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 bytes voor elk knooppunt in de grafiek. Nu, aangezien elke simulatie in de uitbreidingsfase een nieuw knooppunt genereert (en n_sim=200), zouden we een uiteindelijk geheugenverbruik van 200 hebben328252525*4 = 3.2 GB alleen voor het eerste MCTS-knooppunt. In het slechtste geval zou dit tijdens het verkennen van acterende max_rank-knooppunten (waarbij max_rank = 150) resulteren in een totaal geheugenverbruik van 150 * 3.2 GB = 480 GB in RAM-geheugen (of GPU-geheugen als alle tensoren op de GPU waren opgeslagen) . We hebben de training uitgevoerd op ons werkstation met 128 GB RAM en 48 GB GPU-geheugen, dus we moesten het geheugengebruik verminderen.
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.
Nu OpenAlphaTensor open source is, openen zich verschillende opwindende wegen voor verdere ontwikkeling.
Een natuurlijke vooruitgang is de fijnafstemming van OpenAlphaTensor op doelhardwareapparaten. Dit zal naar verwachting leiden tot zeer competitieve rekenprestaties. Ik zal meer publiceren over de prestaties van OpenAlphaTensor op verschillende hardware op GitHub. Op het moment van schrijven van dit artikel volgde OpenAlphaTensor een training.
Een andere belangrijke vooruitgang zou de ondersteuning zijn voor compilatie op afstand, waardoor gebruikers algoritmen kunnen bouwen die zijn geoptimaliseerd voor edge-apparaten. Dit kan worden bereikt door het OpenAlphaTensor-model op een server op te slaan, terwijl het matrixvermenigvuldigingsalgoritme op verschillende hardware wordt geëvalueerd.
Het kan ook belangrijk zijn om de ondersteuning voor verschillende compilers uit te breiden om de op latentie gebaseerde beloningscorrectie te berekenen. Verschillende compilers kunnen leiden tot verschillende geoptimaliseerde algoritmen op bepaalde hardware. De DeepMind-paper liet bijvoorbeeld veelbelovende resultaten zien met behulp van JAX en de XLA-compiler op TPU en Nvidia GPU's. Het zou interessant zijn om dit te evalueren met behulp van NCCL op Nvidia of LLVM op CPU's.
Ten slotte blijft het uitbreiden van het model en het trainingsalgoritme om grotere matrixafmetingen te ondersteunen een grote openstaande uitdaging. Momenteel ondersteunt OpenAlphaTensor een maximale matrixgrootte van 5, maar het kan worden toegepast door grotere matrixvermenigvuldigingen op te splitsen in groepen van kleine MM's met een grootte kleiner dan 5. Deze benadering is suboptimaal en voert de reductie rechtstreeks uit op de grote tensor die overeenkomt met de volledige MM zou in theorie tot betere resultaten kunnen leiden.
Diego Fiori is the CTO of Nebuly AI, a company committed to making AI optimization part of every developer's toolkit.
- Door SEO aangedreven content en PR-distributie. Word vandaag nog versterkt.
- Platoblockchain. Web3 Metaverse Intelligentie. Kennis versterkt. Toegang hier.
- Bron: 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
- :is
- ][P
- $UP
- 1
- 3d
- 8
- a
- in staat
- Over
- boven
- absoluut
- versnellen
- Volgens
- dienovereenkomstig
- Account
- nauwkeurigheid
- Bereiken
- bereikt
- Actie
- acties
- werkelijk
- toegevoegd
- Extra
- gecorrigeerd
- aangenomen
- bevorderen
- Voordeel
- Na
- Agent
- aggregatie
- agressief
- AI
- AI-systemen
- wil
- algoritme
- algoritmen
- Alles
- Het toestaan
- toestaat
- alleen
- al
- alternatief
- onder
- bedragen
- analyseren
- en
- Nog een
- toegepast
- nadering
- benaderingen
- architectuur
- ZIJN
- dit artikel
- kunstmatig
- kunstmatige intelligentie
- AS
- toegewezen
- geassocieerd
- At
- geautomatiseerde
- Beschikbaar
- het vermijden van
- terug
- backup
- gebaseerde
- Eigenlijk
- basis
- BE
- omdat
- wordt
- vaardigheden
- wezen
- onder
- criterium
- BEST
- Betere
- tussen
- Verder
- boord
- Gezelschapsspelletjes
- gebonden
- bredere
- BT
- bouw
- bebouwd
- by
- Dit betekent dat we onszelf en onze geliefden praktisch vergiftigen.
- CAN
- kan niet
- geval
- Veroorzaken
- veroorzaakt
- zeker
- uitdagen
- uitdagend
- verandering
- Kanaal
- ChatGPT
- kind
- keuze
- het kiezen van
- uitgekozen
- aangehaald
- duidelijk
- duidelijk
- code
- verzamelt
- combineren
- toegewijd
- Gemeen
- afstand
- vergeleken
- concurrerend
- complex
- ingewikkeldheid
- componenten
- samengesteld
- compromis
- berekening
- rekenkracht
- Berekenen
- computergebruik
- concepten
- conceptueel
- voorwaarde
- vertrouwen
- Overwegen
- beschouwd
- aangezien
- beschouwt
- consumptie
- bevat
- voortzetten
- contrast
- geconverteerd
- Kern
- Overeenkomend
- komt overeen
- kon
- Counter
- CPU
- creëert
- Wij creëren
- criteria
- CTO
- Actueel
- Op dit moment
- gegevens
- omgang
- beslissen
- beslist
- beslissing
- Besluitvorming
- beslissingen
- deep
- diepgaand leren
- DeepMind
- afhankelijk
- beschreven
- Bepalen
- Ontwikkelaar
- Ontwikkeling
- apparaat
- systemen
- DICT
- anders
- moeilijkheden
- Afmeting
- Afmeting
- direct
- Onthul Nu
- het ontdekken van
- besproken
- distributie
- Verdeeld
- beneden
- Val
- gedurende
- e
- elk
- Vroeger
- gemakkelijk
- rand
- effectief
- doeltreffend
- beide
- element
- geeft je de mogelijkheid
- ingebed
- eindigt
- verbeterde
- enorm
- genoeg
- toegang
- fout
- schatting
- geschat
- Ether (ETH)
- schatten
- geëvalueerd
- evalueren
- evaluaties
- Zelfs
- Alle
- voorbeeld
- voorbeelden
- opwindend
- uitvoering
- uitbreiding
- verwacht
- ervaring
- Verklaren
- uitgelegd
- exploits
- exploratie
- Verken
- Nagegaan
- Verkennen
- uitgedrukt
- verlengen
- verlenging
- uiterst
- tamelijk
- familie
- sneller
- feedback
- weinig
- Figuur
- finale
- het vinden van
- Voornaam*
- Vlotter
- Focus
- volgen
- volgend
- Footprint
- Voor
- formulier
- formaat
- formule
- gevonden
- oppompen van
- vol
- functie
- fundamenteel
- verder
- verdere ontwikkeling
- toekomst
- spel
- Spellen
- voortbrengen
- gegenereerde
- genereert
- het genereren van
- krijgen
- het krijgen van
- Geven
- gegeven
- Vrijgevigheid
- doel
- Goes
- goed
- GPU
- GPU's
- diagram
- groot
- Groep
- Groep
- Groeit
- gids
- hand
- Behandeling
- gebeurd
- gebeurt
- Hardware
- hardware apparaat
- hardware-apparaten
- Hebben
- met
- hier
- hoogst
- horizont
- Hoe
- How To
- Echter
- HTTPS
- reusachtig
- menselijk
- i
- ZIEK
- ideaal
- IDX
- beeld
- Impact
- uitvoeren
- uitvoering
- belangrijk
- verbeteren
- verbeterd
- in
- Laat uw omzet
- Verhoogt
- in toenemende mate
- onafhankelijk
- index
- informatie
- eerste
- invoer
- verkrijgen in plaats daarvan
- instructies
- Intelligentie
- interessant
- intern
- geïntroduceerd
- intuïtie
- IT
- herhaling
- HAAR
- zelf
- jpg
- KDnuggets
- Houden
- sleutel
- kennis
- bekend
- taal
- Groot
- groter
- Achternaam*
- Wachttijd
- laatste
- lagen
- Legkippen
- leiden
- geleerd
- leren
- Lijst
- Kijk
- op zoek
- lot
- gemaakt
- Hoofd
- groot
- maken
- maken
- beheer
- beheren
- veel
- kaart
- in kaart brengen
- Matrix
- Materie
- maximaal
- betekenis
- zinvolle
- middel
- lid
- Geheugen
- methode
- methoden
- metriek
- vermist
- mengsel
- model
- modellen
- gewijzigd
- module
- meer
- efficiënter
- Bovendien
- meest
- beweging
- beweegt
- vermenigvuldigd
- Naturel
- NATUUR
- Nabij
- noodzakelijk
- Noodzaak
- nodig
- negatief
- netwerken
- Neural
- neurale netwerken
- New
- volgende
- nlp
- knooppunt
- knooppunten
- niet-experts
- notie
- aantal
- Nvidia
- verkregen
- of
- Aanbod
- on
- EEN
- open
- open source
- OpenAI
- operatie
- Operations
- optimale
- optimalisatie
- Optimaliseer
- geoptimaliseerde
- Opties
- origineel
- Overige
- anders-
- uitgang
- totaal
- Papier
- papieren
- parameter
- deel
- bijzonder
- vooral
- onderdelen
- Mensen
- prestatie
- uitvoerend
- fase
- stukken
- Plato
- Plato gegevensintelligentie
- PlatoData
- Spelen
- speler
- spelen
- punt
- beleidsmaatregelen door te lezen.
- beleidsmaatregelen
- positie
- mogelijk
- potentieel
- energie
- PRAKTISCH
- praktijk
- voorspellen
- voorspeld
- voorkeuren
- gepresenteerd
- vorig
- waarschijnlijkheid
- waarschijnlijk
- probleem
- processen
- produceren
- geproduceerd
- Product
- progressie
- progressief
- veelbelovend
- voorgestelde
- aantoonbaar
- bewezen
- publiceren
- gepubliceerde
- RAM
- willekeurige
- gelederen
- liever
- bereikt
- Bereikt
- onlangs
- verminderen
- Gereduceerd
- vermindering
- verfijnd
- versterking van leren
- uitgebracht
- resterende
- stoffelijk overschot
- opmerkelijk
- niet vergeten
- vanop
- herhaald
- vertegenwoordigen
- vertegenwoordigt
- nodig
- vereist
- verantwoordelijk
- begrensd
- resultaat
- verkregen
- Resultaten
- terugkeer
- Retourneren
- revolutioneren
- Belonen
- RIJ
- rt
- lopen
- s
- dezelfde
- Bespaar
- besparing
- scenario
- scenario's
- schema
- Ontdek
- Tweede
- sectie
- zaad
- gekozen
- selecteren
- selectie
- ZELF
- reeks
- het instellen van
- verscheidene
- Vorm
- delen
- Bermuda's
- moet
- tonen
- getoond
- Shows
- Signaal
- aanzienlijke
- aanzienlijk
- gelijk
- overeenkomsten
- Eenvoudig
- eenvoud
- eenvoudigweg
- simulatie
- sinds
- situatie
- Maat
- maten
- Klein
- kleinere
- Momentopname
- So
- oplossing
- OPLOSSEN
- Het oplossen van
- sommige
- bron
- Tussenruimte
- ruimten
- specifiek
- specifiek
- gespecificeerd
- snelheid
- besteden
- Uitgaven
- spleet
- vierkant
- gestapeld
- Stadium
- standaard
- begin
- starts
- Land
- state-of-the-art
- bepaald
- Staten
- statistiek
- Stap voor
- Stappen
- stoppen
- shop
- opgeslagen
- winkels
- eenvoudig
- geslaagd
- dergelijk
- ondersteuning
- steunen
- synthetisch
- synthetische gegevens
- synthetisch
- system
- Systems
- op maat gemaakt
- Nemen
- neemt
- het nemen
- doelwit
- Taak
- taken
- technieken
- terminal
- termen
- dat
- De
- De toekomst
- De grafiek
- de informatie
- De Staat
- hun
- Ze
- daarbij
- daarom
- Deze
- Derde
- drie
- driedimensionaal
- Door
- niet de tijd of
- tijdrovend
- naar
- samen
- teken
- tokenization
- getokeniseerd
- tokens
- ook
- toolkit
- top
- fakkel
- Totaal
- traditioneel
- Trainen
- getraind
- Trainingen
- treinen
- traject
- getransformeerd
- behandelen
- waar
- begrijpen
- Universum
- ongebruikt
- bijwerken
- bijgewerkt
- updates
- bijwerken
- upgrade
- us
- Gebruik
- .
- gebruikers
- doorgaans
- waarde
- Values
- divers
- versie
- Video
- video games
- Bezoek
- Bezoeken
- W
- Manier..
- Wat
- welke
- en
- wijd
- Wikipedia
- wil
- het winnen van
- Met
- woorden
- werkstation
- waard
- zou
- het schrijven van
- geschreven
- X
- zephyrnet
- nul