ALBERT encoder: factorised embedding + cross-layer sharing (Geron Ch 15, hmalbt)

The embedding goes V -> E -> H instead of V -> H, and ONE block object is applied n_layers times, so depth costs no extra parameters – which is why n_params is independent of n_layers while n_params_unshared is not. Token ids are 0-based.

Usage

morie_geron_albert(
  X,
  n_layers = 4,
  n_heads = 2,
  d_model = 8,
  d_embed = 4,
  vocab_size = NULL,
  d_ff = NULL,
  seed = 0
)

Arguments

X

0-based token ids, vector or (batch, T) matrix.

n_layers, n_heads, d_model, d_ff

Encoder geometry.

d_embed

Factorised width E in [1, d_model].

vocab_size

Optional; defaults to max(X) + 1.

seed

LCG seed.

Value

List with hidden (array batch x T x d), n_params, n_params_unshared, block_params, embedding_params, embedding_params_direct, shared, d_embed.

Examples

V <- c(1, 2, 3, 4, 5, 6, 7, 8)
morie_geron_albert(V)
#> $hidden
#> , , 1
#> 
#>            [,1]      [,2]      [,3]      [,4]      [,5]     [,6]      [,7]
#> [1,] -0.4857786 0.5282667 -1.249282 0.7538802 -0.257951 -1.24892 0.5186818
#>          [,8]
#> [1,] 1.194362
#> 
#> , , 2
#> 
#>         [,1]     [,2]     [,3]       [,4]       [,5]     [,6]     [,7]
#> [1,] 0.95222 1.889332 1.054112 -0.8664737 -0.7195389 1.166835 1.673234
#>           [,8]
#> [1,] 0.3074001
#> 
#> , , 3
#> 
#>            [,1]       [,2]      [,3]      [,4]      [,5]       [,6]      [,7]
#> [1,] -0.6062372 -0.7641875 0.6946412 0.6764557 0.4903354 -0.6956003 0.1426254
#>           [,8]
#> [1,] 0.4324754
#> 
#> , , 4
#> 
#>          [,1]       [,2]       [,3]     [,4]       [,5]     [,6]     [,7]
#> [1,] 1.212983 -0.9697249 -0.7872076 -1.17903 -0.6053583 0.128578 1.250526
#>          [,8]
#> [1,] 1.398281
#> 
#> , , 5
#> 
#>            [,1]      [,2]     [,3]     [,4]      [,5]      [,6]      [,7]
#> [1,] -0.7381107 0.1369142 1.049472 1.910249 -0.403652 -1.033546 -1.065925
#>            [,8]
#> [1,] -0.8073637
#> 
#> , , 6
#> 
#>           [,1]       [,2]      [,3]       [,4]       [,5]     [,6]       [,7]
#> [1,] -1.078382 -0.6395976 -1.398077 -0.5776292 -0.3926123 1.784118 -0.8007508
#>           [,8]
#> [1,] -1.032377
#> 
#> , , 7
#> 
#>          [,1]      [,2]       [,3]      [,4]       [,5]      [,6]       [,7]
#> [1,] 1.603399 -1.135183 -0.3773036 0.1858058 -0.5898771 0.3503777 -0.6433397
#>           [,8]
#> [1,] -1.613819
#> 
#> , , 8
#> 
#>           [,1]      [,2]     [,3]       [,4]     [,5]       [,6]      [,7]
#> [1,] -0.860094 0.9541803 1.013645 -0.9032569 2.478654 -0.4518425 -1.075052
#>          [,8]
#> [1,] 0.121041
#> 
#> 
#> $n_params
#> [1] 668
#> 
#> $n_params_unshared
#> [1] 2276
#> 
#> $block_params
#> [1] 536
#> 
#> $embedding_params
#> [1] 68
#> 
#> $embedding_params_direct
#> [1] 72
#> 
#> $shared
#> [1] TRUE
#> 
#> $d_embed
#> [1] 4
#> 
#> $estimate
#> [1] 668
#> 
#> $n
#> [1] 8
#> 
#> $method
#> [1] "ALBERT: factorised embedding plus one cross-layer-shared encoder block"
#>