디코딩 API

TensorFlow.org에서 보기 Google Colab에서 실행 GitHub에서 보기 노트북 다운로드

개요

최근에는 자동 회귀 모델을 사용한 언어 생성에 대한 연구가 많이 진행되었습니다. 자동 퇴행 언어 생성에 시간 단계 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 선택 \(id_t = argmax_{w}P(id | id_{1:t-1})\) 각 시간 단계에서의 \(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)