tf.keras.layers.ReversibleEmbedding

An embedding layer which can project backwards to the input dim.

Inherits From: Embedding, Layer, Operation

tf.keras.layers.ReversibleEmbedding(
    input_dim,
    output_dim,
    tie_weights=True,
    embeddings_initializer='uniform',
    embeddings_regularizer=None,
    embeddings_constraint=None,
    mask_zero=False,
    reverse_dtype=None,
    logit_soft_cap=None,
    **kwargs
)

This layer is an extension of keras.layers.Embedding for language models. This layer can be called "in reverse" with reverse=True, in which case the layer will linearly project from output_dim back to input_dim.

By default, the reverse projection will use the transpose of the embeddings weights to project to input_dim (weights are "tied"). If tie_weights=False, the model will use a separate, trainable variable for reverse projection.

This layer has no bias terms.

input_dim Integer. Size of the vocabulary, i.e. maximum integer index + 1.
output_dim Integer. Dimension of the dense embedding.
tie_weights Boolean, whether or not the matrix for embedding and the matrix for the reverse projection should share the same weights.
embeddings_initializer Initializer for the embeddings matrix (see keras.initializers).
embeddings_regularizer Regularizer function applied to the embeddings matrix (see keras.regularizers).
embeddings_constraint Constraint function applied to the embeddings matrix (see keras.constraints).
mask_zero Boolean, whether or not the input value 0 is a special "padding" value that should be masked out.
reverse_dtype The dtype for the reverse projection computation. Defaults to the compute_dtype of the layer.
logit_soft_cap If logit_soft_cap is set and reverse=True, the output logits will be scaled by tanh(logits / logit_soft_cap) * logit_soft_cap. This narrows the range of output logits and can improve training.
**kwargs other keyword arguments passed to keras.layers.Embedding, including name, trainable, dtype etc.

inputs The tensor inputs to the layer.
reverse Boolean. If True the layer will perform a linear projection from output_dim to input_dim, instead of a normal embedding call. Default to False.

Example:

batch_size = 16
vocab_size = 100
hidden_dim = 32
seq_length = 50

# Generate random inputs.
token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))

embedding = keras.layers.ReversibleEmbedding(vocab_size, hidden_dim)
# Embed tokens to shape `(batch_size, seq_length, hidden_dim)`.
hidden_states = embedding(token_ids)
# Project hidden states to shape `(batch_size, seq_length, vocab_size)`.
logits = embedding(hidden_states, reverse=True)

References:

embeddings

input Retrieves the input tensor(s) of a symbolic operation.

Only returns the tensor(s) corresponding to the first time the operation was called.

output Retrieves the output tensor(s) of a layer.

Only returns the tensor(s) corresponding to the first time the operation was called.

variable_serialization_spec Returns a dict mapping quantization modes to variable names in order.

This spec is used by save_own_variables and load_own_variables to determine the correct ordering of variables during serialization for each quantization mode. None means no quantization.

Methods

enable_lora

View source

enable_lora(
    rank,
    lora_alpha=None,
    a_initializer='he_uniform',
    b_initializer='zeros'
)

from_config

View source

@classmethod
from_config(
    config
)

Creates an operation from its config.

This method is the reverse of get_config, capable of instantiating the same operation from the config dictionary.

if "dtype" in config and isinstance(config["dtype"], dict):
    policy = dtype_policies.deserialize(config["dtype"])

Args
config A Python dictionary, typically the output of get_config.

Returns
An operation instance.