디코딩 API

개요

최근에는 자동 회귀 모델을 사용한 언어 생성에 대한 연구가 많이 진행되었습니다. 자동 퇴행 언어 생성에 시간 단계 K에서 토큰의 확률 분포는 단계 K-1까지 모델의 토큰 예측에 따라 달라집니다. 이 모델 들어, 빔 검색, 욕심, 최고-P, 및 최고-K 등의 디코딩 전략 모델의 중요한 구성 요소는 크게 주어진 시간 단계 K에서 토큰 생성 된 출력의 스타일 / 성격에 영향을 미친다.

예를 들어, 빔 검색 할 때마다 단계에서 가설의 가능성이 가장 높은 num_beams을 유지하고 결국 전반적인 가장 높은 확률을 가지고 가설을 선택하여없는 숨겨진 확률이 매우 높은 토큰의 위험을 줄일 수 있습니다. Murray et al. (2018)양 등. (2018) 이 빔 검색 기계 번역 작업에서 잘 작동 보여줍니다. 빔 검색욕심 전략 모두 반복 토큰을 생성하는 가능성이있다.

팬 등 (2018) 등은 대부분 토큰 필터링 된 K와 확률 질량 만 K 토큰 재분배되는 상위 K 샘플링을 도입했다.

아리 홀츠만 외 (2019) 등은 확률 P 개까지 추가 누적 확률 토큰 작은 가능한 세트로부터 선택 상위 P 샘플링을 도입했다. 그런 다음 확률 질량은 이 집합 간에 재분배됩니다. 이러한 방식으로 토큰 세트의 크기는 동적으로 증가하거나 감소할 수 있습니다. 최고-P, 최고-k는 일반적으로 층 세대 등의 작업에 사용됩니다.

Decoding API는 자동 회귀 모델에서 다양한 디코딩 전략을 실험할 수 있는 인터페이스를 제공합니다.

  1. 다음 샘플링 전략은 기본 Decoding 클래스에서 상속되는 sampling_module.py에 제공됩니다.

  2. 빔 검색은 beam_search.py에서 제공됩니다. 깃허브

설정

pip install -q -U tensorflow-text
pip install -q tf-models-nightly
import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf

from official import nlp
from official.nlp.modeling.ops import sampling_module
from official.nlp.modeling.ops import beam_search
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/pkg_resources/__init__.py:119: PkgResourcesDeprecationWarning: 0.18ubuntu0.18.04.1 is an invalid version and will not be supported in a future release
  PkgResourcesDeprecationWarning,

TF-NLP에서 샘플링 모듈을 초기화합니다.

  • symbols_to_logits_fn : 사용이 폐쇄가의 logits을 예측하는 모델을 호출하는 index+1 단계. 이 클로저의 입력 및 출력은 다음과 같습니다.
Args:
 
1] ids : Current decoded sequences. int tensor with shape (batch_size, index + 1 or 1 if padded_decode is True)],
 
2] index [scalar] : current decoded step,
 
3] cache [nested dictionary of tensors] : Only used for faster decoding to store pre-computed attention hidden states for keys and values. More explanation in the cell below.
Returns:
 
1] tensor for next-step logits [batch_size, vocab]
 
2] the updated_cache [nested dictionary of tensors].

캐시는 더 빠른 디코딩에 사용됩니다. 여기서 A는 기준 상기 폐쇄 구현.

  • length_normalization_fn : 사용 길이 정규화 매개 변수를 반환이 폐쇄.
Args: 
 
1] length : scalar for decoded step index.
 
2] dtype : data-type of output tensor
Returns:
 
1] value of length normalization factor.
  • vocab_size : 출력 어휘 크기입니다.

  • max_decode_length : 스칼라 복호 단계의 총 수에 대한.

  • eos_id : 일괄 처리의 모든 출력 디코딩 식별자이 eos_id이있는 경우 디코딩이 중지됩니다.

  • padded_decode : TPU에서 실행되는 경우는 true로 설정합니다. 이것이 True이면 텐서는 max_decoding_length로 채워집니다.

  • top_k :이 값이 경우 top_k이> (1)가 활성화됩니다.

  • top_p :이 값은> 0 <1.0의 경우는 활성화 top_p

  • sampling_temperature이 다시 추정치 softmax를 출력하기 위해 사용된다. 온도는 확률이 높은 토큰 쪽으로 분포를 왜곡하고 꼬리 분포의 질량을 낮춥니다. 값은 양수여야 합니다. 저온은 탐욕에 해당하며 분포를 더 날카롭게 만들고 고온은 더 평평하게 만듭니다.

  • enable_greedy : 기본적으로이 사실이고 욕심 디코딩을 사용할 수 있습니다. 다른 전략을 실험하려면 이것을 False로 설정하십시오.

모델 하이퍼파라미터 초기화

params = {}
params['num_heads'] = 2
params['num_layers'] = 2
params['batch_size'] = 2
params['n_dims'] = 256
params['max_decode_length'] = 4

변압기를 기반으로 같은 자동 회귀 아키텍처에서 인코더 - 디코더 모델, 캐시 빠른 연속 디코딩에 사용됩니다. 모든 계층에 대해 미리 계산된 숨겨진 상태(self-attention 블록 및 cross-attention 블록의 키 및 값)를 저장하는 중첩 사전입니다.

캐시를 초기화합니다.

cache = {
   
'layer_%d' % layer: {
       
'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32),
       
'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32)
       
} for layer in range(params['num_layers'])
   
}
print("cache key shape for layer 1 :", cache['layer_1']['k'].shape)
cache key shape for layer 1 : (2, 4, 2, 128)

필요한 경우 길이 정규화를 위한 클로저를 정의합니다.

생성된 시퀀스의 최종 점수를 정규화하는 데 사용되며 선택 사항입니다.

def length_norm(length, dtype):
 
"""Return length normalization factor."""
 
return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)

model_fn 만들기

실제로, 이것은 같은 실제 모델 구현에 의해 대체됩니다 여기

Args:
i
: Step that is being decoded.
Returns:
  logit probabilities of size
[batch_size, 1, vocab_size]
probabilities = tf.constant([[[0.3, 0.4, 0.3], [0.3, 0.3, 0.4],
                             
[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                           
[[0.2, 0.5, 0.3], [0.2, 0.7, 0.1],
                             
[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]])
def model_fn(i):
 
return probabilities[:, i, :]

symbol_to_logits_fn 초기화

def _symbols_to_logits_fn():
 
"""Calculates logits of the next tokens."""
 
def symbols_to_logits_fn(ids, i, temp_cache):
   
del ids
    logits
= tf.cast(tf.math.log(model_fn(i)), tf.float32)
   
return logits, temp_cache
 
return symbols_to_logits_fn

탐욕스러운

: 욕심쟁이 디코딩 최고 차기 ID와 확률 토큰 ID 선택 idt=argmaxwP(id|id1:t1) 각 시간 단계에서의 t. 다음 스케치는 욕심 많은 디코딩을 보여줍니다.

greedy_obj = sampling_module.SamplingModule(
    length_normalization_fn
=None,
    dtype
=tf.float32,
    symbols_to_logits_fn
=_symbols_to_logits_fn(),
    vocab_size
=3,
    max_decode_length
=params['max_decode_length'],
    eos_id
=10,
    padded_decode
=False)
ids
, _ = greedy_obj.generate(
    initial_ids
=tf.constant([9, 1]), initial_cache=cache)
print("Greedy Decoded Ids:", ids)
Greedy Decoded Ids: tf.Tensor(
[[9 1 2 2 2]
 [1 1 1 2 2]], shape=(2, 5), dtype=int32)

top_k 샘플링

최고-K 샘플링에서 K는 대부분 다음 토큰 ID를 필터링하고 확률 질량은 만 K 식별자 사이에 재분배된다.

top_k_obj = sampling_module.SamplingModule(
    length_normalization_fn
=length_norm,
    dtype
=tf.float32,
    symbols_to_logits_fn
=_symbols_to_logits_fn(),
    vocab_size
=3,
    max_decode_length
=params['max_decode_length'],
    eos_id
=10,
    sample_temperature
=tf.constant(1.0),
    top_k
=tf.constant(3),
    padded_decode
=False,
    enable_greedy
=False)
ids
, _ = top_k_obj.generate(
    initial_ids
=tf.constant([9, 1]), initial_cache=cache)
print("top-k sampled Ids:", ids)
top-k sampled Ids: tf.Tensor(
[[9 1 0 2 2]
 [1 0 1 2 2]], shape=(2, 5), dtype=int32)

top_p 샘플링

대신 샘플링의 가장 가능성 K 토큰 식별자, 누적 확률 확률 (P)을 초과 식별자의 작은 가능한 세트로부터 선택한다면 샘플링 톱 페이지이다.

top_p_obj = sampling_module.SamplingModule(
    length_normalization_fn
=length_norm,
    dtype
=tf.float32,
    symbols_to_logits_fn
=_symbols_to_logits_fn(),
    vocab_size
=3,
    max_decode_length
=params['max_decode_length'],
    eos_id
=10,
    sample_temperature
=tf.constant(1.0),
    top_p
=tf.constant(0.9),
    padded_decode
=False,
    enable_greedy
=False)
ids
, _ = top_p_obj.generate(
    initial_ids
=tf.constant([9, 1]), initial_cache=cache)
print("top-p sampled Ids:", ids)
top-p sampled Ids: tf.Tensor(
[[9 1 1 2 2]
 [1 1 1 0 2]], shape=(2, 5), dtype=int32)

빔 검색 디코딩

빔 검색은 각 시간 단계에서 가장 가능성이 높은 가설 num_beams를 유지하고 결국 전체 확률이 가장 높은 가설을 선택하여 숨겨진 고확률 토큰 ID를 누락할 위험을 줄입니다.

beam_size = 2
params['batch_size'] = 1
beam_cache
= {
   
'layer_%d' % layer: {
       
'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32),
       
'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32)
       
} for layer in range(params['num_layers'])
   
}
print("cache key shape for layer 1 :", beam_cache['layer_1']['k'].shape)
ids
, _ = beam_search.sequence_beam_search(
    symbols_to_logits_fn
=_symbols_to_logits_fn(),
    initial_ids
=tf.constant([9], tf.int32),
    initial_cache
=beam_cache,
    vocab_size
=3,
    beam_size
=beam_size,
    alpha
=0.6,
    max_decode_length
=params['max_decode_length'],
    eos_id
=10,
    padded_decode
=False,
    dtype
=tf.float32)
print("Beam search ids:", ids)
cache key shape for layer 1 : (1, 4, 2, 256)
Beam search ids: tf.Tensor(
[[[9 0 1 2 2]
  [9 1 2 2 2]]], shape=(1, 2, 5), dtype=int32)