Foto por DeepMind on Unsplash
A multiplicação de matrizes é uma operação fundamental usada em muitos sistemas, desde redes neurais até rotinas de computação científica. Encontrar algoritmos eficientes e comprovadamente corretos para a multiplicação de matrizes pode ter um grande impacto em tornar a computação mais rápida e eficiente, mas é uma tarefa muito desafiadora. O espaço de algoritmos possíveis é enorme, e os métodos tradicionais para descobrir algoritmos, como heurística projetada por humanos ou pesquisa combinatória, geralmente não são ideais.
DeepMindA solução recentemente proposta baseada em IA para pesquisa automatizada vai muito além da intuição humana. A solução consiste em um agente de aprendizagem por reforço profundo chamado AlphaTensor, construído sobre AlfaZero. Este agente é treinado para jogar um jogo single-player, TensorGame, onde o objetivo é descobrir algoritmos computacionalmente eficientes para multiplicação de matrizes.
O AlphaTensor é particularmente bom em lidar com grandes matrizes, decompondo grandes multiplicações de matrizes em multiplicações menores. Além disso, o AlphaTensor pode ser usado para obter desempenho de ponta para multiplicação de matrizes, uma vez ajustado em um dispositivo de hardware específico.
O AlphaTensor tem grande potencial para acelerar a computação de aprendizado profundo. No aprendizado profundo, muitas operações demoradas podem ser mapeadas para multiplicações de matrizes. Ao usar o AlphaTensor para otimizar essas operações, o desempenho geral dos modelos de aprendizagem profunda pode ser significativamente melhorado.
Recentemente, OpenAlphaTensor, o primeira implementação de código aberto do AlphaTensor, foi lançado, o que poderia revolucionar o poder computacional dos modelos de aprendizado profundo.
Tensor de multiplicação de matrizes
Para não especialistas em otimização de multiplicação de matrizes, pode não ser simples entender como uma operação como a multiplicação de matrizes pode ser mapeada em um tensor tridimensional. Vou tentar explicar com palavras simples e com exemplos.
Vamos considerar o produto C = A*B, onde para simplificar A e B são matrizes quadradas de tamanho N. A operação de multiplicação pode ser mapeada em um tensor 3D de formato (N^2, N^2, N^2). A primeira dimensão do tensor representa a matriz achatada A, a segunda dimensão a matriz achatada B e a terceira dimensão a matriz achatada C.
O tensor tem apenas valores binários (1 ou 0) para cada entrada. Observe que o tensor representa a operação de multiplicação, portanto é independente dos valores das matrizes A e B.
Cada entrada do tensor corresponde ao coeficiente da operação. Por exemplo, para calcular C[1,1], é necessário multiplicar A[1,1] e B[1,1]. Portanto, a entrada do tensor [0,0,0], que corresponde a A[1,1], B[1,1] e C[1,1], terá valor 1. Em contraste, para calcular C[1,1 ,2,1], A[1] não é necessário. Assim, a linha do tensor T[N+0, :, XNUMX] conterá apenas zeros.
A imagem abaixo mostra um exemplo de tensor para N=2.
Imagem do DeepMind papel publicado em Natureza
Conforme mostrado em (b) e (c) na figura acima, é possível implementar um algoritmo para calcular o produto usando uma decomposição do tensor 3D. Mais especificamente, o algoritmo abaixo pode ser usado para converter uma decomposição tensorial (as matrizes U, V, W) em um algoritmo de multiplicação de matrizes.
Meta-algoritmo parametrizado para calcular o produto matricial C = AB introduzido no DeepMind's papel
O TensorGame
O problema de encontrar algoritmos eficientes para multiplicação de matrizes é extremamente desafiador porque o número de algoritmos possíveis a serem considerados é muito maior que o número de átomos no universo, mesmo para pequenas instâncias de multiplicação de matrizes.
DeepMind converteu esse problema em um jogo para um jogador e o chamou de TensorGame. Neste jogo, o jogador escolhe como combinar diferentes entradas de matrizes para multiplicá-las. Uma pontuação é atribuída com base no número de operações necessárias para atingir o resultado correto da multiplicação. O jogo termina quando o tensor zero é atingido ou quando o número máximo de jogadas é feito. A fatoração final é avaliada com base em uma estimativa da classificação residual e certos critérios de otimização, como complexidade de tempo assintótica ou tempo de execução prático.
A posição inicial no TensorGame corresponde ao Tensor de Multiplicação da Matriz expresso de forma aleatória.
Em cada passo t do jogo, o jogador escreve três vetores
, que especifica os tensores de posto 1 . O estado do jogo é atualizado subtraindo os vetores selecionados pelo jogador:
onde
é o tensor de multiplicação de matrizes.Se o jogo terminar em p etapas, isso significa que o tensor de multiplicação da matriz
pode ser decomposto em tensores p rank-1 , ou seja, tem pelo menos posto p.O TensorGame pode então ser interpretado como um algoritmo de decomposição de classificação e o AlphaTensor pode ser visto como um algoritmo para estimar a classificação do tensor.
Arquitetura AlphaTensor
Até agora aprendemos sobre o TensorGame e esclarecemos como sua solução pode ser vista como um algoritmo de multiplicação de matrizes. Vamos agora explorar os principais conceitos do AlphaTensor, o algoritmo usado para o jogo.
A arquitetura AlphaTensor é basicamente uma arquitetura de transformador codificador-decodificador onde:
- o codificador toma como entrada o estado do jogo , as n ações anteriores executadas pelo modelo (geralmente n=7) e o índice de tempo t da ação atual. As informações são empilhadas em um tensor com forma (n+1, N^2, N^2, N^2). Este tensor é então remodelado e transformado (usando três camadas lineares) em um tensor de forma (N^2, N^2, c) onde c é a dimensão interna do modelo.
- o decodificador gera as ações n_steps a partir do vetor incorporado fornecido pelo codificador de forma auto-regressiva. Cada ação corresponde a um token dos trigêmeos representando um dos trigêmeos decompondo o tensor do jogo (ou seja, reduzindo sua classificação)
O modelo é treinado alternando a retropropagação e a atuação do modelo. A atuação do modelo é usada para gerar dados que são usados para treinar o modelo. Na prática, o modelo é treinado com uma mistura de dados gerados sinteticamente e dados gerados pelo modelo durante a atuação. A etapa de atuação é feita pegando um tensor 3D correspondente a uma operação de matriz e jogando jogos n_atores nele. Cada ator joga um jogo na base padrão ou em uma base alternativa (a mudança de base é aplicada com uma dada probabilidade). Os resultados são coletados e podem ser usados na etapa de treinamento com os dados sintéticos.
A etapa de atuação é baseada no Monte Carlo Tree Search (MCTS) do AlphaZero, modificado para suportar grandes espaços de ação. Resumindo, antes de escolher a ação, os caminhos de n_sims são explorados a partir da saída do modelo com uma exploração futura máxima de 5 etapas. As probabilidades geradas pelo modelo são então ajustadas levando em consideração os caminhos gerados. Em seguida, a ação com o(s) caminho(s) futuro(s) mais promissor(es) é escolhida(s) para continuar o jogo.
Ao treinar o modelo, a recompensa é, na verdade, uma recompensa negativa (penalidade). Seu valor absoluto aumenta a cada passo adicional necessário para resolver o jogo. Se o modelo leva m passos para resolver um TensorGame, a recompensa associada ao jogo é r=-m. Se o modelo não for capaz de resolver o TensorGame em passos de max_rank, a recompensa é calculada estimando a classificação do tensor restante. O posto é estimado como a soma dos postos das matrizes que compõem o tensor. A estimativa é um limite superior na classificação verdadeira do tensor.
Ao ajustar o modelo, a recompensa da penalidade no estado terminal também deve levar em consideração a latência do algoritmo produzido pelo modelo. A fórmula de recompensa torna-se rt'=rt+λbt, onde rt é o esquema de recompensa descrito anteriormente, bt é a recompensa de referência (diferente de zero apenas no estado terminal) e λ é um coeficiente especificado pelo usuário.
Aceleração (%) dos algoritmos descobertos pelo AlphaTensor adaptados para uma GPU e uma TPU, extraídos do artigo da DeepMind. Os aumentos de velocidade são medidos em relação à multiplicação da matriz padrão (por exemplo, cuBLAS para a GPU) no mesmo hardware e comparados com o Algoritmo de Strassen-quadrado. Fonte: DeepMind.
Eu lancei recentemente OpenAlpha Tensor, a primeira implementação de código aberto do AlphaTensor. Nesta seção, vou percorrer a implementação. Como discutimos anteriormente, a arquitetura do AlphaTensor é bastante direta, baseada em um transformador padrão com uma arquitetura de codificador-decodificador. Os componentes mais interessantes do AlphaTensor são a primeira camada na parte do codificador e a maneira como as ações são amostradas.
Vamos começar com a primeira camada de codificação.
# 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
No trecho acima, mostramos como o tensor de entrada é decomposto em três tensores, que são usados como entradas de consulta, chave e valor da camada do transformador.
- Nas três dimensões do tensor que representam as matrizes achatadas (A, B, C), o tensor de entrada é achatado ao longo de cada dimensão junto com a dimensão que representa as ações anteriores. Desta forma, em cada cópia achatada do tensor de entrada, a dimensão selecionada é uma agregação dos últimos valores T-1 e o valor real, para todos os valores S da dimensão selecionada, onde S=N^2. Filosoficamente, é como se, para cada dimensão, nos concentrássemos no que aconteceu nas ações anteriores daquela dimensão.
- Os escalares são mapeados em três espaços diferentes de dimensão S^2, e então remodelados para serem concatenados com os tensores obtidos no ponto anterior. Conceitualmente, os escalares são mapeados para um espaço de incorporação de dimensão S^2 e, em seguida, as informações incorporadas são divididas em S vetores e empilhadas juntas, semelhante ao que acontece com o texto quando tokenizado.
- Os tokens escalares são concatenados com o tensor de entrada reestruturado e, em seguida, fornecidos como entrada para uma camada linear para mapear as informações de foco escalares+histórico do canal na dimensão interna do modelo.
Essas três etapas podem ser interpretadas como uma forma de dar ao modelo tanto informações sobre os escalares (como no passo de tempo do TensorGame) quanto o foco nas ações anteriores de cada canal.
Em relação à forma como as ações são produzidas, é interessante notar que o AlphaTensor gera como saída o trio u, v, w, que visa reduzir o rank do tensor. Os três vetores têm tamanho S e como estão concatenados o modelo tem que produzir um vetor de tamanho 3*S. O AlphaTensor é treinado com um algoritmo RL, então todas as ações possíveis devem ser expressas em termos de probabilidades em um espaço enumerado, ou seja, o modelo produz uma probabilidade sobre as diferentes ações. Isso significa que cada vetor no espaço 3S deve ser mapeado para uma ação diferente. Isso resulta em um espaço de ação de tamanho |F|^(3S), onde |F| é o número de valores diferentes que o elemento de u, v, w pode assumir. Normalmente, os valores são restritos a (-2, -1, 0, 1, 2), resultando em uma cardinalidade de 5 elementos.
Aí vem um grande desafio: para gerar as probabilidades de ação para um produto matricial de matrizes de tamanho 5 precisaríamos de uma memória de 5^75 * 4 bytes, o que significaria ~10^44 GB de memória. Claramente, não podemos administrar um espaço de ação tão grande.
Como resolvemos o problema? Para reduzir o consumo de memória das probabilidades de ação, podemos dividir os trigêmeos em pedaços menores, “tokenizá-los” e tratar os pedaços como tokens gerados na arquitetura do transformador, ou seja, os tokens são fornecidos como entrada para o decodificador em um processo auto-regressivo caminho. No exemplo acima podemos dividir os tripletos em 15 chunks, reduzindo o consumo de memória para 15 * 5^(75/15) * 4, ou seja, 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), )
Acima, mostramos o trecho de código para gerar a ação completa. No código, self.core contém a camada do decodificador e o tensor e representa a saída da camada do codificador. O zero pode ser considerado como o token em modelos NLP e as ações n_steps que representam os chunks n_steps são gerados de forma progressiva.
O modelo retorna três quantidades:
- As ações geradas
- A probabilidade associada à ação completa
- Os logits produzidos para gerar a primeira ação (o primeiro pedaço) que será usado para calcular o valor do modelo.
Vale a pena gastar algumas palavras no parâmetro n_samples. O parâmetro é usado para a etapa de atuação e permite que o modelo gere diferentes versões dos trigêmeos que serão usados para explorar o espaço de ação no algoritmo Monte Carlo Tree Search usado no processo de Atuação. As diferentes ações n_samples são amostradas de acordo com a política gerada pelo modelo.
Etapa de atuação
A parte mais complicada de todo o algoritmo é provavelmente a etapa de atuação usada para resolver o TensorGame. O algoritmo não é explicado profundamente no artigo do AlphaTensor, pois é baseado em vários artigos anteriores do DeepMind que são apenas citados e dados como conhecidos. Aqui, vou reconstruir todas as peças que faltam e explicar passo a passo nossa implementação.
Podemos organizar as etapas de atuação em três componentes diferentes:
- A busca da árvore de Monte-Carlo
- A simulação do jogo
- O cálculo de política aprimorado
Vamos analisá-los um por um.
Pesquisa de árvore de Monte-Carlo (MCTS)
Monte Carlo Tree Search (MCTS) é uma técnica de inteligência artificial amplamente utilizada para jogos, particularmente em jogos de tabuleiro e videogames. O algoritmo cria uma árvore de jogo que simula movimentos e resultados potenciais e usa amostragem aleatória para avaliar a recompensa esperada para cada movimento. O algoritmo então seleciona iterativamente o movimento com a maior recompensa esperada e simula os resultados até atingir um estado terminal ou uma condição de parada especificada. As simulações são usadas para estimar a probabilidade de vitória para cada movimento e orientar o processo de tomada de decisão. O MCTS demonstrou ser eficaz em jogos complexos, onde o número de movimentos e resultados possíveis é grande, e tem sido usado em sistemas de IA de jogos bem-sucedidos, como o AlphaGo.
No AlphaTensor é usada uma versão modificada do MCTS original. Em particular, ao invés de selecionar aleatoriamente a ação de todo o espaço de ação, a ação é selecionada entre um subconjunto gerado diretamente pelo modelo (através dos n_samples apresentados anteriormente). A correção para a atualização da política é então aplicada na etapa de cálculo da política aprimorada.
Em nossa implementação, decidimos manter todas as informações sobre a árvore de Monte-Carlo em um dicionário tendo como chave a versão hash do estado do TensorGame e como valores as informações associadas ao próprio estado. Cada passo de Monte-Carlo parte de um nó e simula minijogos n_sim, explorando o futuro com um horizonte de 5 lances. Caso o nó já tenha sido explorado em simulações anteriores, n_sim é ajustado considerando o número de explorações anteriores. Para cada nó o número de visitas é armazenado no tensor N_s_a, pois este tensor contém o número de visitas por ação filho do nó (dentre as amostradas pelo modelo).
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
O código acima mostra nossa implementação do algoritmo. Para simplificar o código, a correção da política é realizada na função simula_game.
Simulação de Jogo
A função simula_game é responsável por explorar a árvore composta por nós que representam um determinado estado do TensorGame. Ele também executa o modelo sempre que um nó folha é encontrado e armazena todas as informações do nó no dicionário state_dict. Vamos dar uma olhada profunda em sua implementação:
@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)
Cada simulação é dividida em três partes:
- Seleção
- Expansão
- backup
Na parte de seleção, a simulação é executada nos nós de árvore já gerados e o seguinte nó é selecionado usando a seguinte função:
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()]
Na prática, a ação que maximiza a função ucb:
para o estado dado é selecionado. Aqui Q representa os valores Q gerados pelo modelo e π representa a distribuição aleatória sobre as ações amostradas usando a política do modelo. N(s, a) representa o número de visitas do nó à ação a do nó s.
Uma vez que a fase de seleção atinge um nó folha, se a simulação não atingiu uma condição terminal (em termos de exploração máxima, ou seja, horizonte futuro ou final do jogo), o modelo é então usado para selecionar n_samples nós alternativos (eles serão folha nós na iteração sucessiva). Isso é chamado de fase de expansão, pois novos nós são adicionados à árvore. Então, nenhum outro nó é explorado na simulação atual, mas a folha q_value é enviada para a seguinte etapa da simulação: o backup.
O backup é a etapa final de cada simulação. Durante o backup, se o nó folha estiver em estado terminal, a recompensa final é calculada; caso contrário, o valor q da folha é usado como uma recompensa estimada. Em seguida, a recompensa é retropropagada na trajetória de simulação, atualizando os estados q_values e atualizando o contador de visitas N(s, a). No snippet abaixo, mostramos o código para a retropropagação da recompensa.
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
Computação de política aprimorada
Uma vez que todas as simulações foram executadas e o MCTS oferece um instantâneo interessante do futuro próximo, é hora de atualizar a política associada aos nós previstos e devolvê-los, para que possam ser usados durante o treinamento. A política aprimorada, seguindo o método descrito em Hubert e outros, é usado para gerenciar grandes espaços de ação. De fato, para um espaço de busca pequeno, é possível durante o MCTS amostrar uma ação aleatoriamente do espaço de ação e avaliar seu impacto. Uma abordagem semelhante em um espaço de ação muito maior levaria a todas as trajetórias divergentes em caminhos diferentes e seria necessária uma quantidade infinita de trajetórias para obter estatísticas significativas e, em seguida, atualizar a política. Como aqui estamos usando sample-MCTS para evitar a dispersão, ou seja, as ações n_samples são amostradas de acordo com a política do modelo e, em seguida, o MCTS apenas seleciona uma das ações amostradas enquanto explora a árvore, precisamos levar em consideração a correção da amostra ao calcular a política atualizada final que será usada durante o treinamento do modelo.
Na prática, a política melhorada é calculada como
onde
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
Observe que em nossa implementação, depois de calcular a política do tensor N_s_a, precisamos mapeá-la de volta para o tensor de ação original. De fato, N_s_a considera apenas as ações amostradas pelo modelo, enquanto a política final deve conter probabilidades também para as ações não exploradas.
Diferenças em relação ao algoritmo de treinamento ChatGPT
AlphaTensor é o mais recente membro da família AlphaGo/AlphaZero de métodos de inteligência artificial da DeepMind. Esses métodos são baseados no algoritmo Monte Carlo Tree Search (MCTS), que foi refinado e aprimorado pela DeepMind para lidar com tarefas cada vez mais complexas. Outro sistema de IA, o ChatGPT da OpenAI, que tem causado muita agitação por seu desempenho notável, foi treinado com uma abordagem diferente, chamada Aprendizado por Reforço com Feedback Humano (RLHF).
RLHF é uma técnica de ajuste fino usada para ajustar modelos de linguagem para seguir um conjunto de instruções escritas. Ele usa as preferências humanas como um sinal de recompensa para ajustar o modelo, alinhando assim o comportamento do modelo de linguagem com as preferências declaradas de um grupo específico de pessoas, em vez de uma noção mais ampla de 'valores humanos'.
Em contraste, o MCTS é um algoritmo de busca baseado em árvore usado para determinar os movimentos ideais nos jogos. Ele simula movimentos potenciais e atualiza os valores de cada movimento com base em seus resultados, orientando a seleção do melhor movimento.
O RLHF coleta dados de demonstrações escritas por humanos e comparações rotuladas por humanos entre modelos de IA e treina um modelo de recompensa para prever as preferências de um determinado grupo de pessoas. O modelo de recompensa é então usado para ajustar os modelos de IA. O MCTS, por outro lado, usa simulações e avaliações para determinar a melhor decisão.
Embora sejam abordagens diferentes, RLHF e MCTS também apresentam semelhanças. Ambas as técnicas de inteligência artificial usam métodos de tomada de decisão e resolução de problemas, e ambas usam uma abordagem de tentativa e erro para explorar diferentes opções e tomar decisões com base nas informações disponíveis. Ambos também são processos iterativos que melhoram com o tempo à medida que mais informações e experiências são reunidas.
A escolha entre RLHF e MCTS depende da tarefa em questão. O RLHF é ideal quando não há uma métrica clara para avaliar o desempenho do modelo, enquanto o MCTS provou ser eficaz em tarefas semelhantes a jogos em que o conhecimento e a exploração do futuro dão ao modelo uma vantagem significativa.
Otimização de código para treinamento AlphaTensor
A implementação do algoritmo de treinamento AlphaTensor requer encontrar o compromisso perfeito entre velocidade de treinamento e consumo de memória. Como visto na seção Modelo, simplesmente considerar a tokenização de ação pode economizar muita memória, mas uma redução excessivamente agressiva do espaço de ação pode levar à queda na precisão e desempenho mais lento. Este último ocorre porque todos os tokens são gerados sequencialmente de forma autorregressiva pelo decodificador do modelo. Portanto, o tempo de inferência cresce linearmente com o número de tokens por ação, uma vez que o softmax no espaço de ação não é mais o gargalo.
Ao configurar o treinamento do AlphaTensor, as principais dificuldades encontradas foram em lidar com o processo de atuação. Se os tensores não forem armazenados no formato correto, o MCTS pode facilmente causar um crescimento descontrolado do uso da memória. Por outro lado, se o número de tensores armazenados durante cada simulação for reduzido demais, o MCTS pode gastar uma quantidade infinita de tempo recalculando os estados necessários.
Vejamos um exemplo da etapa de simulação do jogo, onde o jogo é explorado observando possíveis cenários futuros. Para cada estado, se não salvarmos as ações geradas pelo modelo e decidirmos salvar apenas a semente aleatória usada para amostrar as ações da política, então, cada vez que explorarmos um nó da árvore, teremos que recalcular a política e em seguida, experimente as ações. Claramente, decidimos armazenar as ações amostradas para economizar tempo e evitar ter que gerenciar o compartilhamento de modelos entre diferentes processos no caso da paralelização da exploração MCTS. Porém, apenas salvar as ações não foi suficiente para obter uma ação de atuação suficientemente eficiente. Na verdade, o tempo para converter as ações n_steps no trio (u, v, w), reduzir o estado do tensor do jogo e criar os novos tensores 3D a partir das ações n_samples seria facilmente um gargalo para todo o treinamento. Em segundo lugar, não queríamos armazenar todos os estados futuros possíveis para cada ação amostrada, pois isso teria um enorme impacto na memória utilizada pelo algoritmo. Suponha que definimos n_samples=32, n=7 e N=5, e vamos lembrar que N é o tamanho do produto da matriz quadrada que queremos reduzir en é o número de ações anteriores lembradas pelo modelo. Nesta situação, cada tensor de estado teria a forma (8, 25, 25, 25), que multiplicado por 32 resultaria em 3282525254 bytes para cada nó no gráfico. Agora, considerando que cada simulação na fase de expansão gera um novo nó (e n_sim=200), teríamos um consumo final de memória de 200328252525*4 = 3.2 GB apenas para o primeiro nó MCTS. No pior cenário, ao explorar os nós max_rank atuantes (onde max_rank=150), isso resultaria em um consumo total de memória de 150 * 3.2 GB = 480 GB na memória RAM (ou memória GPU se todos os tensores fossem armazenados na GPU) . Executamos o treinamento em nossa estação de trabalho com 128 GB de RAM e 48 GB de memória GPU, então tivemos que reduzir o consumo de memória.
Como não queríamos aumentar o tempo de execução, adotamos uma otimização que explora a redundância nos tensores de estado produzidos. Na verdade, os tensores têm em comum n-1 ações anteriores, que podem então ser armazenadas uma vez e não repetidas para cada tensor armazenado. Isto resulta numa redução de memória de 2/7~28%, o que significa que no pior dos casos podem ser armazenados 137GB. Neste ponto, simplesmente podando a parte não utilizada da árvore (como as trajetórias não selecionadas) e armazenando os tensores na memória da CPU, conseguimos evitar qualquer erro de memória durante o treinamento.
Com o OpenAlphaTensor agora sendo de código aberto, vários caminhos empolgantes para desenvolvimento adicional se abrem.
Uma progressão natural é o ajuste fino do OpenAlphaTensor nos dispositivos de hardware de destino. Espera-se que isso leve a um desempenho computacional muito competitivo. Vou publicar mais sobre o desempenho do OpenAlphaTensor em vários hardwares em GitHub. No momento da redação deste artigo, o OpenAlphaTensor estava em treinamento.
Outro avanço importante seria o suporte à compilação remota, permitindo que os usuários construíssem algoritmos otimizados para dispositivos de ponta. Isso pode ser obtido armazenando o modelo OpenAlphaTensor em um servidor, enquanto o algoritmo de multiplicação de matrizes é avaliado em hardware diferente.
Também pode ser importante estender o suporte para diferentes compiladores para calcular a correção de recompensa baseada em latência. Diferentes compiladores podem levar a diferentes algoritmos otimizados em um determinado hardware. Por exemplo, o artigo da DeepMind mostrou resultados promissores usando JAX e o compilador XLA em GPUs TPU e Nvidia. Seria interessante avaliar isso usando NCCL em Nvidia ou LLVM em CPUs.
Por fim, estender o modelo e o algoritmo de treinamento para suportar tamanhos de matriz maiores continua sendo um grande desafio em aberto. Atualmente, o OpenAlphaTensor suporta um tamanho máximo de matriz de 5, mas pode ser aplicado dividindo multiplicações de matrizes maiores em grupos de minúsculos MMs com um tamanho menor que 5. Essa abordagem é abaixo do ideal e realiza a redução diretamente no grande tensor correspondente ao MM completo poderia, teoricamente, levar a melhores resultados.
Diego Fiori é o CTO da Nebuly AI, uma empresa comprometida em tornar a otimização de IA parte do kit de ferramentas de cada desenvolvedor.
- Conteúdo com tecnologia de SEO e distribuição de relações públicas. Seja amplificado hoje.
- Platoblockchain. Inteligência Metaverso Web3. Conhecimento Ampliado. Acesse aqui.
- Fonte: 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
- :é
- ][p
- $UP
- 1
- 3d
- 8
- a
- Capaz
- Sobre
- acima
- absoluto
- acelerando
- Segundo
- conformemente
- Conta
- precisão
- Alcançar
- alcançado
- Açao Social
- ações
- adicionado
- Adicional
- Ajustado
- adotado
- avançar
- Vantagem
- Depois de
- Agente
- agregação
- agressivo
- AI
- Sistemas de IA
- visa
- algoritmo
- algoritmos
- Todos os Produtos
- Permitindo
- permite
- sozinho
- já
- alternativa
- entre
- quantidade
- analisar
- e
- Outro
- aplicado
- abordagem
- se aproxima
- arquitetura
- SOMOS
- artigo
- artificial
- inteligência artificial
- AS
- atribuído
- associado
- At
- Automatizado
- disponível
- evitando
- em caminho duplo
- backup
- baseado
- Basicamente
- base
- BE
- Porque
- torna-se
- antes
- ser
- abaixo
- referência
- MELHOR
- Melhor
- entre
- Pós
- borda
- Board Games
- obrigado
- mais amplo
- BT
- construir
- construído
- by
- chamado
- CAN
- não podes
- casas
- Causar
- causado
- certo
- desafiar
- desafiante
- alterar
- Canal
- ChatGPT
- criança
- escolha
- escolha
- escolhido
- citado
- remover filtragem
- claramente
- código
- coleta
- combinar
- comprometido
- comum
- Empresa
- comparado
- competitivo
- integrações
- complexidade
- componentes
- composta
- compromisso
- computação
- poder computacional
- Computar
- computação
- conceitos
- Conceitualmente
- condição
- confiança
- Considerar
- considerado
- considerando
- considera
- consumo
- contém
- continuar
- contraste
- convertido
- núcleo
- Correspondente
- corresponde
- poderia
- Contador
- CPU
- cria
- Criar
- critérios
- CTO
- Atual
- Atualmente
- dados,
- lidar
- decidir
- decidido
- decisão
- Tomada de Decisão
- decisões
- profundo
- deep learning
- DeepMind
- depende
- descrito
- Determinar
- Developer
- Desenvolvimento
- dispositivo
- Dispositivos/Instrumentos
- DICT
- diferente
- dificuldades
- Dimensão
- dimensões
- diretamente
- descobrir
- descobrindo
- discutido
- distribuição
- dividido
- down
- Cair
- durante
- e
- cada
- Mais cedo
- facilmente
- borda
- Eficaz
- eficiente
- ou
- elemento
- elementos
- incorporado
- termina
- aprimorada
- enorme
- suficiente
- entrada
- erro
- estimativa
- estimado
- Éter (ETH)
- avaliar
- avaliadas
- avaliação
- avaliações
- Mesmo
- Cada
- exemplo
- exemplos
- emocionante
- execução
- expansão
- esperado
- vasta experiência
- Explicação
- explicado
- façanhas
- exploração
- explorar
- Explorado
- Explorando
- expressa
- estender
- estendendo
- extremamente
- bastante
- família
- mais rápido
- retornos
- poucos
- Figura
- final
- descoberta
- Primeiro nome
- Flutuador
- Foco
- seguir
- seguinte
- Pegada
- Escolha
- formulário
- formato
- Fórmula
- encontrado
- da
- cheio
- função
- fundamental
- mais distante
- desenvolvimento adicional
- futuro
- jogo
- Games
- gerar
- gerado
- gera
- gerando
- ter
- obtendo
- OFERTE
- dado
- Dando
- meta
- vai
- Bom estado, com sinais de uso
- GPU
- GPUs
- gráfico
- ótimo
- Grupo
- Do grupo
- Cresce
- Growth
- guia
- mão
- Manipulação
- aconteceu
- acontece
- Hardware
- dispositivo de hardware
- dispositivos de hardware
- Ter
- ter
- SUA PARTICIPAÇÃO FAZ A DIFERENÇA
- mais
- horizonte
- Como funciona o dobrador de carta de canal
- Como Negociar
- Contudo
- HTTPS
- enorme
- humano
- i
- EU VOU
- ideal
- IDX
- imagem
- Impacto
- executar
- implementação
- importante
- melhorar
- melhorado
- in
- Crescimento
- Aumenta
- cada vez mais
- de treinadores em Entrevista Motivacional
- índice
- INFORMAÇÕES
- do estado inicial,
- entrada
- em vez disso
- instruções
- Inteligência
- interessante
- interno
- introduzido
- intuição
- IT
- iteração
- ESTÁ
- se
- jpg
- KDnuggetsGenericName
- Guarda
- Chave
- Conhecimento
- conhecido
- língua
- grande
- Maior
- Sobrenome
- Latência
- mais recente
- camada
- camadas
- conduzir
- aprendido
- aprendizagem
- Lista
- olhar
- procurando
- lote
- moldadas
- a Principal
- principal
- fazer
- Fazendo
- gerencia
- gestão
- muitos
- mapa,
- mapeamento
- Matriz
- Importância
- máximo
- significado
- significativo
- significa
- membro
- Memória
- método
- métodos
- métrico
- desaparecido
- mistura
- modelo
- modelos
- modificada
- módulo
- mais
- mais eficiente
- Além disso
- a maioria
- mover
- movimentos
- multiplicado
- natural
- Natureza
- Perto
- necessário
- você merece...
- necessário
- negativo
- redes
- Neural
- redes neurais
- Novo
- Próximo
- PNL
- nó
- nós
- não especialistas
- Noção
- número
- Nvidia
- obtido
- of
- Oferece
- on
- ONE
- aberto
- open source
- OpenAI
- operação
- Operações
- ideal
- otimização
- Otimize
- otimizado
- Opções
- original
- Outros
- de outra forma
- saída
- global
- Papel
- papéis
- parâmetro
- parte
- particular
- particularmente
- peças
- Pessoas
- perfeita
- atuação
- realização
- fase
- peças
- platão
- Inteligência de Dados Platão
- PlatãoData
- Jogar
- jogador
- jogar
- ponto
- políticas
- Privacidade
- posição
- possível
- potencial
- poder
- Prática
- prática
- predizer
- previsto
- preferências
- apresentado
- anterior
- probabilidade
- provavelmente
- Problema
- processo
- processos
- produzir
- Produzido
- Produto
- progressão
- progressivo
- promissor
- proposto
- provavelmente
- comprovado
- publicar
- publicado
- RAM
- acaso
- fileiras
- em vez
- alcançado
- Chega
- recentemente
- reduzir
- Reduzido
- redução
- refinado
- aprendizagem de reforço
- liberado
- remanescente
- permanece
- notável
- lembrar
- remoto
- repetido
- representando
- representa
- requeridos
- exige
- responsável
- restringido
- resultar
- resultando
- Resultados
- retorno
- Retorna
- revolucionar
- Recompensa
- LINHA
- rt
- Execute
- s
- mesmo
- Salvar
- poupança
- cenário
- cenários
- esquema
- Pesquisar
- Segundo
- Seção
- semente
- selecionado
- selecionando
- doadores,
- AUTO
- conjunto
- contexto
- vários
- Shape
- compartilhando
- Baixo
- rede de apoio social
- mostrar
- mostrando
- Shows
- Signal
- periodo
- de forma considerável
- semelhante
- semelhanças
- simples
- simplicidade
- simplesmente
- simulação
- desde
- situação
- Tamanho
- tamanhos
- pequeno
- menor
- Instantâneo
- So
- solução
- RESOLVER
- Resolvendo
- alguns
- fonte
- Espaço
- espaços
- específico
- especificamente
- especificada
- velocidade
- gastar
- Passar
- divisão
- quadrado
- empilhado
- Etapa
- padrão
- começo
- começa
- Estado
- estado-da-arte
- estabelecido
- Unidos
- estatística
- Passo
- Passos
- paragem
- loja
- armazenadas
- lojas
- franco
- bem sucedido
- tal
- ajuda
- suportes
- sintético
- dados sintéticos
- sinteticamente
- .
- sistemas
- adaptados
- Tire
- toma
- tomar
- Target
- Tarefa
- tarefas
- técnicas
- terminal
- condições
- que
- A
- O Futuro
- The Graph
- as informações
- O Estado
- deles
- Eles
- assim
- assim sendo
- Este
- Terceiro
- três
- tridimensional
- Através da
- tempo
- demorado
- para
- juntos
- token
- tokenization
- tokenized
- Tokens
- também
- kit de ferramentas
- topo
- tocha
- Total
- tradicional
- Trem
- treinado
- Training
- trens
- trajetória
- transformado
- tratar
- verdadeiro
- compreender
- Universo
- não usado
- Atualizar
- Atualizada
- Atualizações
- atualização
- atualização
- us
- Uso
- usar
- usuários
- geralmente
- valor
- Valores
- vário
- versão
- Vídeo
- jogos de vídeo
- Visite a
- Visitas
- W
- Caminho..
- O Quê
- qual
- enquanto
- largamente
- Wikipedia
- precisarão
- vitória
- com
- palavras
- estação de trabalho
- Equivalente há
- seria
- escrita
- escrito
- X
- zefirnet
- zero