#

රේඛීයපක්ෂග්රාහී (අලිබී) සමඟ අවධානය යොමු කරන්න

මෙයක්රියාත්මක කිරීමකි රේඛීය අගතීන් සමඟ අවධානය (අලිබී) කඩදාසි වලින් දුම්රිය කෙටි, දිගු පරීක්ෂණය: රේඛීය අගතීන් සමඟ අවධානය යොමු කිරීම ආදාන දිග නිස්සාරණය සක්රීය කරයි.

මෙයස්ථානීය කේතීකරණ වෙනුවට අවධානය යොමු කිරීමේ ලකුණු වලට එකතු කරන ලද අගතීන් සමඟ ප්රතිස්ථාපනය කරයි (අවධානය යොමු කිරීම්, සොෆ්ට්මැක්ස් වලට පෙර). මෙය ස්වයංක්රීය ප්රතිගාමී කාර්යයන් මත පරීක්ෂා කරන ලද සාපේක්ෂ යෝජනා ක්රමයක් වන අතර සමීප ටෝකන සඳහා නැඹුරුව ඉහළ අගයක් ගන්නා අතර දුර බැහැර ටෝකන සඳහා අඩු වේ. අගතීන් ලොග් පරිමාණයේ රේඛීයව අඩු වේ (එය සොෆ්ට්මැක්ස් වලට පෙර නිසා) සහ එක් එක් හිස වෙනස් බෑවුමක් ඇත.

මෙන්න $i$ -th ටෝකනය සඳහා අවධානය යොමු කිරීමේ සූත්රය,

a_{i} = soft m ax (q_{i} K^{⊤} + m \cdot [- (i - 1), \dots, 1, 0]) = soft m ax (q_{i} K^{⊤} + m \cdot [0, 1, \dots, (i - 1)])

$i$ -th ටෝකනයේ විමසුම $q_{i} \in R^{d}$ $K \in R^{i \times d}$ කොහේද, යතුරු දක්වා ඇත $i$ , සහ $d$ හිසකට විශේෂාංග ගණන. ඉහත සමානාත්මතාවය නතර $soft m ax$ වන බව සලකන්න පරිවර්තන වලට වෙනස් නොවේ (ප්රති result ලය වෙනස් නොකර ඔබට සියලු අංග සඳහා නියත එකතු කළ හැකිය).

අලිබීආකෘතියක් සඳහා පුහුණු කේතය මෙන්න.

35import math
36from typing import Optional
37
38import torch
39from torch import nn
40
41from labml.logger import inspect
42from labml_nn.transformers.mha import MultiHeadAttention

#

එක්එක් හිස $m$ සඳහා හිස විශේෂිත බෑවුම ලබා ගන්න

n_heads යනු අවධානය යොමු කිරීමේ ස්ථරයේ හිස් ගණන $n$

පළමුහිස සඳහා බෑවුම

$\frac{1}{2 ^{\frac{8}{n}}} = 2^{- \frac{8}{n}}$

සෙසුහිස් සඳහා බෑවුම් ඉහත සමාන අනුපාතයක් සහිත ජ්යාමිතික ශ්රේණියක් වේ.

උදාහරණයක්ලෙස හිස් ගණන බෑවුම් $8$ වන විට $\frac{1}{2 ^{1}}, \frac{1}{2 ^{2}}, \dots, \frac{1}{2 ^{8}}$

45def get_slopes(n_heads: int):

#

2ට ආසන්නතම බලය ලබා ගන්න n_heads . 2 බලයක් n_heads නොවේ නම්, අපි මුලින්ම 2 හි සමීපතම (කුඩා) බලයට බෑවුම් ගණනය කර ඉතිරි බෑවුම් එකතු කරමු.

64    n = 2 ** math.floor(math.log2(n_heads))

#

$2^{- \frac{8}{n}}$

66    m_0 = 2.0 ** (-8.0 / n)

#

$2^{- 1 \frac{8}{n}}, 2^{- 2 \frac{8}{n}}, 2^{- 3 \frac{8}{n}}, \dots$

68    m = torch.pow(m_0, torch.arange(1, 1 + n))

#

2බලයක් n_heads නොවේ නම්, අපි ඉතිරි බෑවුම් එකතු කරමු. අපි ඉතිරි බෑවුම් ගණනය කරමු $n * 2$ (කලින් එකතු කරන ලද බෑවුම් මග හැරීම). බෑවුම් ඉහළට ගන්න n_heads .

73    if n < n_heads:

#

$2^{- \frac{8}{2 n}}$

75        m_hat_0 = 2.0 ** (-4.0 / n)

#

$2^{- 1 \frac{8}{2 n}}, 2^{- 3 \frac{8}{2 n}}, 2^{- 5 \frac{8}{2 n}}, \dots$ කලින් එකතු කරන ලද බෑවුම් වළක්වා $2$ ගැනීමට අපි පියවර ගන්නා බව සලකන්න.

78        m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2))

#

ඉතිරිබෑවුම් සහිත බෑවුම් සහිත බෑවුම් සංයුක්ත කරන්න.

80        m = torch.cat([m, m_hat])
81
82    return m

#

අවධානයනැඹුරුව න්යාසය ගණනය

n_heads යනු අවධානය යොමු කිරීමේ ස්ථරයේ හිස් ගණන
mask හැඩයේ අවධානය යොමු කිරීමේ ආවරණයයි [seq_len_q, seq_len_k]

මෙමAliBi අවධානය අගතීන් [seq_len_q, seq_len_k, n_heads, ] සමග හැඩය අනුකෘතියක් නැවත.

85@torch.no_grad()
86def get_alibi_biases(n_heads: int, mask: torch.Tensor):

#

එක්එක් හිස $m$ සඳහා බෑවුම් ලබා ගන්න

97    m = get_slopes(n_heads).to(mask.device)

#

දුරගණනය කරන්න $[0, 1, \dots, N]$ මෙන්න අපි ආවරණ භාවිතා කරන දුර ගණනය කරමු.

එයහේතු වෙස් මුහුණක් බැවින් $[0, 1, \dots, N]$ අපටද භාවිතා කළ හැකිය. distance = torch.arange(mask.shape[1], dtype=torch.long, device=mask.device)[None, :]

104    distance = mask.cumsum(dim=-1)

#

අලිබීනැඹුරුව න්යාසය ලබා ගැනීම සඳහා ඒවා යුගල වශයෙන් ගුණ කරන්න

107    return distance[:, :, None] * m[None, None, :]

#

රේඛීයපක්ෂග්රාහී (අලිබී) සමඟ අවධානය යොමු කරන්න

අපි බහු-ප්රධාන අවධානයඅභිබවා යමු.

110class AlibiMultiHeadAttention(MultiHeadAttention):

#

117    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
118        super().__init__(heads, d_model, dropout_prob)

#

අලිබීඅගතීන් හැඹිලි කිරීමට

121        self.alibi_biases = None

#

query , key සහ value විමසුම, යතුරසහ අගයදෛශික එකතු කිරීම ගබඩා කරන ආතන්ය වේ. ඒවායේ හැඩය ඇත [seq_len, batch_size, d_model] .

mask හැඩය ඇති [seq_len, seq_len, batch_size] අතර කණ්ඩායම සඳහා b , ස්ථානයේ විමසුමට ප්රවේශය i තිබේද යන්න mask[i, j, b] දක්වයි ස්ථානයේ ප්රධාන-අගය j .

123    def forward(self, *,
124                query: torch.Tensor,
125                key: torch.Tensor,
126                value: torch.Tensor,
127                mask: Optional[torch.Tensor] = None):

#

අලිබීක්රියා කරන්නේ හේතු වෙස් මුහුණු සමඟ පමණි.

139        assert mask is not None
140        assert mask.shape[0] == mask.shape[1] and mask.shape[2] == 1

#

query , key value සහ හැඩය [seq_len, batch_size, d_model]

143        seq_len, batch_size, _ = query.shape

#

ආවරණසඳහා හිස මානයක් එක් කර එහි හැඩය පරීක්ෂා කරන්න.

146        mask = self.prepare_mask(mask, query.shape, key.shape)

#

සූදානම්වන්න query , key සහ අවධානය ගණනය කිරීම value සඳහා. මේවාට පසුව හැඩය ඇත [seq_len, batch_size, heads, d_k] .

150        query = self.query(query)
151        key = self.key(key)
152        value = self.value(value)

#

අවධානයලකුණු ගණනය කරන්න $Q K^{⊤}$ . මෙය හැඩයේ ආතතිකයක් ලබා දෙයි [seq_len, seq_len, batch_size, heads] .

156        scores = self.get_scores(query, key)

#

පරිමාණලකුණු $\frac{Q K ^{⊤}}{d _{k}}$

159        scores *= self.scale

#

එයහැඹිලි නොවේ නම් AliBi අගතීන් නිර්මාණය

162        if self.alibi_biases is None or self.alibi_biases.shape[1] < seq_len:

#

mask seq_len, seq_len, 1, 1 හැඩය ඇත

164            self.alibi_biases = get_alibi_biases(scores.shape[-1], mask[:, :, 0, 0])

#

අවධානයලකුණු AliBi අගතීන් එකතු කරන්න. අලිබී අගතීන් හැඩයක් scores ඇති [seq_len, seq_len, n_heads] අතර හැඩයක්

ඇත [seq_len, seq_len, batch_size, n_heads]

169        scores += self.alibi_biases[:seq_len, :seq_len, None, :]

#

වෙස්යොදන්න

172        scores = scores.masked_fill(mask == 0, float('-inf'))

#

$so f t m a x$ ප්රධාන අනුක්රමය මානයක් ඔස්සේ අවධානය $se q so f t m a x (\frac{Q K ^{⊤}}{d _{k}})$

176        attn = self.softmax(scores)

#

අතහැරදැමීම යොදන්න

179        attn = self.dropout(attn)

#

අගයන්අනුව ගුණ කරන්න $se q so f t m a x (\frac{Q K ^{⊤}}{d _{k}}) V$

183        x = torch.einsum("ijbh,jbhd->ibhd", attn, value)

#

බහුහිස් සංයුක්ත කරන්න

186        x = x.reshape(seq_len, batch_size, -1)

#

ප්රතිදානස්ථරය

189        return self.output(x)

#

බෑවුම්දැකීමට සරල පරීක්ෂණ කාර්යය.

192def _test_alibi():

#

196    inspect(get_slopes(12).tolist(), _n=-1)
197    from labml_nn.transformers.utils import subsequent_mask
198
199    mask = subsequent_mask(8)[:, :, 0]
200    inspect(mask)
201
202    inspect(get_alibi_biases(12, mask)[:, :, 3], _n=-1)

#

206if __name__ == '__main__':
207    _test_alibi()

රේඛීයපක්ෂග්රාහී (අලිබී) සමඟ අවධානය යොමු කරන්න

එක්එක් හිස m සඳහා හිස විශේෂිත බෑවුම ලබා ගන්න

අවධානයනැඹුරුව න්යාසය ගණනය

රේඛීයපක්ෂග්රාහී (අලිබී) සමඟ අවධානය යොමු කරන්න

එක්එක් හිස $m$ සඳහා හිස විශේෂිත බෑවුම ලබා ගන්න