Photo by DeepMind on Unsplash
Mnożenie macierzy jest podstawową operacją stosowaną w wielu systemach, od sieci neuronowych po naukowe procedury obliczeniowe. Znalezienie wydajnych i dających się udowodnić poprawnych algorytmów mnożenia macierzy może mieć ogromny wpływ na przyspieszenie i zwiększenie wydajności obliczeń, ale jest to bardzo trudne zadanie. Przestrzeń możliwych algorytmów jest ogromna, a tradycyjne metody odkrywania algorytmów, takie jak heurystyka zaprojektowana przez człowieka czy wyszukiwanie kombinatoryczne, są często nieoptymalne.
DeepMindNiedawno zaproponowane przez firmę rozwiązanie oparte na sztucznej inteligencji do automatycznego wyszukiwania wykracza daleko poza ludzką intuicję. Rozwiązanie składa się z agenta uczenia się metodą głębokiego wzmacniania o nazwie AlphaTensor AlfaZero. Ten agent jest przeszkolony do gry w grę dla jednego gracza, TensorGame, w której celem jest odkrycie wydajnych obliczeniowo algorytmów mnożenia macierzy.
AlphaTensor jest szczególnie dobry w obsłudze dużych macierzy poprzez rozkładanie dużych mnożeń macierzy na mniejsze mnożenia. Co więcej, AlphaTensor może być wykorzystany do osiągnięcia najnowocześniejszej wydajności mnożenia macierzy po dokładnym dostrojeniu na określonym urządzeniu sprzętowym.
AlphaTensor ma ogromny potencjał w zakresie przyspieszania przetwarzania głębokiego uczenia. W uczeniu głębokim wiele czasochłonnych operacji można odwzorować na mnożenie macierzy. Używając AlphaTensor do optymalizacji tych operacji, można znacznie poprawić ogólną wydajność modeli głębokiego uczenia.
Ostatnio OpenAlphaTensor, pierwsza implementacja open source AlphaTensor, która może zrewolucjonizować moc obliczeniową modeli głębokiego uczenia się.
Tensor mnożenia macierzy
Dla osób niebędących ekspertami w optymalizacji mnożenia macierzy zrozumienie, w jaki sposób operację taką jak mnożenie macierzy można odwzorować w trójwymiarowym tensorze, może nie być proste. Postaram się wyjaśnić to prostymi słowami i przykładami.
Rozważmy iloczyn C = A*B, gdzie dla uproszczenia zarówno A, jak i B są macierzami kwadratowymi o rozmiarze N. Operację mnożenia można odwzorować w trójwymiarowym tensorze kształtu (N^3, N^2, N^2). Pierwszy wymiar tensorowy reprezentuje spłaszczoną macierz A, drugi wymiar spłaszczoną macierz B, a trzeci wymiar spłaszczoną macierz C.
Tensor ma tylko wartości binarne (1 lub 0) dla każdego wpisu. Zauważ, że tensor reprezentuje operację mnożenia, więc jest niezależny od wartości macierzy A i B.
Każdemu wpisowi tensora odpowiada współczynnik operacji. Na przykład, aby obliczyć C[1,1], konieczne jest pomnożenie zarówno A[1,1], jak i B[1,1]. Dlatego wpis tensorowy [0,0,0], który odpowiada A[1,1], B[1,1] i C[1,1], będzie miał wartość 1. Natomiast obliczenie C[1,1 ,2,1], A[1] nie jest potrzebne. Zatem rząd tensora T[N+0, :, XNUMX] będzie zawierał tylko zera.
Poniższy rysunek przedstawia przykład tensora dla N=2.
Zdjęcie z DeepMind papier opublikowane w Natura
Jak pokazano w (b) i (c) na powyższym rysunku, możliwe jest zaimplementowanie algorytmu obliczania iloczynu za pomocą dekompozycji tensora 3D. Dokładniej, poniższy algorytm może być użyty do przekształcenia rozkładu tensorowego (macierze U, V, W) na algorytm mnożenia macierzy.
Metaalgorytm sparametryzowany do obliczania iloczynu macierzowego C=AB wprowadzony w DeepMind papier
Gra Tensor
Problem znalezienia wydajnych algorytmów mnożenia macierzy jest niezwykle trudny, ponieważ liczba możliwych algorytmów do rozważenia jest znacznie większa niż liczba atomów we wszechświecie, nawet w przypadku małych przypadków mnożenia macierzy.
DeepMind przekształcił ten problem w grę dla jednego gracza i nazwał ją TensorGame. W tej grze gracz wybiera sposób łączenia różnych wpisów macierzy, aby je pomnożyć. Wynik jest przypisywany na podstawie liczby operacji wymaganych do uzyskania prawidłowego wyniku mnożenia. Gra kończy się, gdy osiągnięty zostanie tensor zera lub gdy zostanie wykonana maksymalna liczba ruchów. Ostateczna faktoryzacja jest oceniana na podstawie oszacowania rezydualnego rzędu i pewnych kryteriów optymalizacji, takich jak asymptotyczna złożoność czasowa lub praktyczny czas wykonywania.
Pozycja początkowa w grze TensorGame odpowiada tensorowi mnożenia macierzy wyrażonemu w sposób losowy.
W każdym kroku t gry gracz zapisuje trzy wektory
, który określa tensory rangi 1 . Stan gry jest aktualizowany poprzez odjęcie wybranych przez gracza wektorów:
gdzie
jest tensorem mnożenia macierzy.Jeśli gra kończy się w p krokach, oznacza to, że tensor mnożenia macierzy
można rozłożyć na tensory p rang-1 , czyli ma co najmniej rangę p.TensorGame można wtedy interpretować jako algorytm dekompozycji rang, a AlphaTensor można postrzegać jako algorytm szacowania rangi tensora.
Architektura AlphaTensor
Do tej pory poznaliśmy TensorGame i wyjaśniliśmy, w jaki sposób jej rozwiązanie można postrzegać jako algorytm mnożenia macierzy. Przyjrzyjmy się teraz głównym koncepcjom AlphaTensor, algorytmu używanego w grze.
Architektura AlphaTensor jest zasadniczo architekturą transformatora enkodera-dekodera, w której:
- koder przyjmuje jako dane wejściowe stan gry , n poprzednich działań podjętych przez model (zwykle n=7) oraz indeks czasowy t bieżącej akcji. Informacje są układane razem w tensor o kształcie (n+1, N^2, N^2, N^2). Tensor ten jest następnie przekształcany i przekształcany (za pomocą trzech warstw liniowych) w tensor o kształcie (N^2, N^2, c), gdzie c jest wewnętrznym wymiarem modelu.
- dekoder generuje akcje n_kroków z osadzonego wektora podanego przez koder w sposób autoregresywny. Każda akcja odpowiada żetonowi trójek reprezentujący jedną z trójek rozkładających tensor gry (tj. obniżający jej rangę)
Model jest szkolony przez naprzemienną propagację wsteczną i działanie modelu. Działanie modelu służy do generowania danych, które są następnie wykorzystywane do uczenia modelu. W praktyce model jest szkolony za pomocą mieszanki danych generowanych syntetycznie i danych generowanych przez model podczas działania. Krok aktorski jest wykonywany poprzez pobranie tensora 3D odpowiadającego operacji macierzowej i rozegranie na nim gier n_aktorów. Każdy aktor gra albo na zasadach standardowych, albo na zasadach alternatywnych (zmiana podstawy jest stosowana z określonym prawdopodobieństwem). Wyniki są następnie gromadzone i mogą być wykorzystane na etapie szkolenia z danymi syntetycznymi.
Krok działania opiera się na wyszukiwarce Monte Carlo Tree Search (MCTS) firmy AlphaZero, zmodyfikowanej w celu obsługi dużych przestrzeni akcji. Krótko mówiąc, przed wybraniem akcji ścieżki n_sims są badane na podstawie danych wyjściowych modelu, przy czym maksymalna przyszła eksploracja wynosi 5 kroków. Wygenerowane przez model prawdopodobieństwa są następnie korygowane z uwzględnieniem wygenerowanych ścieżek. Następnie wybierana jest akcja z najbardziej obiecującą przyszłą ścieżką (ścieżkami), aby kontynuować grę.
Podczas szkolenia modelu nagroda jest w rzeczywistości nagrodą ujemną (karą). Jego bezwzględna wartość wzrasta z każdym dodatkowym krokiem wymaganym do rozwiązania gry. Jeśli model wykonuje m kroków, aby rozwiązać grę TensorGame, nagroda związana z grą wynosi r=-m. Jeśli model nie jest w stanie rozwiązać TensorGame w krokach max_rank, nagroda jest obliczana poprzez oszacowanie rangi pozostałego tensora. Ranga jest szacowana jako suma rang macierzy składających się na tensor. Oszacowanie jest górną granicą rzeczywistego rzędu tensora.
Podczas dostrajania modelu nagroda karna w stanie końcowym powinna również uwzględniać opóźnienie algorytmu wytworzonego przez model. Formuła nagrody przyjmuje postać rt'=rt+λbt, gdzie rt jest opisanym wcześniej schematem nagrody, bt jest nagrodą wzorcową (niezerową tylko w stanie końcowym) oraz λ jest współczynnikiem określonym przez użytkownika.
Przyspieszenia (%) algorytmów odkrytych przez AlphaTensor dostosowanych do GPU i TPU, wyodrębnione z artykułu DeepMind. Przyspieszenia są mierzone względem standardowego mnożenia macierzy (np. cuBLAS dla GPU) na tym samym sprzęcie i porównywane z Algorytm Strassena-kwadrat, źródło: DeepMind.
niedawno zwolniłem OpenAlphaTensor, pierwsza implementacja Open Source AlphaTensor. W tej sekcji omówię implementację. Jak omówiliśmy wcześniej, architektura AlphaTensor jest dość prosta, oparta na standardowym transformatorze z architekturą enkoder-dekoder. Najciekawsze komponenty AlphaTensor to pierwsza warstwa w części enkodera oraz sposób próbkowania akcji.
Zacznijmy od pierwszej warstwy kodowania.
# 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
W powyższym fragmencie pokazujemy, jak tensor wejściowy jest rozkładany na trzy tensory, które są następnie używane jako dane wejściowe zapytania, klucza i wartości warstwy transformatora.
- W trzech wymiarach tensorowych reprezentujących spłaszczone macierze (A, B, C) tensor wejściowy jest spłaszczany wzdłuż każdego wymiaru wraz z wymiarem reprezentującym poprzednie działania. W ten sposób w każdej spłaszczonej kopii tensora wejściowego wybrany wymiar jest agregacją ostatnich wartości T-1 i wartości rzeczywistej, dla wszystkich wartości S wybranego wymiaru, gdzie S=N^2. Z filozoficznego punktu widzenia jest tak, jakbyśmy dla każdego wymiaru skupiali się na tym, co wydarzyło się w poprzednich działaniach w tym wymiarze.
- Skalary są odwzorowywane w trzech różnych przestrzeniach o wymiarze S^2, a następnie przekształcane w celu połączenia z tensorami otrzymanymi w poprzednim punkcie. Koncepcyjnie skalary są odwzorowywane na przestrzeń osadzania o wymiarze S^2, a następnie osadzone informacje są dzielone na wektory S i układane razem, podobnie jak dzieje się to z tekstem po tokenizacji.
- Tokeny skalarne są łączone ze zrestrukturyzowanym tensorem wejściowym, a następnie podawane jako dane wejściowe do warstwy liniowej w celu mapowania informacji o skupieniu skalarów + historii kanału w wewnętrznym wymiarze modelu.
Te trzy kroki można interpretować jako sposób przekazania modelowi zarówno informacji o skalarach (jak w kroku czasowym TensorGame), jak i skupienia się na poprzednich działaniach dla każdego kanału.
Jeśli chodzi o sposób tworzenia akcji, warto zauważyć, że AlphaTensor generuje na wyjściu trójkę u, v, w, która ma na celu zmniejszenie rangi tensorowej. Trzy wektory mają rozmiar S, a ponieważ są one połączone, model musi generować wektor o rozmiarze 3*S. AlphaTensor jest trenowany za pomocą algorytmu RL, więc wszystkie możliwe działania muszą być wyrażone jako prawdopodobieństwa w wyliczonej przestrzeni, tj. model generuje prawdopodobieństwo dla różnych działań. Oznacza to, że każdy wektor w przestrzeni 3S powinien być odwzorowany na inną akcję. Daje to przestrzeń akcji o rozmiarze |F|^(3S), gdzie |F| to liczba różnych wartości, jakie może przyjąć element u, v, w. Zwykle wartości są ograniczone do (-2, -1, 0, 1, 2), co daje liczność 5 elementów.
Oto główne wyzwanie: aby wygenerować prawdopodobieństwa działania dla iloczynu macierzowego macierzy o rozmiarze 5, potrzebowalibyśmy pamięci 5^75 * 4 bajtów, co oznaczałoby ~10^44 GB pamięci. Najwyraźniej nie jesteśmy w stanie zarządzać tak dużą przestrzenią działania.
Jak rozwiązać problem? Aby zmniejszyć zużycie pamięci przez prawdopodobieństwa działania, możemy podzielić trójki na mniejsze części, „tokenizować” je i traktować te części jako wygenerowane tokeny w architekturze transformatora, tj. tokeny są podawane jako dane wejściowe do dekodera w auto-regresji sposób. W powyższym przykładzie możemy podzielić trójki na 15 kawałków, zmniejszając zużycie pamięci do 15 * 5^(75/15) * 4, czyli 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), )
Powyżej pokazujemy fragment kodu do generowania pełnej akcji. W kodzie self.core zawiera warstwę dekodera, a tensor e reprezentuje dane wyjściowe warstwy kodera. Zero można uznać za tzw token w modelach NLP, a akcje n_steps reprezentujące porcje n_steps są generowane w sposób progresywny.
Model zwraca trzy wielkości:
- Wygenerowane akcje
- Prawdopodobieństwo związane z pełną akcją
- Logi utworzone w celu wygenerowania pierwszej akcji (pierwszej porcji), która zostanie wykorzystana do obliczenia wartości modelu.
Warto poświęcić kilka słów parametrowi n_samples. Parametr jest używany dla kroku aktorskiego i pozwala modelowi generować różne wersje trójek, które następnie zostaną wykorzystane do eksploracji przestrzeni akcji w algorytmie Monte Carlo Tree Search używanym w procesie aktorskim. Różne działania n_samples są próbkowane zgodnie z zasadami wygenerowanymi przez model.
Krok aktorski
Najtrudniejszą częścią całego algorytmu jest prawdopodobnie krok aktorski użyty do rozwiązania TensorGame. Algorytm nie jest dogłębnie wyjaśniony w artykule AlphaTensor, ponieważ opiera się na kilku poprzednich artykułach DeepMind, które są tylko cytowane i podane jako znane. Tutaj zrekonstruuję wszystkie brakujące elementy i wyjaśnię krok po kroku naszą implementację.
Etapy aktorskie możemy podzielić na trzy różne komponenty:
- Poszukiwanie drzewa Monte-Carlo
- Symulacja gry
- Ulepszone obliczenia polityki
Przeanalizujmy je jeden po drugim.
Wyszukiwanie drzewa Monte-Carlo (MCTS)
Monte Carlo Tree Search (MCTS) to szeroko stosowana technika sztucznej inteligencji w grach, szczególnie w grach planszowych i grach wideo. Algorytm tworzy drzewo gry, które symuluje potencjalne ruchy i wyniki oraz wykorzystuje losowe pobieranie próbek do oceny oczekiwanej nagrody za każdy ruch. Następnie algorytm iteracyjnie wybiera ruch z najwyższą oczekiwaną nagrodą i symuluje wyniki, aż osiągnie stan końcowy lub określony warunek zatrzymania. Symulacje służą do oszacowania prawdopodobieństwa wygranej dla każdego ruchu i ukierunkowania procesu decyzyjnego. Wykazano, że MCTS jest skuteczny w złożonych grach, w których liczba możliwych ruchów i wyników jest duża, i został wykorzystany w udanych systemach sztucznej inteligencji, takich jak AlphaGo.
W AlphaTensor używana jest zmodyfikowana wersja oryginalnego MCTS. W szczególności, zamiast losowego wybierania akcji z całej przestrzeni akcji, akcja jest wybierana spośród podzbioru generowanego bezpośrednio przez model (poprzez przedstawione wcześniej n_próbki). Korekta aktualizacji polisy jest następnie stosowana w kroku obliczania ulepszonej polisy.
W naszej implementacji zdecydowaliśmy się przechowywać wszystkie informacje o drzewie Monte-Carlo w słowniku, którego kluczem jest wersja skrótu stanu TensorGame, a jako wartościami informacje związane z samym stanem. Każdy krok Monte-Carlo zaczyna się od węzła i symuluje mini-gry n_sim, eksplorując przyszłość z horyzontem 5 ruchów. Jeśli węzeł był już eksplorowany w poprzednich symulacjach, n_sim jest korygowane o liczbę poprzednich eksploracji. Dla każdego węzła liczba odwiedzin jest przechowywana w tensorze N_s_a, ponieważ ten tensor zawiera liczbę odwiedzin przypadających na akcję potomną węzła (wśród tych próbkowanych przez model).
def monte_carlo_tree_search( model: torch.nn.Module, state: torch.Tensor, n_sim: int, t_time: int, n_steps: int, game_tree: Dict, state_dict: Dict,
): """Runs the monte carlo tree search algorithm. Args: model (torch.nn.Module): The model to use for the simulation. state (torch.Tensor): The initial state. n_sim (int): The number of simulations to run. t_time (int): The current time step. n_steps (int): The maximum number of steps to simulate. game_tree (Dict): The game tree. state_dict (Dict): The dictionary containing the states. """ state_hash = to_hash(extract_present_state(state)) if state_hash in state_dict: with torch.no_grad(): N_s_a = state_dict[state_hash][3] n_sim -= int(N_s_a.sum()) n_sim = max(n_sim, 0) for _ in range(n_sim): simulate_game(model, state, t_time, n_steps, game_tree, state_dict) # return next state possible_states_dict, _, repetitions, N_s_a, q_values, _ = state_dict[ state_hash ] possible_states = _recompose_possible_states(possible_states_dict) next_state_idx = select_future_state( possible_states, q_values, N_s_a, repetitions, return_idx=True ) next_state = possible_states[next_state_idx] return next_state
Powyższy kod pokazuje naszą implementację algorytmu. Dla uproszczenia kodu korekta zasad jest wykonywana w funkcji simulate_game.
Symulacja gry
Funkcja simulate_game odpowiada za eksplorację drzewa złożonego z węzłów reprezentujących określony stan gry TensorGame. Uruchamia również model za każdym razem, gdy napotkany zostanie węzeł liścia, i przechowuje wszystkie informacje o węzłach w słowniku state_dict. Przyjrzyjmy się bliżej jego implementacji:
@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)
Każda symulacja jest podzielona na trzy części:
- Wybór
- Ekspansja
- backup
W części wyboru symulacja jest uruchamiana na już wygenerowanych węzłach drzewa, a następny węzeł jest wybierany za pomocą następującej funkcji:
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()]
W praktyce działanie maksymalizujące funkcję ucb:
dla danego stanu jest wybrany. Tutaj Q reprezentuje wartości Q wygenerowane przez model, a π reprezentuje losowy rozkład działań próbkowanych przy użyciu polityki modelu. N(s, a) reprezentuje liczbę wizyt węzła do działania a z węzła s.
Gdy faza wyboru dotrze do węzła liścia, jeśli symulacja nie osiągnęła warunku końcowego (pod względem maksymalnej eksploracji, tj. przyszłego horyzontu lub zakończenia gry), model jest następnie używany do wyboru n_próbek alternatywnych węzłów (będą węzłów w kolejnej iteracji). Nazywa się to fazą ekspansji, ponieważ do drzewa dodawane są nowe węzły. Następnie żaden kolejny węzeł nie jest eksplorowany w bieżącej symulacji, ale liść q_value jest wysyłany do następnego kroku symulacji: tworzenia kopii zapasowej.
Kopia zapasowa jest ostatnim etapem każdej symulacji. Podczas tworzenia kopii zapasowej, jeśli węzeł liścia był w stanie końcowym, obliczana jest ostateczna nagroda; w przeciwnym razie wartość liścia q jest używana jako szacunkowa nagroda. Następnie nagroda jest wstecznie propagowana na trajektorii symulacji, aktualizując zarówno stany q_values, jak i aktualizując licznik odwiedzin N(s, a). W poniższym fragmencie pokazujemy kod wstecznej propagacji nagrody.
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
Ulepszone obliczanie zasad
Gdy wszystkie symulacje zostaną uruchomione, a MCTS zaoferuje interesującą migawkę najbliższej przyszłości, nadszedł czas, aby zaktualizować politykę związaną z przewidywanymi węzłami i zwrócić je, aby można było z nich korzystać podczas szkolenia. Ulepszona polityka, zgodnie z metodą opisaną w Huberta i in, służy do zarządzania dużymi przestrzeniami akcji. W rzeczywistości, w przypadku małej przestrzeni poszukiwań, podczas MCTS możliwe jest losowe próbkowanie akcji z przestrzeni akcji i ocena jej wpływu. Podobne podejście w znacznie większej przestrzeni działania doprowadziłoby do rozbieżności wszystkich trajektorii w różnych ścieżkach i wymagałoby nieskończonej liczby trajektorii w celu uzyskania miarodajnych statystyk, a następnie zaktualizowania polityki. Ponieważ tutaj używamy sample-MCTS w celu uniknięcia rozproszenia, tj. n_samples akcji jest próbkowanych zgodnie z polityką modelu, a następnie MCTS po prostu wybiera jedną z próbkowanych akcji podczas eksploracji drzewa, musimy wziąć pod uwagę poprawkę próbki podczas obliczania ostateczna zaktualizowana polityka, która będzie używana podczas uczenia modelu.
W praktyce ulepszona polityka jest obliczana jako
gdzie
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
Zauważ, że w naszej implementacji po obliczeniu polityki z tensora N_s_a musimy ją zmapować z powrotem do pierwotnego tensora akcji. W rzeczywistości N_s_a bierze pod uwagę tylko akcje próbkowane przez model, podczas gdy ostateczna polityka musi zawierać prawdopodobieństwa również dla akcji niezbadanych.
Różnice w odniesieniu do algorytmu szkoleniowego ChatGPT
AlphaTensor to najnowszy członek rodziny metod sztucznej inteligencji AlphaGo/AlphaZero firmy DeepMind. Metody te opierają się na algorytmie wyszukiwania drzewa Monte Carlo (MCTS), który został udoskonalony i udoskonalony przez DeepMind, aby sprostać coraz bardziej złożonym zadaniom. Inny system sztucznej inteligencji, ChatGPT OpenAI, który wywołał wiele szumu ze względu na swoje niezwykłe działanie, został przeszkolony przy użyciu innego podejścia, zwanego uczeniem się przez wzmacnianie za pomocą informacji zwrotnej od ludzi (RLHF).
RLHF to technika precyzyjnego dostrajania używana do dostrajania modeli językowych w celu przestrzegania zestawu pisemnych instrukcji. Wykorzystuje ludzkie preferencje jako sygnał nagrody w celu dostrojenia modelu, dopasowując w ten sposób zachowanie modelu językowego do określonych preferencji określonej grupy ludzi, a nie do szerszego pojęcia „wartości ludzkich”.
Z kolei MCTS to oparty na drzewie algorytm wyszukiwania używany do określania optymalnych ruchów w grach. Symuluje potencjalne ruchy i aktualizuje wartości każdego ruchu na podstawie ich wyników, kierując wyborem najlepszego ruchu.
RLHF zbiera dane z demonstracji napisanych przez ludzi i oznaczonych przez ludzi porównań między modelami AI i trenuje model nagrody, aby przewidywać preferencje danej grupy ludzi. Model nagrody jest następnie używany do dostrajania modeli AI. Z drugiej strony MCTS wykorzystuje symulacje i oceny, aby określić najlepszą decyzję.
Chociaż są to różne podejścia, RLHF i MCTS mają również podobieństwa. Obie techniki sztucznej inteligencji wykorzystują metody podejmowania decyzji i rozwiązywania problemów oraz obie wykorzystują podejście prób i błędów w celu zbadania różnych opcji i podejmowania decyzji na podstawie dostępnych informacji. Oba są również procesami iteracyjnymi, które z czasem ulegają poprawie w miarę gromadzenia większej ilości informacji i doświadczeń.
Wybór między RLHF a MCTS zależy od wykonywanego zadania. RLHF jest idealny, gdy nie ma jasnej miary do oceny wydajności modelu, podczas gdy MCTS okazał się skuteczny w zadaniach przypominających gry, w których wiedza i eksploracja przyszłości dają modelowi znaczącą przewagę.
Optymalizacja kodu na potrzeby szkolenia AlphaTensor
Implementacja algorytmu treningowego AlphaTensor wymaga znalezienia idealnego kompromisu pomiędzy szybkością treningu a zużyciem pamięci. Jak widać w sekcji Model, samo rozważenie tokenizacji akcji może zaoszczędzić dużo pamięci, ale zbyt agresywna redukcja przestrzeni akcji może prowadzić zarówno do spadku dokładności, jak i spowolnienia działania. To ostatnie dzieje się, ponieważ wszystkie tokeny są generowane sekwencyjnie w sposób autoregresyjny przez dekoder modelu. Dlatego czas wnioskowania rośnie liniowo wraz z liczbą tokenów na akcję, gdy softmax na polu akcji nie jest już wąskim gardłem.
Podczas konfigurowania treningu AlphaTensor główne trudności napotkano w radzeniu sobie z procesem aktorskim. Jeśli tensory nie są przechowywane we właściwym formacie, MCTS może łatwo spowodować niekontrolowany wzrost wykorzystania pamięci. Z drugiej strony, jeśli liczba tensorów przechowywanych podczas każdej symulacji zostanie zbytnio zmniejszona, MCTS może poświęcić nieskończoną ilość czasu na ponowne obliczenie wymaganych stanów.
Weźmy przykład etapu symulacji gry, podczas którego bada się grę, przyglądając się możliwym przyszłym scenariuszom. Jeśli dla każdego stanu nie zapiszemy działań wygenerowanych przez model i zdecydujemy się zapisać tylko losowe ziarno użyte do próbkowania działań z polityki, to za każdym razem, gdy będziemy eksplorować węzeł drzewa, musielibyśmy ponownie obliczyć politykę i następnie wypróbuj działania. Oczywiście zdecydowaliśmy się przechowywać próbkowane działania, aby zaoszczędzić czas i uniknąć konieczności zarządzania udostępnianiem modeli pomiędzy różnymi procesami w przypadku równoległości eksploracji MCTS. Jednak samo zapisanie akcji nie wystarczyło, aby uzyskać wystarczająco skuteczny krok aktorski. W rzeczywistości czas na konwersję akcji n_steps na trójkę (u, v, w), redukcję stanu tensora gry i utworzenie nowych tensorów 3D z akcji n_samples z łatwością byłby wąskim gardłem dla całego szkolenia. Po drugie, nie chcieliśmy przechowywać wszystkich możliwych przyszłych stanów dla każdej próbkowanej akcji, ponieważ miałoby to ogromny wpływ na pamięć wykorzystywaną przez algorytm. Załóżmy, że ustawiliśmy n_samples=32, n=7 i N=5 i pamiętajmy, że N to rozmiar iloczynu macierzy kwadratowej, który chcemy zmniejszyć, a n to liczba poprzednich działań zapamiętanych przez model. W tej sytuacji każdy tensor stanu miałby postać (8, 25, 25, 25), co pomnożone przez 32 dałoby 3282525254 bajty dla każdego węzła w grafie. Teraz, biorąc pod uwagę, że każda symulacja w fazie ekspansji generuje nowy węzeł (i n_sim=200), mielibyśmy ostateczne zużycie pamięci na poziomie 200328252525*4 = 3.2 GB tylko dla pierwszego węzła MCTS. W najgorszym przypadku podczas eksploracji działających węzłów max_rank (gdzie max_rank=150) skutkowałoby to całkowitym zużyciem pamięci 150 * 3.2 GB = 480 GB w pamięci RAM (lub pamięci GPU, jeśli wszystkie tensory były przechowywane w GPU) . Szkolenie prowadziliśmy na naszej stacji roboczej ze 128 GB RAM i 48 GB pamięci GPU, więc musieliśmy zmniejszyć zużycie pamięci.
Ponieważ nie chcieliśmy wydłużać czasu wykonania, przyjęliśmy optymalizację, która wykorzystuje redundancję wytworzonych tensorów stanu. W rzeczywistości tensory mają wspólne poprzednie działania n-1, które można następnie zapisać raz i nie powtarzać dla każdego zapisanego tensora. Powoduje to redukcję pamięci o 2/7 ~ 28%, co oznacza, że w najgorszym przypadku można zapisać 137 GB. W tym momencie, po prostu przycinając nieużywaną część drzewa (np. niewybrane trajektorie) i przechowując tensory w pamięci procesora, byliśmy w stanie uniknąć błędów pamięci podczas uczenia.
Ponieważ OpenAlphaTensor jest teraz open source, otwiera się kilka ekscytujących dróg dalszego rozwoju.
Naturalnym postępem jest dostrajanie OpenAlphaTensor na docelowych urządzeniach sprzętowych. Oczekuje się, że doprowadzi to do bardzo konkurencyjnej wydajności obliczeniowej. Opublikuję więcej informacji na temat wydajności OpenAlphaTensor na różnych urządzeniach GitHub. W momencie pisania tego artykułu OpenAlphaTensor przechodził szkolenie.
Kolejnym ważnym postępem byłaby obsługa zdalnej kompilacji, umożliwiająca użytkownikom tworzenie algorytmów zoptymalizowanych pod kątem urządzeń brzegowych. Można to osiągnąć, przechowując model OpenAlphaTensor na serwerze, podczas gdy algorytm mnożenia macierzy jest oceniany na innym sprzęcie.
Ważne może być również rozszerzenie obsługi różnych kompilatorów w celu obliczenia korekty nagrody opartej na opóźnieniu. Różne kompilatory mogą prowadzić do różnych zoptymalizowanych algorytmów na danym sprzęcie. Na przykład artykuł DeepMind wykazał obiecujące wyniki przy użyciu JAX i kompilatora XLA na procesorach graficznych TPU i Nvidia. Interesujące byłoby oszacowanie tego przy użyciu NCCL na Nvidii lub LLVM na procesorach.
Wreszcie, rozszerzenie modelu i algorytmu szkoleniowego w celu obsługi większych rozmiarów macierzy pozostaje głównym otwartym wyzwaniem. Obecnie OpenAlphaTensor obsługuje maksymalny rozmiar macierzy 5, ale można go zastosować, dzieląc większe mnożenia macierzy na grupy małych MM o rozmiarze mniejszym niż 5. Takie podejście jest nieoptymalne i przeprowadzanie redukcji bezpośrednio na dużym tensorze odpowiadającym pełny MM mógłby teoretycznie prowadzić do lepszych wyników.
Diego Fioriego jest dyrektorem technicznym Nebuly AI, firmy, która dąży do tego, aby optymalizacja AI stała się częścią zestawu narzędzi każdego programisty.
- Dystrybucja treści i PR oparta na SEO. Uzyskaj wzmocnienie już dziś.
- Platoblockchain. Web3 Inteligencja Metaverse. Wzmocniona wiedza. Dostęp tutaj.
- Źródło: 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
- :Jest
- ][P
- $W GÓRĘ
- 1
- 3d
- 8
- a
- Zdolny
- O nas
- powyżej
- bezwzględny
- przyspieszenie
- Stosownie
- odpowiednio
- Konto
- precyzja
- Osiągać
- osiągnięty
- Działania
- działania
- faktycznie
- w dodatku
- Dodatkowy
- Skorygowana
- przyjęty
- awansować
- Korzyść
- Po
- Agent
- zbiór
- agresywny
- AI
- Systemy SI
- Cele
- algorytm
- Algorytmy
- Wszystkie kategorie
- Pozwalać
- pozwala
- sam
- już
- alternatywny
- wśród
- ilość
- w czasie rzeczywistym sprawiają,
- i
- Inne
- stosowany
- podejście
- awanse
- architektura
- SĄ
- artykuł
- sztuczny
- sztuczna inteligencja
- AS
- przydzielony
- powiązany
- At
- zautomatyzowane
- dostępny
- unikając
- z powrotem
- backup
- na podstawie
- Gruntownie
- podstawa
- BE
- bo
- staje się
- zanim
- jest
- poniżej
- Benchmark
- BEST
- Ulepsz Swój
- pomiędzy
- Poza
- deska
- Gry planszowe
- związany
- szerszy
- BT
- budować
- wybudowany
- by
- nazywa
- CAN
- nie może
- walizka
- Spowodować
- powodowany
- pewien
- wyzwanie
- wyzwanie
- zmiana
- Kanał
- ChatGPT
- dziecko
- wybór
- Wybierając
- wybrany
- cytowane
- jasny
- wyraźnie
- kod
- zbiera
- połączyć
- zobowiązany
- wspólny
- sukcesy firma
- w porównaniu
- konkurencyjny
- kompleks
- kompleksowość
- składniki
- w składzie
- kompromis
- obliczenia
- moc obliczeniowa
- obliczać
- computing
- Koncepcje
- Koncepcyjnie
- warunek
- pewność siebie
- Rozważać
- za
- wobec
- rozważa
- konsumpcja
- zawiera
- kontynuować
- kontrast
- przeliczone
- rdzeń
- Odpowiedni
- odpowiada
- mógłby
- Przeciwdziałać
- CPU
- tworzy
- Tworzenie
- Kryteria
- CTO
- Aktualny
- Obecnie
- dane
- czynienia
- zdecydować
- postanowiła
- decyzja
- Podejmowanie decyzji
- Decyzje
- głęboko
- głęboka nauka
- DeepMind
- zależy
- opisane
- Ustalać
- Deweloper
- oprogramowania
- urządzenie
- urządzenia
- DICT
- różne
- trudności
- Wymiary
- Wymiary
- bezpośrednio
- odkryj
- odkrywanie
- omówione
- 分配
- podzielony
- na dół
- Spadek
- podczas
- e
- każdy
- Wcześniej
- z łatwością
- krawędź
- Efektywne
- wydajny
- bądź
- element
- Elementy
- osadzone
- kończy się
- wzmocnione
- ogromny
- dość
- wejście
- błąd
- oszacowanie
- szacunkowa
- Eter (ETH)
- oceniać
- oceniane
- oceny
- oceny
- Parzyste
- Każdy
- przykład
- przykłady
- ekscytujący
- egzekucja
- ekspansja
- spodziewany
- doświadczenie
- Wyjaśniać
- wyjaśnione
- exploity
- eksploracja
- odkryj
- zbadane
- Exploring
- wyrażone
- rozciągać się
- rozsuwalny
- niezwykle
- dość
- członków Twojej rodziny
- szybciej
- informacja zwrotna
- kilka
- Postać
- finał
- znalezieniu
- i terminów, a
- pływak
- Skupiać
- obserwuj
- następujący
- Ślad stopy
- W razie zamówieenia projektu
- Nasz formularz
- format
- formuła
- znaleziono
- od
- pełny
- funkcjonować
- fundamentalny
- dalej
- dalszy rozwój
- przyszłość
- gra
- Games
- Generować
- wygenerowane
- generuje
- generujący
- otrzymać
- miejsce
- Dać
- dany
- Dający
- cel
- Goes
- dobry
- GPU
- GPU
- wykres
- wspaniały
- Zarządzanie
- Grupy
- Rośnie
- Wzrost
- poprowadzi
- ręka
- Prowadzenie
- się
- dzieje
- sprzęt komputerowy
- urządzenie sprzętowe
- urządzenia sprzętowe
- Have
- mający
- tutaj
- Najwyższa
- horyzont
- W jaki sposób
- How To
- Jednak
- HTTPS
- olbrzymi
- człowiek
- i
- CHORY
- idealny
- IDX
- obraz
- Rezultat
- wdrożenia
- realizacja
- ważny
- podnieść
- ulepszony
- in
- Zwiększać
- Zwiększenia
- coraz bardziej
- niezależny
- wskaźnik
- Informacja
- początkowy
- wkład
- zamiast
- instrukcje
- Inteligencja
- ciekawy
- wewnętrzny
- wprowadzono
- intuicja
- IT
- iteracja
- JEGO
- samo
- jpg
- Knuggety
- Trzymać
- Klawisz
- wiedza
- znany
- język
- duży
- większe
- Nazwisko
- Utajenie
- firmy
- warstwa
- nioski
- prowadzić
- dowiedziałem
- nauka
- Lista
- Popatrz
- poszukuje
- Partia
- zrobiony
- Główny
- poważny
- robić
- Dokonywanie
- zarządzanie
- zarządzający
- wiele
- mapa
- mapowanie
- Matrix
- Materia
- maksymalny
- znaczenie
- wymowny
- znaczy
- członek
- Pamięć
- metoda
- metody
- metryczny
- brakujący
- mieszanina
- model
- modele
- zmodyfikowano
- moduł
- jeszcze
- bardziej wydajny
- Ponadto
- większość
- ruch
- porusza się
- pomnożona
- Naturalny
- Natura
- Blisko
- niezbędny
- Potrzebować
- potrzebne
- ujemny
- sieci
- Nerwowy
- sieci neuronowe
- Nowości
- Następny
- nlp
- węzeł
- węzły
- nie-eksperci
- Pojęcie
- numer
- Nvidia
- uzyskane
- of
- Oferty
- on
- ONE
- koncepcja
- open source
- OpenAI
- działanie
- operacje
- Optymalny
- optymalizacja
- Optymalizacja
- zoptymalizowane
- Opcje
- oryginalny
- Inne
- Inaczej
- wydajność
- ogólny
- Papier
- Papiery
- parametr
- część
- szczególny
- szczególnie
- strony
- Ludzie
- doskonały
- jest gwarancją najlepszej jakości, które mogą dostarczyć Ci Twoje monitory,
- wykonywania
- faza
- sztuk
- plato
- Analiza danych Platona
- PlatoDane
- Grać
- gracz
- gra
- punkt
- polityka
- polityka
- position
- możliwy
- potencjał
- power
- Praktyczny
- praktyka
- przewidzieć
- Przewiduje
- preferencje
- przedstawione
- poprzedni
- prawdopodobieństwo
- prawdopodobnie
- Problem
- wygląda tak
- procesów
- produkować
- Wytworzony
- Produkt
- progresja
- progresywny
- obiecujący
- zaproponowane
- prawdopodobnie
- Sprawdzony
- publikować
- opublikowany
- RAM
- przypadkowy
- szeregi
- raczej
- osiągnięty
- Osiąga
- niedawno
- zmniejszyć
- Zredukowany
- redukcja
- rafinowany
- uczenie się wzmacniania
- wydany
- pozostały
- szczątki
- znakomity
- pamiętać
- zdalny
- powtórzony
- reprezentowanie
- reprezentuje
- wymagany
- Wymaga
- odpowiedzialny
- ograniczony
- dalsze
- wynikły
- Efekt
- powrót
- powraca
- zrewolucjonizować
- Nagradzać
- RZĄD
- rt
- run
- s
- taki sam
- Zapisz
- oszczędność
- scenariusz
- scenariusze
- schemat
- Szukaj
- druga
- Sekcja
- nasienie
- wybrany
- wybierając
- wybór
- SAMEGO SIEBIE
- zestaw
- ustawienie
- kilka
- Shape
- dzielenie
- Short
- powinien
- pokazać
- pokazane
- Targi
- Signal
- znaczący
- znacznie
- podobny
- podobieństwa
- Prosty
- prostota
- po prostu
- symulacja
- ponieważ
- sytuacja
- Rozmiar
- rozmiary
- mały
- mniejszy
- Migawka
- So
- rozwiązanie
- ROZWIĄZANIA
- Rozwiązywanie
- kilka
- Źródło
- Typ przestrzeni
- obowiązuje
- specyficzny
- swoiście
- określony
- prędkość
- wydać
- Spędzanie
- dzielić
- Kwadratowa
- ułożone w stos
- STAGE
- standard
- początek
- rozpocznie
- Stan
- state-of-the-art
- stwierdził,
- Zjednoczone
- statystyka
- Ewolucja krok po kroku
- Cel
- zatrzymanie
- sklep
- przechowywany
- sklep
- bezpośredni
- udany
- taki
- wsparcie
- podpory
- syntetyczny
- dane syntetyczne
- syntetycznie
- system
- systemy
- dostosowane
- Brać
- trwa
- biorąc
- cel
- Zadanie
- zadania
- Techniki
- terminal
- REGULAMIN
- że
- Połączenia
- Przyszłość
- Wykres
- Informacje
- Państwo
- ich
- Im
- a tym samym
- w związku z tym
- Te
- Trzeci
- trzy
- trójwymiarowy
- Przez
- czas
- czasochłonne
- do
- razem
- żeton
- tokenizacja
- tokenizowany
- Żetony
- także
- Zestaw narzędzi
- Top
- pochodnia
- Kwota produktów:
- tradycyjny
- Pociąg
- przeszkolony
- Trening
- pociągi
- trajektoria
- przekształcony
- leczyć
- prawdziwy
- zrozumieć
- Wszechświat
- nieużywana
- Aktualizacja
- zaktualizowane
- Nowości
- aktualizowanie
- uaktualnienie
- us
- Stosowanie
- posługiwać się
- Użytkownicy
- zazwyczaj
- wartość
- Wartości
- różnorodny
- wersja
- Wideo
- gier wideo
- Odwiedzić
- Wizyty
- W
- Droga..
- Co
- który
- Podczas
- szeroko
- Wikipedia
- będzie
- zwycięski
- w
- słowa
- stacja robocza
- wartość
- by
- pisanie
- napisany
- X
- zefirnet
- zero