GEMM o Coração Matemático

A curiosa leitora precisa saber que a multiplicação de matrizes está no coração dos Transformers, mas a expressão multiplicar matrizes esconde quase toda a dificuldade prática. A definição matemática da multiplicação de matrizes cabe em uma linha. Isso não é suficiente. Qualquer implementação capaz de aproveitar uma CPU, ou GPU moderna, exige compreender hierarquias de memória, ordem dos acessos, blocagem, vetorização, paralelismo, precisão numérica e o custo de transportar cada byte. Cada byte conta.

Em álgebra linear de alto desempenho, uma multiplicação economizada pode ajudar. Um acesso evitado à memória pode ajudar ainda mais.

Nosso objeto de estudo será a operação GEMM, de General Matrix-Matrix Multiplication. A atenta leitora verá como a mesma equação atravessa diferentes camadas da máquina: memória principal, caches, registradores vetoriais, memória compartilhada da GPU e unidades matriciais especializadas. Ao final, poderemos olhar para um Transformer e reconhecer quais operações são limitadas por computação e quais são limitadas pelo movimento de dados. Isso, vai fazer toda a diferença nos seus projetos futuros. Uma vantagem competitiva em tempos de agentes de inteligência artificial.

1. Da multiplicação de matrizes à GEMM

Sejam $A$, $B$ e $C$ três matrizes tal que:

$$A\in {\mathbb{R}}^{M\times K},B\in {\mathbb{R}}^{K\times N},C\in {\mathbb{R}}^{M\times N}.$$

Se $C$ for a matriz produto, a multiplicação de matrizes clássica, das matrizes $A$ e $B$, define cada elemento da matriz $C$ como:

$$C_{ij}=\sum _{p=0}^{K-1}{A}_{ip}{B}_{pj}.$$

Vale a pena desmontar essa fórmula antes de generalizá-la.

O elemento $C_{ij}$ é o produto escalar entre a linha $i$ de $A$ e a coluna $j$ de $B$: percorremos as duas em paralelo, multiplicamos os pares de valores e somamos os $K$ produtos. Daí a exigência sobre as dimensões internas: a linha de $A$ tem $K$ elementos e a coluna de $B$ também, e sem essa coincidência o emparelhamento não existe. As dimensões externas, $M$ e $N$, apenas contam quantas linhas e quantas colunas há para combinar, e é delas que $C$ herda o formato $M\times N$.

Um exemplo pequeno deixa o mecanismo visível. Com $M=2$, $K=3$ e $N=2$:

O elemento $C_{00}$ combina a linha $0$ de $A$ com a coluna $0$ de $B$:

$$C_{00}=1\cdot 4+2\cdot 2+0\cdot 1=8.$$

Repetindo o procedimento para as quatro posições:

Duas propriedades dessa definição sustentarão tudo o que veremos adiante, e ambas se provam em uma linha, olhando para um elemento genérico. Primeiro, multiplicar por um escalar comuta com o produto:

$${[(\alpha A)B]}_{ij}=\sum _p(\alpha A_{ip})B_{pj}=\alpha \sum _p{A}_{ip}{B}_{pj}={[\alpha (AB)]}_{ij}.$$

Segundo, a soma de matrizes de mesmo formato é definida elemento a elemento, $[X+Y{]}_{ij}=X_{ij}+Y_{ij}$; portanto, somar ao produto $AB$ qualquer outra matriz $M\times N$ é uma operação bem definida e barata, que não envolve nenhuma multiplicação de matrizes adicional. Guarde essas duas observações: elas são a licença algébrica para tudo o que a GEMM acrescenta.

Falta apresentar o personagem que usaremos o artigo inteiro: a transposta. A transposta de uma matriz $X\in {\mathbb{R}}^{m\times n}$, denotada $X^{\mathrm{\top }}$, é a matriz em ${\mathbb{R}}^{n\times m}$ obtida trocando o papel de linhas e colunas:

$$(X^{\mathrm{\top }}{)}_{ij}=X_{ji}.$$

No nosso exemplo, $B\in {\mathbb{R}}^{3\times 2}$ tem transposta

Nenhum valor muda; muda apenas qual índice conta linhas e qual conta colunas. A transposta aparecerá com frequência nos Transformers, pois os scores de atenção da Seção 13 são $Q{K}^{\mathrm{\top }}$, e aparece nas bibliotecas por um motivo adicional, mundano: dependendo de como uma matriz está disposta na memória, é mais barato lê-la como se estivesse transposta do que fabricar fisicamente uma cópia transposta.

Com esses ingredientes, podemos enunciar a operação que as bibliotecas BLAS (Basic Linear Algebra Subprograms, Subprogramas Básicos de Álgebra Linear) de fato implementam: