Foto: Deepmind on Unsplash
Matrismultiplikation är en grundläggande operation som används i många system, från neurala nätverk till vetenskapliga beräkningsrutiner. Att hitta effektiva och bevisligen korrekta algoritmer för matrismultiplikation kan ha en enorm inverkan på att göra beräkningen snabbare och mer effektiv, men är en mycket utmanande uppgift. Utrymmet för möjliga algoritmer är enormt, och traditionella metoder för att upptäcka algoritmer, såsom mänskligt utformad heuristik eller kombinatorisk sökning, är ofta suboptimala.
Deepminds nyligen föreslagna AI-baserade lösning för automatiserad sökning går långt utöver mänsklig intuition. Lösningen består av en djup förstärkningsinlärningsagent kallad AlphaTensor, byggd ovanpå Alpha Zero. Denna agent är utbildad för att spela ett enspelarspel, TensorGame, där målet är att upptäcka beräkningseffektiva algoritmer för matrismultiplikation.
AlphaTensor är särskilt bra på att hantera stora matriser genom att bryta ner stora matrismultiplikationer till mindre multiplikationer. Dessutom kan AlphaTensor användas för att uppnå toppmodern prestanda för matrismultiplikation när den väl har finjusterats på en specifik hårdvaruenhet.
AlphaTensor har stor potential för att accelerera djupinlärning av datoranvändning. Vid djupinlärning kan många tidskrävande operationer mappas till matrismultiplikationer. Genom att använda AlphaTensor för att optimera dessa operationer kan den övergripande prestandan hos modeller för djupinlärning förbättras avsevärt.
Nyligen öppnade OpenAlphaTensor, den första implementeringen av öppen källkod av AlphaTensor, släpptes, vilket skulle kunna revolutionera beräkningskraften hos modeller för djupinlärning.
Matrix Multiplikation Tensor
För icke-experter på matrismultiplikationsoptimering kanske det inte är enkelt att förstå hur en operation som matrismultiplikation kan mappas i en tredimensionell tensor. Jag ska försöka förklara det i enkla ord och med exempel.
Låt oss betrakta produkten C = A*B, där både A och B för enkelhets skull är kvadratiska matriser av storlek N. Multiplikationsoperationen kan mappas i en 3D-tensor av form (N^2, N^2, N^2). Den första tensordimensionen representerar den tillplattade matrisen A, den andra dimensionen den tillplattade matrisen B och den tredje dimensionen den tillplattade matrisen C.
Tensorn har bara binära värden (antingen 1 eller 0) för varje post. Observera att tensorn representerar multiplikationsoperationen, så den är oberoende av värdena på matriserna A och B.
Varje inmatning av tensorn motsvarar operationens koefficient. Till exempel, för att beräkna C[1,1] är det nödvändigt att multiplicera både A[1,1] och B[1,1]. Därför kommer tensorposten [0,0,0], som motsvarar A[1,1], B[1,1] och C[1,1], att ha värdet 1. Däremot för att beräkna C[1,1 ,2,1], A[1] behövs inte. Tensorraden T[N+0, :, XNUMX] kommer alltså endast att innehålla nollor.
Bilden nedan visar ett exempel på en tensor för N=2.
Bild från DeepMind's papper som publicerades i Natur
Som visas i (b) och (c) i figuren ovan är det möjligt att implementera en algoritm för att beräkna produkten med hjälp av en nedbrytning av 3D-tensorn. Mer specifikt kan algoritmen nedan användas för att omvandla en tensorupplösning (matriserna U, V, W) till en matrismultiplikationsalgoritm.
Metaalgoritm parametriserad för att beräkna matrisprodukten C=AB introducerad i DeepMinds papper
Tensorspelet
Problemet med att hitta effektiva algoritmer för matrismultiplikation är extremt utmanande eftersom antalet möjliga algoritmer att överväga är mycket större än antalet atomer i universum, även för små fall av matrismultiplikation.
DeepMind konverterade detta problem till ett enspelarspel och kallade det TensorGame. I det här spelet väljer spelaren hur man kombinerar olika matrisposter för att multiplicera dem. En poäng tilldelas baserat på antalet operationer som krävs för att uppnå rätt multiplikationsresultat. Spelet slutar när nolltensorn uppnås eller när det maximala antalet drag har gjorts. Den slutliga faktoriseringen utvärderas baserat på en uppskattning av den resterande rangordningen och vissa optimeringskriterier, såsom asymptotisk tidskomplexitet eller praktisk körtid.
Den initiala positionen i TensorGame motsvarar Matrix Multiplication Tensor uttryckt på någon slumpmässig basis.
I varje steg t i spelet skriver spelaren ner tre vektorer
, som specificerar rang-1-tensorerna . Spelets tillstånd uppdateras genom att subtrahera vektorerna som valts av spelaren:
var
är Matrix Multiplication Tensor.Om spelet slutar i p steg, betyder detta att Matrix Multiplication Tensor
kan delas upp i p rank-1 tensorer , dvs den har åtminstone rang p.TensorGame kan sedan tolkas som en rank-nedbrytningsalgoritm och AlphaTensor kan ses som en algoritm för att uppskatta rankningen av tensorn.
AlphaTensor-arkitektur
Hittills har vi lärt oss om TensorGame och klargjort hur dess lösning kan ses som en matrismultiplikationsalgoritm. Låt oss nu utforska huvudkoncepten för AlphaTensor, algoritmen som används för spelet.
AlphaTensor-arkitektur är i grunden en kodar-avkodar-transformatorarkitektur där:
- kodaren tar som indata speltillståndet , de n föregående åtgärderna som utförts av modellen (vanligtvis n=7), och tidsindex t för den aktuella åtgärden. Information staplas ihop i en tensor med form (n+1, N^2, N^2, N^2). Denna tensor omformas sedan och transformeras (med hjälp av tre linjära lager) till en tensor av form (N^2, N^2, c) där c är modellens inre dimension.
- avkodaren genererar n_steps-åtgärderna från den inbäddade vektorn som ges av kodaren på ett autoregressivt sätt. Varje åtgärd motsvarar en token av trillingarna representerar en av trillingarna som bryter ner speltensorn (dvs. minskar dess rang)
Modellen tränas genom att omväxlande backpropagation och modellagerande. Model acting används för att generera data som sedan används för att träna modellen. I praktiken tränas modellen med en blandning av syntetiskt genererad data och data som genereras av modellen under agerandet. Skådespelarsteget görs genom att ta en 3D-tensor som motsvarar en matrisoperation och spela n_actors-spel på den. Varje skådespelare spelar ett spel antingen på standardbasis eller på alternativ basis (förändringen av basen tillämpas med en given sannolikhet). Resultaten samlas sedan in och kan användas i träningssteget med den syntetiska datan.
Skådespelarsteget är baserat på AlphaZeros Monte Carlo Tree Search (MCTS), modifierat för att stödja stora actionutrymmen. Kort sagt, innan du väljer åtgärd, utforskas n_sims-vägar från modellutgången med en maximal framtida utforskning av 5 steg. Sannolikheterna som genereras av modellen justeras sedan med hänsyn till de genererade vägarna. Sedan väljs handlingen med de mest lovande framtidsvägarna för att fortsätta spelet.
När du tränar modellen är belöningen faktiskt en negativ belöning (straff). Dess absoluta värde ökar för varje ytterligare steg som krävs för att lösa spelet. Om modellen tar m steg för att lösa ett TensorGame, är belöningen förknippad med spelet r=-m. Om modellen inte kan lösa TensorGame i max_rank-steg, beräknas belöningen genom att uppskatta rankningen för den återstående tensorn. Rangen uppskattas som summan av rangorden av de matriser som utgör tensorn. Uppskattningen är en övre gräns för tensorns sanna rang.
Vid finjustering av modellen bör straffbelöningen vid terminaltillståndet också ta hänsyn till latensen för algoritmen som produceras av modellen. Belöningsformeln blir rt'=rt+λbt, där rt är belöningsschemat som beskrivits tidigare, bt är riktmärkesbelöningen (inte noll endast vid terminaltillståndet), och λ är en användarspecificerad koefficient.
Snabbare (%) av AlphaTensor-upptäckta algoritmer skräddarsydda för en GPU och en TPU, extraherad från DeepMinds papper. Speed-ups mäts i förhållande till standardmatrismultiplikation (t.ex. cuBLAS för GPU) på samma hårdvara och jämförs med Strassen-kvadratalgoritm. Källa: Deepmind.
Jag släppte nyligen OpenAlphaTensor, den första implementeringen av AlphaTensor med öppen källkod. I det här avsnittet kommer jag att gå igenom implementeringen. Som vi diskuterade tidigare är AlphaTensor-arkitekturen ganska okomplicerad, baserad på en standardtransformator med en kodar-avkodararkitektur. De mest intressanta komponenterna i AlphaTensor är det första lagret i kodardelen och hur åtgärderna samplas.
Låt oss börja med det första kodningsskiktet.
# 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 ovan visar vi hur ingångstensorn delas upp i tre tensorer, som sedan används som fråge-, nyckel- och värdeingångar för transformatorlagret.
- Över de tre tensordimensionerna som representerar de tillplattade matriserna (A, B, C), är ingångstensorn tillplattad längs varje dimension tillsammans med dimensionen som representerar de föregående åtgärderna. På detta sätt, i varje tillplattad kopia av inmatningstensorn, är den valda dimensionen en aggregering av de sista T-1-värdena och det faktiska värdet, för alla S-värden för den valda dimensionen, där S=N^2. Filosofiskt är det som om vi för varje dimension fokuserar på vad som hände i de tidigare handlingarna i den dimensionen.
- Skalärerna mappas i tre olika utrymmen med dimensionen S^2 och omformas sedan för att sammanlänkas med tensorerna som erhölls vid föregående punkt. Begreppsmässigt mappas skalärerna till ett inbäddningsutrymme med dimensionen S^2, och sedan delas den inbäddade informationen i S-vektorer och staplas ihop, liknande vad som händer med text när den tokeniseras.
- Skalära tokens sammanlänkas med den omstrukturerade ingångstensorn och ges sedan som indata till ett linjärt lager för att kartlägga fokusinformationen för skalärer+kanalhistorik i modellens interna dimension.
Dessa tre steg kan tolkas som ett sätt att ge modellen både information om skalärerna (som i TensorGame-tidssteget) och fokus på de tidigare åtgärderna för varje kanal.
När det gäller hur åtgärderna produceras är det intressant att notera att AlphaTensor genererar tripletten u, v, w som utdata, som syftar till att minska tensorrankningen. De tre vektorerna har storlek S och eftersom de är sammanlänkade måste modellen producera en vektor med storlek 3*S. AlphaTensor är tränad med en RL-algoritm, så alla möjliga åtgärder måste uttryckas i termer av sannolikheter i ett uppräknat utrymme, dvs modellen producerar en sannolikhet över de olika åtgärderna. Detta innebär att varje vektor i 3S-utrymmet bör mappas till en annan handling. Detta resulterar i ett handlingsutrymme med storleken |F|^(3S), där |F| är antalet olika värden som elementet u, v, w kan ta. Vanligtvis är värdena begränsade till (-2, -1, 0, 1, 2), vilket resulterar i en kardinalitet av 5 element.
Här kommer en stor utmaning: för att generera åtgärdssannolikheterna för en matrisprodukt av matriser av storlek 5 skulle vi behöva ett minne på 5^75 * 4 byte, vilket skulle innebära ~10^44 GB minne. Det är klart att vi inte kan hantera ett så stort handlingsutrymme.
Hur löser vi problemet? För att minska minnesavtrycket för handlingssannolikheterna kan vi dela upp trillingarna i mindre bitar, "tokenisera" dem och behandla bitarna som genererade tokens i transformatorarkitekturen, dvs. tokens ges som indata till avkodaren i en autoregressiv sätt. I exemplet ovan kan vi dela upp trillingarna i 15 bitar, vilket minskar minnesförbrukningen till 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), )
Ovan visar vi kodavsnittet för att generera hela åtgärden. I koden innehåller self.core avkodarskiktet och tensorn e representerar utdata från kodarskiktet. Noll kan betraktas som token i NLP-modeller och n_steps-åtgärderna som representerar n_steps-bitarna genereras på ett progressivt sätt.
Modellen returnerar tre kvantiteter:
- De genererade åtgärderna
- Sannolikheten för den fullständiga åtgärden
- De logit som skapas för att generera den första åtgärden (den första biten) som kommer att användas för att beräkna modellvärdet.
Det är värt att lägga några ord på parametern n_samples. Parametern används för handlingssteget och den tillåter modellen att generera olika versioner av trillingarna som sedan kommer att användas för att utforska handlingsutrymmet i Monte Carlo Tree Search-algoritmen som används i Acting-processen. De n_samples olika åtgärderna samplas enligt den policy som genereras av modellen.
Tillförordnad steg
Den mest knepiga delen av hela algoritmen är förmodligen det agerande steget som används för att lösa TensorGame. Algoritmen är inte djupt förklarad i AlphaTensor-artikeln, eftersom den är baserad på flera DeepMinds tidigare artiklar som bara citeras och ges som kända. Här kommer jag att rekonstruera alla de saknade bitarna och steg för steg förklara vår implementering.
Vi kan organisera skådespelarstegen i tre olika komponenter:
- Monte-Carlo Tree Search
- Spelsimuleringen
- Den förbättrade policyberäkningen
Låt oss analysera dem en efter en.
Monte-Carlo Tree Search (MCTS)
Monte Carlo Tree Search (MCTS) är en mycket använd teknik för artificiell intelligens för spel, särskilt i brädspel och videospel. Algoritmen skapar ett spelträd som simulerar potentiella drag och resultat och använder slumpmässigt urval för att utvärdera den förväntade belöningen för varje drag. Algoritmen väljer sedan iterativt draget med den högsta förväntade belöningen och simulerar utfall tills den når ett terminaltillstånd eller ett specificerat stopptillstånd. Simuleringarna används för att uppskatta sannolikheten att vinna för varje drag och vägleda beslutsprocessen. MCTS har visat sig vara effektivt i komplexa spel där antalet möjliga drag och utfall är stort, och det har använts i framgångsrika spelande AI-system, som AlphaGo.
I AlphaTensor används en modifierad version av original MCTS. I synnerhet, istället för att slumpmässigt välja åtgärden från hela åtgärdsutrymmet, väljs åtgärden bland en delmängd som genereras direkt av modellen (genom de n_samples som presenterades tidigare). Korrigeringen av policyuppgraderingen tillämpas sedan i steget förbättrad policyberäkning.
I vår implementering bestämde vi oss för att behålla all information om Monte-Carlo-trädet i en ordbok som har hash-versionen av TensorGame-tillståndet som nyckel och som värden informationen som är associerad med själva staten. Varje Monte-Carlo-steg startar från en nod och simulerar n_sim-minispel och utforskar framtiden med en horisont på 5 drag. Om noden redan har utforskats i tidigare simuleringar, justeras n_sim med hänsyn till antalet tidigare undersökningar. För varje nod lagras antalet besök i N_s_a-tensorn, eftersom denna tensor innehåller antalet besök per nodunderordnad åtgärd (bland de som samplades 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 ovan visar vår implementering av algoritmen. För enkelhetens skull utförs policykorrigeringen i simulate_game-funktionen.
Spelsimulering
Funktionen simulate_game är ansvarig för att utforska trädet som består av noder som representerar ett visst tillstånd i TensorGame. Den kör också modellen närhelst en lövnod påträffas och den lagrar all nodinformation i state_dict-ordboken. Låt oss ge en djup titt på dess genomförande:
@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)
Varje simulering är uppdelad i tre delar:
- Urval
- Expansion
- säkerhetskopiering
I urvalsdelen körs simuleringen på de redan genererade trädnoderna, och följande nod väljs med följande 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 praktiken är åtgärden som maximerar ucb-funktionen:
för det givna tillståndet är valt. Här representerar Q Q-värdena som genereras av modellen och π representerar den slumpmässiga fördelningen över de åtgärder som samplats med hjälp av modellpolicyn. N(s, a) representerar antalet besök av noden till åtgärd a från nod s.
När väl urvalsfasen når en lövnod, om simuleringen inte har nått ett terminaltillstånd (i termer av antingen maximal utforskning, dvs framtida horisont eller spelslut), används modellen för att välja n_samples alternativa noder (de kommer att vara löv noder i den successiva iterationen). Detta kallas expansionsfasen, eftersom nya noder läggs till i trädet. Sedan utforskas ingen ytterligare nod i den aktuella simuleringen, utan bladet q_value skickas till följande simuleringssteg: backupen.
Backup är det sista steget i varje simulering. Under backup, om lövnoden var ett terminaltillstånd, beräknas den slutliga belöningen; annars används bladets q-värde som en uppskattad belöning. Sedan sprids belöningen tillbaka på simuleringsbanan, varvid både tillstånden q_values uppdateras och besöksräknaren N(s, a) uppdateras. I utdraget nedan visar vi koden för återförökningen av belöningen.
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
Förbättrad policyberäkning
När alla simuleringar har körts och MCTS erbjuder en intressant ögonblicksbild av den närmaste framtiden är det dags att uppdatera policyn som är kopplad till de förutspådda noderna och returnera dem, så att de kan användas under träning. Den förbättrade policyn, enligt metoden som beskrivs i Hubert et al, används för att hantera stora actionutrymmen. Faktum är att för litet sökutrymme är det möjligt under MCTS att ta ett slumpmässigt urval av en åtgärd från åtgärdsutrymmet och utvärdera dess effekt. Ett liknande tillvägagångssätt i ett mycket större handlingsutrymme skulle leda till att alla banor skiljer sig åt i olika banor och det skulle behövas en oändlig mängd banor för att få meningsfull statistik och sedan uppdatera policyn. Eftersom vi här använder sample-MCTS för att undvika spridningen, dvs n_samples-åtgärder samplas i enlighet med modellpolicyn och sedan väljer MCTS bara en av de samplade åtgärderna medan vi utforskar trädet, måste vi ta hänsyn till sample-korrigeringen vid beräkning den slutgiltiga uppdaterade policyn som kommer att användas när modellen tränas.
I praktiken beräknas den förbättrade policyn som
var
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
Observera att i vår implementering efter att ha beräknat policyn från N_s_a-tensorn måste vi mappa tillbaka den till den ursprungliga handlingstensorn. Faktum är att N_s_a bara tar hänsyn till de åtgärder som samplats av modellen, medan den slutliga policyn måste innehålla sannolikheter även för de icke-utforskade åtgärderna.
Skillnader med avseende på ChatGPT-träningsalgoritm
AlphaTensor är den senaste medlemmen i AlphaGo/AlphaZero-familjen av metoder för artificiell intelligens från DeepMind. Dessa metoder är baserade på Monte Carlo Tree Search (MCTS) algoritm, som har förfinats och förbättrats av DeepMind för att hantera allt mer komplexa uppgifter. Ett annat AI-system, OpenAIs ChatGPT, som har orsakat mycket surr för sin enastående prestanda, tränades med ett annat tillvägagångssätt, kallat Reinforcement Learning with Human Feedback (RLHF).
RLHF är en finjusteringsteknik som används för att ställa in språkmodeller för att följa en uppsättning skriftliga instruktioner. Den använder mänskliga preferenser som en belöningssignal för att finjustera modellen, och därigenom anpassa språkmodellens beteende med de uttalade preferenserna för en specifik grupp människor, snarare än någon bredare uppfattning om "mänskliga värderingar".
Däremot är MCTS en trädbaserad sökalgoritm som används för att bestämma de optimala dragen i spel. Den simulerar potentiella drag och uppdaterar värdena för varje drag baserat på deras resultat, vilket vägleder valet av det bästa draget.
RLHF samlar in data från människoskrivna demonstrationer och människomärkta jämförelser mellan AI-modeller, och tränar en belöningsmodell för att förutsäga preferenser för en given grupp människor. Belöningsmodellen används sedan för att finjustera AI-modellerna. MCTS, å andra sidan, använder simuleringar och utvärderingar för att bestämma det bästa beslutet.
Även om de är olika tillvägagångssätt har RLHF och MCTS också likheter. Båda teknikerna för artificiell intelligens använder sig av beslutsfattande och problemlösningsmetoder, och båda använder sig av en trial-and-error-metod för att utforska olika alternativ och fatta beslut baserat på tillgänglig information. Båda är också iterativa processer som förbättras över tid när mer information och erfarenhet samlas in.
Valet mellan RLHF och MCTS beror på uppgiften. RLHF är idealiskt när det inte finns något tydligt mått för att utvärdera modellens prestanda, medan MCTS har visat sig vara effektivt i spelliknande uppgifter där kunskap och utforskande av framtiden ger modellen en betydande fördel.
Kodoptimering för AlphaTensor-träning
Att implementera AlphaTensor-träningsalgoritmen kräver att man hittar den perfekta kompromissen mellan träningshastighet och minnesförbrukning. Som framgår av modellavsnittet kan det spara mycket minne genom att bara överväga handlingstokeniseringen, men en alltför aggressiv minskning av handlingsutrymmet kan leda till både minskning av noggrannheten och långsammare prestanda. Det senare händer eftersom alla tokens genereras sekventiellt på ett autoregressivt sätt av modellavkodaren. Därför växer slutledningstiden linjärt med antalet tokens per åtgärd när softmax på åtgärdsutrymmet inte längre är flaskhalsen.
När man satte upp AlphaTensor-utbildningen fann man de största svårigheterna med att hantera skådespelarprocessen. Om tensorerna inte lagras i rätt format kan MCTS lätt orsaka okontrollerad minnesanvändningsökning. Å andra sidan, om antalet lagrade tensorer under varje simulering reduceras för mycket, kan MCTS spendera oändligt mycket tid på att beräkna om de erforderliga tillstånden.
Låt oss ta ett exempel på spelsimuleringssteget, där spelet utforskas genom att titta på möjliga framtidsscenarier. För varje tillstånd, om vi inte sparar åtgärderna som genereras av modellen och vi bestämmer oss för att endast spara det slumpmässiga fröet som används för att ta prov på åtgärderna från policyn, så måste vi varje gång vi utforskar en trädnod beräkna om policyn och prova sedan åtgärderna. Helt klart bestämde vi oss för att lagra de samplade åtgärderna för att spara tid och för att undvika att behöva hantera modelldelning mellan olika processer i fallet med MCTS-utforskningsparallellisering. Det räckte dock inte att bara spara åtgärderna för att få ett tillräckligt effektivt agerande steg. Faktum är att tiden för att konvertera n_steps-åtgärderna till (u, v, w) tripletten, minska speltensortillståndet och skapa de nya 3D-tensorerna från n_samples-åtgärderna skulle lätt bli en flaskhals för hela träningen. För det andra ville vi inte lagra alla möjliga framtida tillstånd för varje samplade åtgärd, eftersom detta skulle ha en enorm inverkan på minnet som används av algoritmen. Anta att vi sätter n_samples=32, n=7 och N=5, och låt oss komma ihåg att N är storleken på den kvadratiska matrisprodukten vi vill reducera och n är antalet tidigare åtgärder som modellen minns. I denna situation skulle varje tillståndstensor ha formen (8, 25, 25, 25), vilket multiplicerat med 32 skulle resultera i 3282525254 byte för varje nod i grafen. Med tanke på att varje simulering i expansionsfasen genererar en ny nod (och n_sim=200), skulle vi ha en slutlig minnesförbrukning på 200328252525*4 = 3.2 GB för den första MCTS-noden ensam. I värsta fall, när man utforskar verkande max_rank-noder (där max_rank=150), skulle detta resultera i en total minnesförbrukning på 150 * 3.2 GB = 480 GB i RAM-minne (eller GPU-minne om alla tensorer var lagrade på GPU) . Vi körde utbildningen på vår arbetsstation med 128 GB RAM och 48 GB GPU-minne, så vi var tvungna att minska minnesförbrukningen.
Eftersom vi inte ville öka exekveringstiden antog vi en optimering som utnyttjar redundansen i de producerade statliga tensorerna. Faktum är att tensorerna har n-1 tidigare åtgärder gemensamma, som sedan kan lagras en gång och inte upprepas för varje lagrad tensor. Detta resulterar i en minnesminskning på 2/7~28%, vilket innebär att i värsta fall kan 137 GB lagras. Vid det här laget, genom att helt enkelt beskära den oanvända delen av trädet (som de ovalda banorna) och lagra tensorerna i CPU-minnet, kunde vi undvika minnesfel under träningen.
När OpenAlphaTensor nu är öppen källkod öppnar sig flera spännande vägar för vidareutveckling.
En naturlig utveckling är finjusteringen av OpenAlphaTensor på målhårdvaruenheter. Detta förväntas leda till mycket konkurrenskraftiga beräkningsprestanda. Jag kommer att publicera mer om prestanda för OpenAlphaTensor på olika hårdvara på GitHub. När denna artikel skrevs, genomgick OpenAlphaTensor utbildning.
Ett annat viktigt framsteg skulle vara stödet för fjärrkompilering, vilket gör det möjligt för användare att bygga algoritmer optimerade för edge-enheter. Detta kan uppnås genom att lagra OpenAlphaTensor-modellen på en server, medan matrismultiplikationsalgoritmen utvärderas på olika hårdvara.
Det kan också vara viktigt att utöka stödet för olika kompilatorer för att beräkna den latensbaserade belöningskorrigeringen. Olika kompilatorer kan leda till olika optimerade algoritmer på en given hårdvara. Till exempel visade DeepMind-papperet lovande resultat med JAX och XLA-kompilatorn på TPU och Nvidia GPU:er. Det skulle vara intressant att utvärdera detta med NCCL på Nvidia eller LLVM på processorer.
Slutligen är det fortfarande en stor öppen utmaning att utöka modellen och träningsalgoritmen för att stödja större matrisstorlekar. För närvarande stöder OpenAlphaTensor en maximal matrisstorlek på 5, men den kan tillämpas genom att dela upp större matrismultiplikationer i grupper av små MM:s med en storlek mindre än 5. Detta tillvägagångssätt är suboptimalt och utför reduktionen direkt på den stora tensorn som motsvarar fullständig MM skulle teoretiskt kunna leda till bättre resultat.
Diego Fiori är CTO för Nebuly AI, ett företag som har åtagit sig att göra AI-optimering till en del av varje utvecklares verktygslåda.
- SEO-drivet innehåll och PR-distribution. Bli förstärkt idag.
- Platoblockchain. Web3 Metaverse Intelligence. Kunskap förstärkt. Tillgång här.
- Källa: 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
- :är
- ][s
- $UPP
- 1
- 3d
- 8
- a
- Able
- Om oss
- ovan
- Absolut
- accelererande
- Enligt
- i enlighet med detta
- Konto
- noggrannhet
- Uppnå
- uppnås
- Handling
- åtgärder
- faktiskt
- lagt till
- Annat
- justerat
- antagen
- avancera
- Fördel
- Efter
- Recensioner
- aggregation
- aggressiv
- AI
- AI-system
- Syftet
- algoritm
- algoritmer
- Alla
- tillåta
- tillåter
- ensam
- redan
- alternativ
- bland
- mängd
- analysera
- och
- Annan
- tillämpas
- tillvägagångssätt
- tillvägagångssätt
- arkitektur
- ÄR
- Artikeln
- konstgjord
- artificiell intelligens
- AS
- delad
- associerad
- At
- Automatiserad
- tillgänglig
- undvika
- tillbaka
- säkerhetskopiering
- baserat
- I grund och botten
- grund
- BE
- därför att
- blir
- innan
- Där vi får lov att vara utan att konstant prestera,
- nedan
- riktmärke
- BÄST
- Bättre
- mellan
- Bortom
- ombord
- Brädspel
- bunden
- bredare
- BT
- SLUTRESULTAT
- byggt
- by
- kallas
- KAN
- kan inte
- Vid
- Orsak
- orsakas
- vissa
- utmanar
- utmanande
- byta
- Kanal
- ChatGPT
- barn
- val
- välja
- valda
- citerade
- klar
- klart
- koda
- samlar
- kombinera
- engagerad
- Gemensam
- företag
- jämfört
- konkurrenskraftig
- komplex
- Komplexiteten
- komponenter
- sammansatt
- kompromiss
- beräkning
- beräkningskraft
- Compute
- databehandling
- Begreppen
- konceptuellt
- tillstånd
- förtroende
- Tänk
- anses
- med tanke på
- anser
- konsumtion
- innehåller
- fortsätta
- Däremot
- konverterad
- Kärna
- Motsvarande
- motsvarar
- kunde
- Motverka
- CPU
- skapar
- Skapa
- kriterier
- CTO
- Aktuella
- För närvarande
- datum
- som handlar om
- beslutar
- beslutade
- Beslutet
- Beslutsfattande
- beslut
- djup
- djupt lärande
- Deepmind
- beror
- beskriven
- Bestämma
- Utvecklare
- Utveckling
- anordning
- enheter
- DICT
- olika
- svårigheter
- Dimensionera
- dimensioner
- direkt
- Upptäck
- upptäcka
- diskuteras
- fördelning
- dividerat
- ner
- Drop
- under
- e
- varje
- Tidigare
- lätt
- kant
- Effektiv
- effektiv
- antingen
- elementet
- element
- inbäddade
- slutar
- förbättrad
- enorm
- tillräckligt
- inträde
- fel
- uppskatta
- beräknad
- Eter (ETH)
- utvärdera
- utvärderade
- utvärdering
- utvärderingar
- Även
- Varje
- exempel
- exempel
- spännande
- utförande
- expansionen
- förväntat
- erfarenhet
- Förklara
- förklarade
- bedrifter
- utforskning
- utforska
- utforskas
- Utforska
- uttryckt
- förlänga
- sträcker
- extremt
- ganska
- familj
- snabbare
- återkoppling
- få
- Figur
- slutlig
- finna
- Förnamn
- Flyta
- Fokus
- följer
- efter
- Fotavtryck
- För
- formen
- format
- formeln
- hittade
- från
- full
- fungera
- grundläggande
- ytterligare
- ytterligare utveckling
- framtida
- lek
- Games
- generera
- genereras
- genererar
- generera
- skaffa sig
- få
- Ge
- ges
- Ge
- Målet
- Går
- god
- GPU
- GPUs
- diagram
- stor
- Grupp
- Gruppens
- Växer
- Tillväxt
- styra
- sidan
- Arbetsmiljö
- hänt
- händer
- hårdvara
- hårdvara
- hårdvaruenheter
- Har
- har
- här.
- högsta
- horisonten
- Hur ser din drömresa ut
- How To
- Men
- HTTPS
- stor
- humant
- i
- SJUK
- idealisk
- IDX
- bild
- Inverkan
- genomföra
- genomförande
- med Esport
- förbättra
- förbättras
- in
- Öka
- Ökar
- alltmer
- oberoende
- index
- informationen
- inledande
- ingång
- istället
- instruktioner
- Intelligens
- intressant
- inre
- introducerade
- intuition
- IT
- iteration
- DESS
- sig
- jpg
- KDnuggets
- Ha kvar
- Nyckel
- kunskap
- känd
- språk
- Large
- större
- Efternamn
- Latens
- senaste
- lager
- skikt
- leda
- lärt
- inlärning
- Lista
- se
- du letar
- Lot
- gjord
- Huvudsida
- större
- göra
- Framställning
- hantera
- hantera
- många
- karta
- kartläggning
- Matris
- Materia
- maximal
- betyder
- meningsfull
- betyder
- medlem
- Minne
- metod
- metoder
- metriska
- saknas
- blandning
- modell
- modeller
- modifierad
- modul
- mer
- mer effektiv
- Dessutom
- mest
- flytta
- förflyttar
- multiplicerat
- Natural
- Natur
- Nära
- nödvändigt för
- Behöver
- behövs
- negativ
- nätverk
- neural
- neurala nätverk
- Nya
- Nästa
- nlp
- nod
- noder
- icke-experter
- Begrepp
- antal
- Nvidia
- erhållna
- of
- Erbjudanden
- on
- ONE
- öppet
- öppen källkod
- OpenAI
- drift
- Verksamhet
- optimala
- optimering
- Optimera
- optimerad
- Tillbehör
- ursprungliga
- Övriga
- annat
- produktion
- övergripande
- Papper
- papper
- parameter
- del
- särskilt
- särskilt
- reservdelar till din klassiker
- Personer
- perfekt
- prestanda
- utför
- fas
- bitar
- plato
- Platon Data Intelligence
- PlatonData
- Spela
- Spelaren
- i
- Punkt
- Strategier
- policy
- placera
- möjlig
- potentiell
- kraft
- Praktisk
- praktiken
- förutse
- förutsagda
- preferenser
- presenteras
- föregående
- Sannolikheten
- förmodligen
- Problem
- process
- processer
- producera
- producerad
- Produkt
- progression
- progressiv
- lovande
- föreslagen
- bevisligen
- beprövade
- publicera
- publicerade
- RAM
- slumpmässig
- leden
- snarare
- kommit fram till
- når
- nyligen
- minska
- Minskad
- reducerande
- raffinerade
- förstärkning lärande
- frigörs
- Återstående
- resterna
- anmärkningsvärd
- ihåg
- avlägsen
- upprepade
- representerar
- representerar
- Obligatorisk
- Kräver
- ansvarig
- begränsad
- resultera
- resulterande
- Resultat
- avkastning
- återgår
- revolutionera
- Belöna
- RAD
- rt
- Körning
- s
- Samma
- Save
- sparande
- scenario
- scenarier
- ordningen
- Sök
- Andra
- §
- frö
- vald
- väljer
- Val
- SJÄLV
- in
- inställning
- flera
- Forma
- delning
- Kort
- skall
- show
- visas
- Visar
- Signal
- signifikant
- signifikant
- liknande
- Likheterna
- Enkelt
- enkelhet
- helt enkelt
- simulering
- eftersom
- Situationen
- Storlek
- storlekar
- Small
- mindre
- Snapshot
- So
- lösning
- LÖSA
- Lösa
- några
- Källa
- Utrymme
- utrymmen
- specifik
- specifikt
- specificerade
- fart
- spendera
- Spendera
- delas
- kvadrat
- staplade
- Etapp
- standard
- starta
- startar
- Ange
- state-of-the-art
- anges
- Stater
- statistik
- Steg
- Steg
- stoppa
- lagra
- lagras
- lagrar
- okomplicerad
- framgångsrik
- sådana
- stödja
- Stöder
- syntetisk
- syntetiska data
- syntetiskt
- system
- System
- skräddarsydd
- Ta
- tar
- tar
- Målet
- uppgift
- uppgifter
- tekniker
- terminal
- villkor
- den där
- Smakämnen
- Framtiden
- Grafen
- den information
- Staten
- deras
- Dem
- vari
- därför
- Dessa
- Tredje
- tre
- tredimensionella
- Genom
- tid
- tidskrävande
- till
- tillsammans
- token
- tokenization
- befogenhetsbeskrivas
- tokens
- alltför
- toolkit
- topp
- brännaren
- Totalt
- traditionell
- Tåg
- tränad
- Utbildning
- tåg
- bana
- transformerad
- behandla
- sann
- förstå
- Universum
- oanvänd
- Uppdatering
- uppdaterad
- Uppdateringar
- uppdatering
- uppgradera
- us
- Användning
- användning
- användare
- vanligen
- värde
- Värden
- olika
- version
- Video
- videospel
- Besök
- Besök
- W
- Sätt..
- Vad
- som
- medan
- brett
- wikipedia
- kommer
- vinna
- med
- ord
- arbetsstation
- värt
- skulle
- skrivning
- skriven
- X
- zephyrnet
- noll-