Photo par DeepMind on Unsplash
La multiplication matricielle est une opération fondamentale utilisée dans de nombreux systèmes, des réseaux de neurones aux routines de calcul scientifique. Trouver des algorithmes efficaces et prouvés corrects pour la multiplication de matrices peut avoir un impact énorme sur la rapidité et l'efficacité des calculs, mais c'est une tâche très difficile. L'espace des algorithmes possibles est énorme et les méthodes traditionnelles de découverte d'algorithmes, telles que l'heuristique conçue par l'homme ou la recherche combinatoire, sont souvent sous-optimales.
DeepMindLa solution basée sur l'IA récemment proposée pour la recherche automatisée va bien au-delà de l'intuition humaine. La solution consiste en un agent d'apprentissage par renforcement profond appelé AlphaTensor, construit sur AlphaZero. Cet agent est formé pour jouer à un jeu solo, TensorGame, où le but est de découvrir des algorithmes de calcul efficaces pour la multiplication matricielle.
AlphaTensor est particulièrement efficace pour gérer les grandes matrices en décomposant les grandes multiplications matricielles en plus petites multiplications. De plus, AlphaTensor peut être utilisé pour obtenir des performances de pointe pour la multiplication matricielle une fois affiné sur un périphérique matériel spécifique.
AlphaTensor a un grand potentiel pour accélérer l'informatique d'apprentissage en profondeur. Dans l'apprentissage en profondeur, de nombreuses opérations chronophages peuvent être mappées sur des multiplications matricielles. En utilisant AlphaTensor pour optimiser ces opérations, les performances globales des modèles d'apprentissage en profondeur peuvent être considérablement améliorées.
Récemment, OpenAlphaTensor, le première implémentation open source d'AlphaTensor, a été publié, ce qui pourrait révolutionner la puissance de calcul des modèles d'apprentissage en profondeur.
Tenseur de multiplication matricielle
Pour les non-experts en optimisation de la multiplication matricielle, il peut ne pas être simple de comprendre comment une opération telle que la multiplication matricielle peut être mappée dans un tenseur tridimensionnel. Je vais essayer de l'expliquer avec des mots simples et avec des exemples.
Considérons le produit C = A*B, où pour simplifier A et B sont des matrices carrées de taille N. L'opération de multiplication peut être mappée dans un tenseur 3D de forme (N^2, N^2, N^2). La première dimension du tenseur représente la matrice aplatie A, la deuxième dimension la matrice aplatie B et la troisième dimension la matrice aplatie C.
Le tenseur n'a que des valeurs binaires (1 ou 0) pour chaque entrée. Notez que le tenseur représente l'opération de multiplication, il est donc indépendant des valeurs des matrices A et B.
Chaque entrée du tenseur correspond au coefficient de l'opération. Par exemple, pour calculer C[1,1], il est nécessaire de multiplier à la fois A[1,1] et B[1,1]. Par conséquent, l'entrée du tenseur [0,0,0], qui correspond à A[1,1], B[1,1] et C[1,1], aura la valeur 1. En revanche, pour calculer C[1,1 ,2,1], A[1] n'est pas nécessaire. Ainsi, la ligne du tenseur T[N+0, :, XNUMX] ne contiendra que des zéros.
L'image ci-dessous montre un exemple de tenseur pour N=2.
Image de DeepMind papier publié dans Nature
Comme illustré en (b) et (c) dans la figure ci-dessus, il est possible de mettre en œuvre un algorithme de calcul du produit à l'aide d'une décomposition du tenseur 3D. Plus précisément, l'algorithme ci-dessous peut être utilisé pour convertir une décomposition tensorielle (les matrices U, V, W) en un algorithme de multiplication matricielle.
Méta-algorithme paramétré pour calculer le produit matriciel C=AB introduit dans DeepMind papier
Le TensorGame
Le problème de trouver des algorithmes efficaces pour la multiplication matricielle est extrêmement difficile car le nombre d'algorithmes possibles à considérer est beaucoup plus grand que le nombre d'atomes dans l'univers, même pour de petites instances de multiplication matricielle.
DeepMind a converti ce problème en un jeu solo et l'a appelé le TensorGame. Dans ce jeu, le joueur choisit comment combiner différentes entrées de matrices pour les multiplier. Un score est attribué en fonction du nombre d'opérations nécessaires pour obtenir le bon résultat de multiplication. Le jeu se termine lorsque le tenseur zéro est atteint ou lorsque le nombre maximum de coups a été effectué. La factorisation finale est évaluée sur la base d'une estimation du rang résiduel et de certains critères d'optimisation, tels que la complexité temporelle asymptotique ou le temps d'exécution pratique.
La position initiale dans le TensorGame correspond au Matrix Multiplication Tensor exprimé sur une base aléatoire.
A chaque étape t du jeu, le joueur écrit trois vecteurs
, qui spécifie les tenseurs de rang 1 . L'état du jeu est mis à jour en soustrayant les vecteurs sélectionnés par le joueur :
De
est le tenseur de multiplication matricielle.Si le jeu se termine en p étapes, cela signifie que le tenseur de multiplication matricielle
peut être décomposé en p tenseurs de rang 1 , c'est-à-dire qu'il a au moins le rang p.Le TensorGame peut alors être interprété comme un algorithme de décomposition de rang et AlphaTensor peut être vu comme un algorithme d'estimation du rang du tenseur.
Architecture AlphaTensor
Jusqu'à présent, nous avons découvert le TensorGame et expliqué comment sa solution peut être considérée comme un algorithme de multiplication matricielle. Explorons maintenant les principaux concepts d'AlphaTensor, l'algorithme utilisé pour le jeu.
L'architecture AlphaTensor est essentiellement une architecture de transformateur d'encodeur-décodeur où :
- l'encodeur prend en entrée l'état du jeu , les n actions précédentes prises par le modèle (généralement n=7), et l'indice de temps t de l'action en cours. Les informations sont empilées dans un tenseur de forme (n+1, N^2, N^2, N^2). Ce tenseur est ensuite remodelé et transformé (à l'aide de trois couches linéaires) en un tenseur de forme (N^2, N^2, c) où c est la dimension interne du modèle.
- le décodeur génère les actions n_steps à partir du vecteur embarqué donné par l'encodeur de manière auto-régressive. Chaque action correspond à un jeton des triplets représentant l'un des triplets décomposant le tenseur de jeu (ie réduisant son rang)
Le modèle est formé en alternant la rétropropagation et l'action du modèle. L'action du modèle est utilisée pour générer des données qui sont ensuite utilisées pour former le modèle. En pratique, le modèle est formé avec un mélange de données générées synthétiquement et de données générées par le modèle pendant l'action. L'étape d'action se fait en prenant un tenseur 3D correspondant à une opération matricielle et en jouant à des jeux n_actors dessus. Chaque acteur joue un jeu soit sur la base standard, soit sur une base alternative (le changement de base est appliqué avec une probabilité donnée). Les résultats sont ensuite collectés et peuvent être utilisés dans l'étape d'apprentissage avec les données synthétiques.
L'étape d'action est basée sur le Monte Carlo Tree Search (MCTS) d'AlphaZero, modifié pour prendre en charge de grands espaces d'action. En bref, avant de choisir l'action, les chemins n_sims sont explorés à partir de la sortie du modèle avec une exploration future maximale de 5 étapes. Les probabilités générées par le modèle sont ensuite ajustées en tenant compte des chemins générés. Ensuite, l’action présentant le ou les chemins futurs les plus prometteurs est choisie pour continuer le jeu.
Lors de la formation du modèle, la récompense est en fait une récompense négative (pénalité). Sa valeur absolue augmente avec chaque étape supplémentaire requise pour résoudre le jeu. Si le modèle prend m étapes pour résoudre un TensorGame, la récompense associée au jeu est r=-m. Si le modèle n'est pas capable de résoudre le TensorGame en étapes max_rank, la récompense est calculée en estimant le rang du tenseur restant. Le rang est estimé comme la somme des rangs des matrices qui composent le tenseur. L'estimation est une borne supérieure sur le vrai rang du tenseur.
Lors du réglage fin du modèle, la pénalité à l’état terminal doit également prendre en compte la latence de l’algorithme produit par le modèle. La formule de récompense devient rt'=rt+λbt, où rt est le système de récompense décrit précédemment, bt est la récompense de référence (non nulle uniquement à l'état terminal), et λ est un coefficient spécifié par l'utilisateur.
Accélérations (%) des algorithmes découverts par AlphaTensor et adaptés pour un GPU et un TPU, extraits de l'article de DeepMind. Les accélérations sont mesurées par rapport à la multiplication matricielle standard (par exemple cuBLAS pour le GPU) sur le même matériel et comparées à la Algorithme de Strassen-carré. La source: DeepMind.
j'ai récemment sorti OuvrirAlphaTensor, la première implémentation open source d'AlphaTensor. Dans cette section, je vais parcourir la mise en œuvre. Comme nous l'avons vu précédemment, l'architecture AlphaTensor est assez simple, basée sur un transformateur standard avec une architecture encodeur-décodeur. Les composants les plus intéressants d'AlphaTensor sont la première couche de la partie encodeur et la façon dont les actions sont échantillonnées.
Commençons par la première couche d'encodage.
# 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
Dans l'extrait ci-dessus, nous montrons comment le tenseur d'entrée est décomposé en trois tenseurs, qui sont ensuite utilisés comme entrées de requête, de clé et de valeur de la couche de transformateur.
- Sur les trois dimensions du tenseur représentant les matrices aplaties (A, B, C), le tenseur d'entrée est aplati le long de chaque dimension avec la dimension représentant les actions précédentes. Ainsi, dans chaque copie aplatie du tenseur d'entrée, la dimension sélectionnée est une agrégation des dernières valeurs T-1 et de la valeur réelle, pour toutes les valeurs S de la dimension sélectionnée, où S=N^2. Philosophiquement, c'est comme si, pour chaque dimension, nous nous concentrions sur ce qui s'est passé dans les actions précédentes dans cette dimension.
- Les scalaires sont cartographiés dans trois espaces différents de dimension S^2, puis remodelés pour être concaténés avec les tenseurs obtenus au point précédent. Conceptuellement, les scalaires sont mappés sur un espace d'intégration de dimension S ^ 2, puis les informations intégrées sont découpées en vecteurs S et empilées, comme ce qui arrive au texte lorsqu'il est symbolisé.
- Les jetons scalaires sont concaténés avec le tenseur d'entrée restructuré, puis donnés en entrée à une couche linéaire pour mapper les informations de mise au point scalaires + historique de canal dans la dimension interne du modèle.
Ces trois étapes peuvent être interprétées comme un moyen de donner au modèle à la fois des informations sur les scalaires (comme dans le pas de temps TensorGame) et le focus sur les actions précédentes pour chaque canal.
Concernant la façon dont les actions sont produites, il est intéressant de noter qu'AlphaTensor génère en sortie le triplet u, v, w, qui vise à réduire le rang du tenseur. Les trois vecteurs ont une taille S et puisqu'ils sont concaténés, le modèle doit produire un vecteur de taille 3*S. AlphaTensor est formé avec un algorithme RL, donc toutes les actions possibles doivent être exprimées en termes de probabilités dans un espace énuméré, c'est-à-dire que le modèle produit une probabilité sur les différentes actions. Cela signifie que chaque vecteur dans l'espace 3S doit être mappé à une action différente. Il en résulte un espace d'action de taille |F|^(3S), où |F| est le nombre de valeurs différentes que peut prendre l'élément de u, v, w. Habituellement, les valeurs sont limitées à (-2, -1, 0, 1, 2), ce qui donne une cardinalité de 5 éléments.
Voici un défi majeur : pour générer les probabilités d'action pour un produit matriciel de matrices de taille 5, nous aurions besoin d'une mémoire de 5^75 * 4 octets, ce qui signifierait ~10^44 Go de mémoire. De toute évidence, nous ne pouvons pas gérer un si grand espace d'action.
Comment résolvons-nous le problème ? Pour réduire l'empreinte mémoire des probabilités d'action, nous pouvons diviser les triplets en plus petits morceaux, les "tokéniser" et traiter les morceaux comme des jetons générés dans l'architecture du transformateur, c'est-à-dire que les jetons sont donnés en entrée au décodeur dans un auto-régressif chemin. Dans l'exemple ci-dessus, nous pouvons diviser les triplets en 15 morceaux, réduisant la consommation de mémoire à 15 * 5 ^ (75/15) * 4, soit 187.5 Ko.
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), )
Ci-dessus, nous montrons l'extrait de code pour générer l'action complète. Dans le code, self.core contient la couche décodeur et le tenseur e représente la sortie de la couche encodeur. Zéro peut être considéré comme le dans les modèles NLP et les actions n_steps représentant les morceaux n_steps sont générées de manière progressive.
Le modèle renvoie trois quantités :
- Les actions générées
- La probabilité associée à l'action complète
- Les logits produits pour générer la première action (le premier morceau) qui sera utilisé pour calculer la valeur du modèle.
Cela vaut la peine de passer quelques mots sur le paramètre n_samples. Le paramètre est utilisé pour l'étape d'action et il permet au modèle de générer différentes versions des triplets qui seront ensuite utilisées pour explorer l'espace d'action dans l'algorithme Monte Carlo Tree Search utilisé dans le processus d'action. Les différentes actions n_samples sont échantillonnées en fonction de la politique générée par le modèle.
Étape par intérim
La partie la plus délicate de tout l'algorithme est probablement l'étape Acting utilisée pour résoudre le TensorGame. L'algorithme n'est pas expliqué en profondeur dans l'article AlphaTensor, car il est basé sur plusieurs articles précédents de DeepMind qui sont juste cités et donnés comme connus. Ici, je vais reconstruire toutes les pièces manquantes et expliquer étape par étape notre mise en œuvre.
Nous pouvons organiser les étapes d'action en trois composantes différentes :
- L'arborescence de Monte-Carlo
- La simulation de jeu
- Le calcul de politique amélioré
Analysons-les un par un.
Recherche arborescente de Monte-Carlo (MCTS)
Monte Carlo Tree Search (MCTS) est une technique d'intelligence artificielle largement utilisée pour le jeu, en particulier dans les jeux de société et les jeux vidéo. L'algorithme crée un arbre de jeu qui simule les mouvements et les résultats potentiels et utilise un échantillonnage aléatoire pour évaluer la récompense attendue pour chaque mouvement. L'algorithme sélectionne ensuite de manière itérative le mouvement avec la récompense attendue la plus élevée et simule les résultats jusqu'à ce qu'il atteigne un état terminal ou une condition d'arrêt spécifiée. Les simulations sont utilisées pour estimer la probabilité de gagner pour chaque coup et guider le processus de prise de décision. MCTS s'est avéré efficace dans les jeux complexes où le nombre de mouvements et de résultats possibles est important, et il a été utilisé dans des systèmes d'IA de jeu réussis, tels qu'AlphaGo.
Dans AlphaTensor, une version modifiée du MCTS original est utilisée. En particulier, au lieu de sélectionner aléatoirement l'action dans tout l'espace des actions, l'action est sélectionnée parmi un sous-ensemble généré directement par le modèle (à travers les n_échantillons présentés précédemment). La correction apportée à la mise à niveau de la politique est ensuite appliquée à l'étape de calcul de la politique améliorée.
Dans notre implémentation, nous avons décidé de conserver toutes les informations sur l'arbre de Monte-Carlo dans un dictionnaire ayant comme clé la version hachée de l'état TensorGame et comme valeurs les informations associées à l'état lui-même. Chaque étape de Monte-Carlo part d'un nœud et simule des mini-jeux n_sim, explorant le futur avec un horizon de 5 coups. Si le nœud a déjà été exploré dans des simulations précédentes, n_sim est ajusté en tenant compte du nombre d'explorations précédentes. Pour chaque nœud, le nombre de visites est stocké dans le tenseur N_s_a, puisque ce tenseur contient le nombre de visites par action enfant de nœud (parmi celles échantillonnées par le modèle).
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
Le code ci-dessus montre notre implémentation de l'algorithme. Pour une question de simplicité de code, la correction de politique est effectuée dans la fonction simulator_game.
Simulation de jeu
La fonction simulator_game est chargée d'explorer l'arbre composé de nœuds représentant un état particulier du TensorGame. Il exécute également le modèle chaque fois qu'un nœud feuille est rencontré et il stocke toutes les informations de nœud dans le dictionnaire state_dict. Examinons en profondeur sa mise en œuvre :
@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)
Chaque simulation est divisée en trois parties :
- Sélection
- Expansion
- sauvegarde
Dans la partie sélection, la simulation est exécutée sur les nœuds d'arbre déjà générés, et le nœud suivant est sélectionné à l'aide de la fonction suivante :
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()]
En pratique, l'action maximisant la fonction ucb :
pour l'état donné est sélectionné. Ici, Q représente les valeurs Q générées par le modèle et π représente la distribution aléatoire sur les actions échantillonnées à l'aide de la politique du modèle. N(s, a) représente le nombre de visites du nœud à l'action a depuis le nœud s.
Une fois que la phase de sélection atteint un nœud feuille, si la simulation n'a pas atteint une condition terminale (en termes d'exploration maximale, c'est-à-dire d'horizon futur ou de fin de partie), le modèle est alors utilisé pour sélectionner n_échantillons nœuds alternatifs (ils seront feuille nœuds dans l'itération successive). C'est ce qu'on appelle la phase d'expansion, car de nouveaux nœuds sont ajoutés à l'arbre. Ensuite, aucun autre nœud n'est exploré dans la simulation en cours, mais la feuille q_value est envoyée à l'étape de simulation suivante : la sauvegarde.
La sauvegarde est la dernière étape de chaque simulation. Pendant la sauvegarde, si le nœud feuille était un état terminal, la récompense finale est calculée ; sinon, la valeur q de la feuille est utilisée comme récompense estimée. Ensuite, la récompense est rétro-propagée sur la trajectoire de simulation en mettant à jour à la fois les états q_values et en mettant à jour le compteur de visites N(s, a). Dans l'extrait ci-dessous, nous montrons le code pour la rétro-propagation de la récompense.
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
Calcul de politique amélioré
Une fois que toutes les simulations ont été exécutées et que le MCTS offre un aperçu intéressant du futur proche, il est temps de mettre à jour la politique associée aux nœuds prédits et de les renvoyer, afin qu'ils puissent être utilisés pendant la formation. La politique améliorée, suivant la méthode décrite dans Hubert et coll., est utilisé pour gérer de grands espaces d'action. En fait, pour un petit espace de recherche, il est possible pendant MCTS d'échantillonner une action au hasard dans l'espace d'action et d'évaluer son impact. Une approche similaire dans un espace d'action beaucoup plus large conduirait à toutes les trajectoires divergentes dans des chemins différents et il faudrait une quantité infinie de trajectoires pour obtenir des statistiques significatives et ensuite mettre à jour la politique. Étant donné qu'ici, nous utilisons sample-MCTS pour éviter la dispersion, c'est-à-dire que les actions n_samples sont échantillonnées conformément à la politique du modèle, puis que MCTS sélectionne simplement l'une des actions échantillonnées lors de l'exploration de l'arbre, nous devons prendre en compte la correction de l'échantillon lors du calcul la stratégie finale mise à jour qui sera utilisée lors de la formation du modèle.
En pratique, la politique améliorée est calculée comme
De
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
Notez que dans notre implémentation, après avoir calculé la politique à partir du tenseur N_s_a, nous devons la remettre en correspondance avec le tenseur d'action d'origine. En fait, N_s_a ne considère que les actions échantillonnées par le modèle, tandis que la politique finale doit également contenir des probabilités pour les actions non explorées.
Différences par rapport à l'algorithme d'entraînement ChatGPT
AlphaTensor est le dernier membre de la famille AlphaGo/AlphaZero de méthodes d'intelligence artificielle de DeepMind. Ces méthodes sont basées sur l'algorithme Monte Carlo Tree Search (MCTS), qui a été affiné et amélioré par DeepMind pour s'attaquer à des tâches de plus en plus complexes. Un autre système d'IA, ChatGPT d'OpenAI, qui a fait beaucoup de bruit pour ses performances remarquables, a été formé avec une approche différente, appelée Reinforcement Learning with Human Feedback (RLHF).
RLHF est une technique de réglage fin utilisée pour ajuster les modèles de langage afin de suivre un ensemble d'instructions écrites. Il utilise les préférences humaines comme signal de récompense pour affiner le modèle, alignant ainsi le comportement du modèle linguistique sur les préférences déclarées d'un groupe spécifique de personnes, plutôt que sur une notion plus large de «valeurs humaines».
En revanche, MCTS est un algorithme de recherche arborescent utilisé pour déterminer les mouvements optimaux dans les jeux. Il simule les mouvements potentiels et met à jour les valeurs de chaque mouvement en fonction de leurs résultats, guidant la sélection du meilleur mouvement.
RLHF collecte des données à partir de démonstrations écrites par l'homme et de comparaisons étiquetées par l'homme entre des modèles d'IA, et forme un modèle de récompense pour prédire les préférences d'un groupe de personnes donné. Le modèle de récompense est ensuite utilisé pour affiner les modèles d'IA. MCTS, d'autre part, utilise des simulations et des évaluations pour déterminer la meilleure décision.
Bien qu'il s'agisse d'approches différentes, RLHF et MCTS présentent également des similitudes. Les deux techniques d'intelligence artificielle utilisent des méthodes de prise de décision et de résolution de problèmes, et toutes deux utilisent une approche par essais et erreurs pour explorer différentes options et prendre des décisions en fonction des informations disponibles. Les deux sont également des processus itératifs qui s'améliorent au fil du temps à mesure que davantage d'informations et d'expériences sont recueillies.
Le choix entre RLHF et MCTS dépend de la tâche à accomplir. RLHF est idéal lorsqu'il n'y a pas de métrique claire pour évaluer les performances du modèle, tandis que MCTS s'est avéré efficace dans les tâches de type jeu où la connaissance et l'exploration de l'avenir donnent au modèle un avantage significatif.
Optimisation du code pour la formation AlphaTensor
La mise en œuvre de l'algorithme d'entraînement AlphaTensor nécessite de trouver le parfait compromis entre la vitesse d'entraînement et la consommation de mémoire. Comme indiqué dans la section Modèle, le simple fait de considérer la tokenisation de l'action peut économiser beaucoup de mémoire, mais une réduction trop agressive de l'espace d'action peut entraîner à la fois une baisse de la précision et un ralentissement des performances. Ce dernier se produit parce que tous les jetons sont générés séquentiellement de manière autorégressive par le décodeur de modèle. Par conséquent, le temps d'inférence augmente de manière linéaire avec le nombre de jetons par action une fois que le softmax sur l'espace d'action n'est plus le goulot d'étranglement.
Lors de la mise en place de la formation AlphaTensor, les principales difficultés ont été rencontrées dans la gestion du processus d'action. Si les tenseurs ne sont pas stockés dans le bon format, le MCTS peut facilement provoquer une croissance incontrôlée de l'utilisation de la mémoire. En revanche, si le nombre de tenseurs stockés lors de chaque simulation est trop réduit, le MCTS peut passer un temps infini à recalculer les états requis.
Prenons un exemple de l'étape de simulation de jeu, où le jeu est exploré en examinant de possibles scénarios futurs. Pour chaque état, si nous ne sauvegardons pas les actions générées par le modèle et que nous décidons de sauvegarder uniquement la graine aléatoire utilisée pour échantillonner les actions de la politique, alors chaque fois que nous explorons un nœud d'arborescence, nous devrons recalculer la politique et puis échantillonnez les actions. En clair, nous avons décidé de stocker les actions échantillonnées pour gagner du temps et éviter d'avoir à gérer le partage de modèles entre différents processus dans le cas de la parallélisation d'exploration MCTS. Cependant, il ne suffisait pas de sauvegarder les actions pour obtenir une démarche d'action suffisamment efficace. En fait, le temps nécessaire pour convertir les actions n_steps en triplet (u, v, w), réduire l'état du tenseur du jeu et créer les nouveaux tenseurs 3D à partir des actions n_samples serait facilement un goulot d'étranglement pour l'ensemble de la formation. Deuxièmement, nous ne voulions pas stocker tous les états futurs possibles pour chaque action échantillonnée, car cela aurait un impact énorme sur la mémoire utilisée par l'algorithme. Supposons que nous définissons n_samples=32, n=7 et N=5, et rappelons-nous que N est la taille du produit matriciel carré que nous voulons réduire et n est le nombre d'actions précédentes mémorisées par le modèle. Dans cette situation, chaque tenseur d'état aurait la forme (8, 25, 25, 25), ce qui multiplié par 32 donnerait 3282525254 octets pour chaque nœud du graphe. Maintenant, en considérant que chaque simulation dans la phase d'expansion génère un nouveau nœud (et n_sim=200), nous aurions une consommation mémoire finale de 200328252525*4 = 3.2 Go pour le premier nœud MCTS seul. Dans le pire des cas, tout en explorant les nœuds max_rank actifs (où max_rank = 150), cela entraînerait une consommation totale de mémoire de 150 * 3.2 Go = 480 Go en mémoire RAM (ou en mémoire GPU si tous les tenseurs étaient stockés sur le GPU) . Nous avons exécuté la formation sur notre poste de travail avec 128 Go de RAM et 48 Go de mémoire GPU, nous avons donc dû réduire la consommation de mémoire.
Comme nous ne voulions pas augmenter le temps d'exécution, nous avons adopté une optimisation qui exploite la redondance des tenseurs d'état produits. En fait, les tenseurs ont n-1 actions précédentes en commun, qui peuvent ensuite être stockées une fois et non répétées pour chaque tenseur stocké. Cela entraîne une réduction de la mémoire de 2/7 à 28 %, ce qui signifie que dans le pire des cas, 137 Go peuvent être stockés. À ce stade, en élaguant simplement la partie inutilisée de l'arbre (comme les trajectoires non sélectionnées) et en stockant les tenseurs dans la mémoire CPU, nous avons pu éviter toute erreur de mémoire lors de l'entraînement.
Avec OpenAlphaTensor étant désormais open source, plusieurs avenues passionnantes pour un développement ultérieur s'ouvrent.
Une progression naturelle est le réglage fin d'OpenAlphaTensor sur les périphériques matériels cibles. Cela devrait conduire à des performances de calcul très compétitives. Je publierai plus sur les performances d'OpenAlphaTensor sur divers matériels sur GitHub. Au moment de la rédaction de cet article, OpenAlphaTensor était en cours de formation.
Une autre avancée importante serait la prise en charge de la compilation à distance, permettant aux utilisateurs de créer des algorithmes optimisés pour les appareils périphériques. Ceci peut être réalisé en stockant le modèle OpenAlphaTensor sur un serveur, tandis que l'algorithme de multiplication matricielle est évalué sur un matériel différent.
Il pourrait également être important d'étendre la prise en charge de différents compilateurs pour calculer la correction de récompense basée sur la latence. Différents compilateurs peuvent conduire à différents algorithmes optimisés sur un matériel donné. Par exemple, l'article DeepMind a montré des résultats prometteurs en utilisant JAX et le compilateur XLA sur les GPU TPU et Nvidia. Il serait intéressant d'évaluer cela en utilisant NCCL sur Nvidia ou LLVM sur les CPU.
Enfin, l'extension du modèle et de l'algorithme de formation pour prendre en charge des tailles de matrice plus grandes reste un défi ouvert majeur. Actuellement, OpenAlphaTensor prend en charge une taille de matrice maximale de 5, mais il peut être appliqué en divisant des multiplications matricielles plus importantes en groupes de minuscules MM d'une taille inférieure à 5. Cette approche est sous-optimale et effectue la réduction directement sur le grand tenseur correspondant au un MM complet pourrait théoriquement conduire à de meilleurs résultats.
Diego Fiori est le CTO de Nebuly AI, une entreprise qui s'engage à intégrer l'optimisation de l'IA à la boîte à outils de chaque développeur.
- Contenu propulsé par le référencement et distribution de relations publiques. Soyez amplifié aujourd'hui.
- Platoblockchain. Intelligence métaverse Web3. Connaissance Amplifiée. Accéder ici.
- La source: 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
- :est
- ][p
- $UP
- 1
- 3d
- 8
- a
- Capable
- À propos
- au dessus de
- Absolute
- accélérer
- Selon
- en conséquence
- Compte
- précision
- atteindre
- atteint
- Action
- actes
- actually
- ajoutée
- Supplémentaire
- Ajusté
- adopté
- avancer
- Avantage
- Après
- Agent
- agrégation
- agressif
- AI
- Systèmes d'IA
- vise
- algorithme
- algorithmes
- Tous
- Permettre
- permet
- seul
- déjà
- alternative
- parmi
- montant
- il analyse
- ainsi que
- Une autre
- appliqué
- une approche
- approches
- architecture
- SONT
- article
- artificiel
- intelligence artificielle
- AS
- attribué
- associé
- At
- Automatisation
- disponibles
- en évitant
- RETOUR
- sauvegarde
- basé
- En gros
- base
- BE
- car
- devient
- before
- va
- ci-dessous
- référence
- LES MEILLEURS
- Améliorée
- jusqu'à XNUMX fois
- Au-delà
- planche
- Jeux de société
- lié
- plus large
- BT
- construire
- construit
- by
- appelé
- CAN
- ne peut pas
- maisons
- Causes
- causé
- certaines
- challenge
- difficile
- Change
- Développement
- ChatGPT
- enfant
- le choix
- choose
- choisi
- cité
- clair
- clairement
- code
- recueille
- combiner
- engagé
- Commun
- Société
- par rapport
- compétitif
- complexe
- complexité
- composants électriques
- composé
- compromis
- calcul
- puissance de calcul
- calcul
- informatique
- concepts
- Conceptuellement
- condition
- confiance
- Considérer
- considéré
- considérant
- considère
- consommation
- contient
- continuer
- contraste
- converti
- Core
- Correspondant
- correspond
- pourriez
- Counter
- Processeur
- crée des
- La création
- critères
- CTO
- Courant
- Lecture
- données
- traitement
- décider
- décidé
- décision
- La prise de décision
- décisions
- profond
- l'apprentissage en profondeur
- DeepMind
- dépend
- décrit
- Déterminer
- Développeur
- Développement
- dispositif
- Compatibles
- DICT
- différent
- problèmes
- Dimension
- dimensions
- directement
- découvrez
- découverte
- discuté
- distribution
- divisé
- down
- Goutte
- pendant
- e
- chacun
- Plus tôt
- même
- Edge
- Efficace
- efficace
- non plus
- élément
- éléments
- intégré
- se termine
- améliorée
- énorme
- assez
- entrée
- erreur
- estimation
- estimé
- Ether (ETH)
- évaluer
- évalué
- évaluer
- évaluations
- Pourtant, la
- Chaque
- exemple
- exemples
- passionnant
- exécution
- avec des données
- attendu
- Découvrez
- Expliquer
- expliqué
- exploits
- exploration
- explorez
- Exploré
- Explorant
- exprimé
- étendre
- extension
- extrêmement
- équitablement
- famille
- plus rapide
- Réactions
- few
- Figure
- finale
- trouver
- Prénom
- flotteur
- Focus
- suivre
- Abonnement
- numérique
- Pour
- formulaire
- le format
- formule
- trouvé
- de
- plein
- fonction
- fondamental
- plus
- la poursuite du développement
- avenir
- jeu
- Games
- générer
- généré
- génère
- générateur
- obtenez
- obtention
- Donner
- donné
- Don
- objectif
- Goes
- Bien
- GPU
- GPU
- graphique
- l'
- Réservation de groupe
- Groupes
- Pousse
- Croissance
- guide
- main
- Maniabilité
- arrivé
- arrive
- Matériel
- périphérique matériel
- périphériques matériels
- Vous avez
- ayant
- ici
- le plus élevé
- horizon
- Comment
- How To
- Cependant
- HTTPS
- majeur
- humain
- i
- MAUVAIS
- idéal
- IDX
- image
- Impact
- Mettre en oeuvre
- la mise en oeuvre
- important
- améliorer
- amélioré
- in
- Améliore
- Augmente
- de plus en plus
- indépendant
- indice
- d'information
- initiale
- contribution
- plutôt ;
- Des instructions
- Intelligence
- intéressant
- interne
- introduit
- intuition
- IT
- itération
- SES
- lui-même
- jpg
- KDnuggetsGenericName
- XNUMX éléments à
- clés / KEY :
- spécialisées
- connu
- langue
- gros
- plus importantes
- Nom de famille
- Latence
- Nouveautés
- couche
- poules pondeuses
- conduire
- savant
- apprentissage
- Liste
- Style
- recherchez-
- Lot
- LES PLANTES
- Entrée
- majeur
- faire
- Fabrication
- gérer
- les gérer
- de nombreuses
- Localisation
- cartographie
- Matrice
- Matière
- maximales
- sens
- significative
- veux dire
- membre
- Mémoire
- méthode
- méthodes
- métrique
- manquant
- mélange
- modèle
- numériques jumeaux (digital twin models)
- modifié
- module
- PLUS
- plus efficace
- Par ailleurs
- (en fait, presque toutes)
- Bougez
- se déplace
- multiplié
- Nature
- Nature
- Près
- nécessaire
- Besoin
- nécessaire
- négatif
- réseaux
- Neural
- les réseaux de neurones
- Nouveauté
- next
- nlp
- nœud
- nœuds
- non-spécialistes
- Notion
- nombre
- Nvidia
- obtenu
- of
- Offres Speciales
- on
- ONE
- ouvert
- open source
- OpenAI
- opération
- Opérations
- optimaux
- à mettre en œuvre pour gérer une entreprise rentable. Ce guide est basé sur trois décennies d'expérience
- Optimiser
- optimisé
- Options
- original
- Autre
- autrement
- sortie
- global
- Papier
- papiers
- paramètre
- partie
- particulier
- particulièrement
- les pièces
- Personnes
- parfaite
- performant
- effectuer
- phase
- pièces
- Platon
- Intelligence des données Platon
- PlatonDonnées
- Jouez
- joueur
- jouer
- Point
- politiques
- politique
- position
- possible
- défaillances
- power
- Méthode
- pratique
- prévoir
- prédit
- préférences
- présenté
- précédent
- probabilité
- Probablement
- Problème
- processus
- les process
- produire
- Produit
- Produit
- progression
- progressif
- prometteur
- proposé
- prouvable
- proven
- publier
- publié
- RAM
- aléatoire
- rangs
- plutôt
- atteint
- atteint
- récemment
- réduire
- Prix Réduit
- réduire
- raffiné
- apprentissage par renforcement
- libéré
- restant
- reste
- remarquables
- rappeler
- éloigné
- répété
- représentation
- représente
- conditions
- a besoin
- responsables
- limité
- résultat
- résultant
- Résultats
- retourner
- Retours
- révolutionner
- Récompenser
- RANGÉE
- rt
- Courir
- s
- même
- Épargnez
- économie
- scénario
- scénarios
- programme
- Rechercher
- Deuxièmement
- Section
- seed
- choisi
- la sélection
- sélection
- AUTO
- set
- mise
- plusieurs
- Forme
- partage
- Shorts
- devrait
- montrer
- montré
- Spectacles
- Signal
- significative
- de façon significative
- similaires
- similitudes
- étapes
- simplicité
- simplement
- simulation
- depuis
- situation
- Taille
- tailles
- petit
- faibles
- Instantané
- So
- sur mesure
- RÉSOUDRE
- Résoudre
- quelques
- Identifier
- Space
- espaces
- groupe de neurones
- spécifiquement
- spécifié
- vitesse
- passer
- Dépenses
- scission
- carré
- empilé
- Étape
- Standard
- Commencer
- départs
- Région
- state-of-the-art
- A déclaré
- États
- statistiques
- étapes
- Étapes
- arrêt
- Boutique
- stockée
- STORES
- simple
- réussi
- tel
- Support
- Les soutiens
- haute
- données synthétiques
- synthétiquement
- combustion propre
- Système
- Prenez
- prend
- prise
- Target
- Tâche
- tâches
- techniques
- terminal
- conditions
- qui
- La
- El futuro
- Le graphique
- les informations
- L'État
- leur
- Les
- ainsi
- donc
- Ces
- Troisièmement
- trois
- tridimensionnel
- Avec
- fiable
- long
- à
- ensemble
- jeton
- tokenization
- tokenisé
- Tokens
- trop
- Boîte à outils
- top
- torche
- Total
- traditionnel
- Train
- qualifié
- Formation
- les trains
- trajectoire
- transformé
- traiter
- oui
- comprendre
- Univers
- inutilisé
- Mises à jour
- a actualisé
- Actualités
- la mise à jour
- améliorer
- us
- Utilisation
- utilisé
- utilisateurs
- d'habitude
- Plus-value
- Valeurs
- divers
- version
- Vidéo
- jeux vidéo
- Visiter
- Visites
- W
- Façon..
- Quoi
- qui
- tout en
- largement
- Wikipédia
- sera
- une équipe qui gagne ?
- avec
- des mots
- poste de travail
- vaut
- pourra
- écriture
- code écrit
- X
- zéphyrnet
- zéro