Contents

nn.Embedding

Lookup table은 자료구조로 자주 사용되기도 하지만 그 자체로 학습에 사용되어야 하는 경우도 있다. 대표적인 것이 언어모델에서 token을 만드는 과정이다. 자주 사용되는 만큼 당연히 torch에서도 이를 지원하며 이는 nn.Embedding을 사용하면 된다.

nn.Embedding

torch에서 nn.Embedding은 Sparse Layers항목에 있다. 추정컨데 여기서 sparse는 sparse matrix의 의미로 사용되었으리라 생각한다. Sparse matrix는 이름에서 나타내듯 대부분의 값이 0과 같은 값으로 채워져있고 값이 채워져있는 위치가 드물 때 해당 index에 어떤 값이 있는지를 표현하는 방식으로 0을 잔뜩 채워넣어 표기하는 것보다 효율적인 방법이다. Embedding도 이에 딱 맞지는 않지만 index를 통해 직접 조회하므로 유사한 점이 있고 값을 회수하는 것도 효율적이라는 공통점을 갖는다.

nn.Embedding은 lookup table로 주어진 index에 대해 해당하는 벡터를 return해주는 함수로 생각해 볼 수 있다. nn.Embeddingnum_embeddingsembedding_dim을 필수적으로 받으며 각각은 embedding할 대상의 크기와 embedding의 차원을 의미한다. 언어모델에서 생각해보면 num_embeddings에는 단어 전체의 크기인 vocab_size가 들어가고 embedding_dim은 해당 단어에 해당하는 벡터가 된다. 재미있는 점은 nn.Embeddingweight 자체가 바로 lookup table의 value라는 점이다.

1
>>> embedding_layer = nn.Embedding(num_embeddings=8, embedding_dim=4)

이 때 embedding_layer의 weight은 다음으로 조회할 수 있다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> embedding_layer.weight
Parameter containing:
tensor([[-3.3307e-01,  3.6703e-01,  1.0469e+00,  4.7747e-01],
        [ 1.2580e+00, -8.4635e-01, -4.5423e-01,  7.3989e-01],
        [-9.6321e-01, -5.6505e-01,  1.4820e+00, -1.4327e+00],
        [ 1.3749e+00, -1.1486e+00,  1.7615e+00, -1.1193e+00],
        [ 5.0136e-01,  3.7170e-02, -1.4641e+00,  5.2615e-01],
        [-1.3963e+00,  7.1644e-02,  1.0268e+00,  2.0797e+00],
        [ 1.2618e+00,  9.1458e-06, -3.5095e-02, -7.5799e-01],
        [-9.7169e-02, -1.2832e-01,  8.4959e-01, -1.2375e+00]],
       requires_grad=True)

Index로 query를 하면 해당 index의 값을 그대로 얻게 된다.

1
2
3
4
5
>>> idx = torch.IntTensor([1, 3, 4])
>>> embedding_layer(idx)
tensor([[ 1.2580, -0.8463, -0.4542,  0.7399],
        [ 1.3749, -1.1486,  1.7615, -1.1193],
        [ 0.5014,  0.0372, -1.4641,  0.5262]], grad_fn=<EmbeddingBackward>)

비교해보면 확인할 수 있듯, index에 대한 query결과는 weight과 같다.

위의 형태로도 대부분의 경우 embedding은 충분히 사용이 가능하다. nn.Embedding에서 padding_idx라는 parameter가 있는데 이는 말 그대로 padding인 index를 알려주어 해당 index에 대해서는 학습을 하지 않도록 할 수 있다. padding_idx로 지정되면 weight은 0으로 고정이 된다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> embedding_layer = nn.Embedding(num_embeddings=8, embedding_dim=4, padding_idx=2, max_norm=1, norm_type=2)
>>> embedding_layer.weight
Parameter containing:
tensor([[-0.2921,  0.3198, -1.4463, -0.0211],
        [ 1.6262, -1.3835, -0.6565, -0.7844],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [-0.2434,  0.1026,  0.3376,  2.4716],
        [-1.1623, -0.2063,  0.7496,  0.3677],
        [-1.6133, -1.6416,  0.2389, -0.4451],
        [-1.8385, -0.1169,  0.0202,  0.8692],
        [ 1.0441, -0.5111, -0.8074, -0.1289]], requires_grad=True)

또한 max_normnorm_typenn.Embedding의 크기를 조절하는데 사용할 수 있다. 대부분 딥러닝 기반 학습에서 스케일을 일정하게 유지하는 것은 학습에 중요한 요소이며 embedding을 사용하는 경우 이를 이용해 크기를 일정수준 이하로 제한할 수 있다. Norm의 type은 기본적으로 2-norm이며 norm_type을 사용해 바꿀 수 있다. max_norm으로 지정된 크기로 normalize를 해주어 embedding의 스케일을 지정한 크기 이하로 유지할 수 있다. 특이한 점은 query를 하는 시점에 대해 normalization을 수행한다는 점이다. 아래 예시를 보자.

1
2
3
>>> embedding_layer.weight.norm(dim=-1)
tensor([1.5099, 2.3675, 0.0000, 2.5085, 1.4459, 2.3564, 2.0371, 1.4212],
       grad_fn=<CopyBackwards>)

분명 위에서 max_norm=1로 설정하였음에도 norm의 크기가 1을 넘는다. 이 때 query를 수행해보자.

1
2
3
4
>>> embedding_layer(idx)
tensor([[ 0.6869, -0.5844, -0.2773, -0.3313],
        [-0.0970,  0.0409,  0.1346,  0.9853],
        [-0.8039, -0.1427,  0.5184,  0.2543]], grad_fn=<EmbeddingBackward>)

위와 비교해보면 idx[1, 3, 4]였음에도 weight과 다른 값이 나온 것을 볼 수 있다. 이에 대한 norm을 확인해보자.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
>>> embedding_layer.weight
Parameter containing:
tensor([[-0.2921,  0.3198, -1.4463, -0.0211],
        [ 0.6869, -0.5844, -0.2773, -0.3313],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [-0.0970,  0.0409,  0.1346,  0.9853],
        [-0.8039, -0.1427,  0.5184,  0.2543],
        [-1.6133, -1.6416,  0.2389, -0.4451],
        [-1.8385, -0.1169,  0.0202,  0.8692],
        [ 1.0441, -0.5111, -0.8074, -0.1289]], requires_grad=True)

>>> embedding_layer.weight.norm(dim=-1)
tensor([1.5099, 1.0000, 0.0000, 1.0000, 1.0000, 2.3564, 2.0371, 1.4212],

조회한 [1, 3, 4]에 대해서는 noramlization이 된 것을 볼 수 있다. 즉 nn.Embedding의 normalization은 조회가 될 경우 그 때 그 때 적용된다.

여기에 사용된 실제 실행 코드는 https://github.com/jihoonerd/pytorch-in-action/blob/main/embedding/embedding.ipynb에서 확인할 수 있다.

Reference