Pytorch, singkirkan loop for saat menambahkan permutasi satu vektor ke entri matriks?

Question

Saya mencoba menerapkan makalah ini, dan terjebak dengan langkah sederhana ini. Meskipun ini berkaitan dengan perhatian, hal yang membuat saya terjebak hanyalah bagaimana menerapkan permutasi vektor yang ditambahkan ke matriks tanpa menggunakan for loop.

Skor perhatian telah ditambahkan vektor bias yang dipelajari, teorinya adalah ia mengkodekan posisi relatif (ji) dari dua token yang diwakili oleh skor tersebut.

jadi alpha adalah matriks T x T, T bergantung pada batch yang diteruskan, dan B adalah vektor bias yang dipelajari yang panjangnya harus tetap dan sebesar 2T. Implementasi saya saat ini yang saya yakini sesuai dengan saran makalah ini adalah:

    def __init__(...):
       ...
        self.bias = torch.nn.Parameter(torch.randn(config.n),requires_grad = True)
        stdv = 1. / math.sqrt(self.bias.data.size(0))
        self.bias.data.uniform_(-stdv, stdv)
     def forward(..)
        ...
        #n = 201  (2* max_seq_len + 1)

        B_matrix = torch.zeros(self.T, self.T) # 60 x 60
        for i in range(self.T):
          B_matrix[i] = self.bias[torch.arange(start=n//2-i, end=n//2-i+T)])]

        attention_scores = attention_scores + B_matrix.unsqueeze(0)
        # 64 x 60 x 60   
        ...

Ini satu-satunya bagian yang relevan

B_matrix = torch.zeros(self.T, self.T) # 60 x 60
        for i in range(self.T):
          B_matrix[i] = self.bias[torch.arange(start=n//2-i, end=n//2-i+T)])]

pada dasarnya mencoba untuk tidak menggunakan loop for untuk melewati setiap baris.

tapi saya tahu ini pasti sangat tidak efisien, dan mahal bila model ini sangat besar. Saya melakukan perulangan eksplisit untuk setiap baris untuk mendapatkan permutasi dari vektor bias yang dipelajari.

Adakah yang bisa membantu saya dengan cara yang lebih baik, mungkin melalui penyiaran cerdas?

Setelah memikirkannya, saya tidak perlu membuat instance matriks nol, tetapi masih tidak bisa menghilangkan loop for? dan tidak dapat menggunakan gathering karena B_matrix berukuran berbeda dari vektor b ubin.

functor = lambda i : bias[torch.arange(start=n//2-i, end=n//2-i+T)]
B_matrix = torch.stack([functor(i) for i in torch.arange(T)])

Answer 1

Saya tidak tahu apa nyang seharusnya ada dalam kode Anda, tetapi saya pikir contoh berikut torch.meshgridmenyediakan apa yang Anda cari.

Seandainya

n, m = 10, 20   # arbitrary
a = torch.randn(n, m)
b = torch.randn(n + m)

kemudian

for i in range(n):
    for j in range(m):
        a[i, j] = a[i, j] + b[n - i + j]

setara dengan

ii, jj = torch.meshgrid(torch.arange(n), torch.arange(m))
a = a + b[n - ii + jj]

meskipun yang terakhir adalah operasi yang tidak pada tempatnya, yang biasanya merupakan hal yang baik. Jika Anda benar-benar menginginkan operasi di tempat, gantilah a =dengan a[...] =.

Perhatikan bahwa ini adalah contoh pengindeksan array integer di mana kita mengindeks bmenggunakan tensor yang bentuknya sama dengan a.