TensorFlow中LSTM单元状态与权重的默认初始化及手动设置咨询

Great question! Let’s break this down clearly—LSTM initialization plays a huge role in training stability, so it’s smart to dig into the details.

Default Initialization for LSTMs in TensorFlow #

Weights #

You’re totally right about tf.get_variable() defaulting to glorot_uniform_initializer (aka Xavier uniform initialization)! TensorFlow’s LSTM implementations—whether you’re using Keras’ LSTM layer or the native LSTMCell—rely on tf.get_variable() to create all trainable weights (including the weight matrices for the input gate, forget gate, output gate, and cell state).

So all trainable LSTM weights use Glorot uniform initialization by default. This method scales weights based on the input and output dimensions, which helps mitigate gradient vanishing/exploding issues early in training.

Hidden & Cell States #

As for the LSTM’s internal states (hidden state h and cell state c), the default initialization is all-zero tensors:

• In Keras, if you don’t specify the initial_state parameter when calling an LSTM layer, it automatically generates all-zero initial states shaped (batch_size, units) at the start of each batch.
• For native tf.nn.rnn_cell.LSTMCell, calling the cell without passing initial_state will trigger cell.zero_state() to create all-zero states.

Easy Ways to Manually Set Initializers #

For Weights #

Both Keras and native TensorFlow let you override weight initializers with straightforward parameters—no complicated hacks needed:

Keras Example

import tensorflow as tf
from tensorflow.keras.layers import LSTM
from tensorflow.keras.initializers import HeNormal, Constant

# Customize initializers for an LSTM layer
lstm_layer = LSTM(
    units=64,
    kernel_initializer=HeNormal(),  # Use He normal initialization for input weights
    recurrent_initializer='orthogonal',  # Orthogonal init for recurrent weights (string or instance works)
    bias_initializer=Constant(value=0.1)  # Initialize all biases to 0.1
)

Native TensorFlow Example

from tensorflow.nn.rnn_cell import LSTMCell

lstm_cell = LSTMCell(
    num_units=64,
    kernel_initializer=tf.initializers.glorot_normal(),  # Glorot normal init for input weights
    recurrent_initializer=tf.initializers.orthogonal()  # Orthogonal init for recurrent weights
)

For Initial States #

If you want to move beyond all-zero initial states, here are simple methods—including making initial states trainable:

Non-Trainable Initial State (Keras)

batch_size = 32
hidden_units = 64

# Manually create random initial states
initial_h = tf.random.normal(shape=(batch_size, hidden_units))
initial_c = tf.random.normal(shape=(batch_size, hidden_units))

# Pass states to the LSTM layer
output, final_state = lstm_layer(inputs, initial_state=[initial_h, initial_c])

Trainable Initial State (Keras)

For scenarios where you want initial states to adapt during training (great for small datasets or specific use cases):

# Define trainable initial states (start with single-sample shape, then expand to batch)
trainable_initial_h = tf.Variable(tf.random.normal(shape=(1, hidden_units)), trainable=True)
trainable_initial_c = tf.Variable(tf.random.normal(shape=(1, hidden_units)), trainable=True)

# Expand to match the current batch size
initial_h_batch = tf.tile(trainable_initial_h, [batch_size, 1])
initial_c_batch = tf.tile(trainable_initial_c, [batch_size, 1])

# Use the trainable states in the LSTM layer
output, final_state = lstm_layer(inputs, initial_state=[initial_h_batch, initial_c_batch])

Native TensorFlow Example

from tensorflow.nn.rnn_cell import LSTMCell, LSTMStateTuple

batch_size = 32
hidden_units = 64

lstm_cell = LSTMCell(num_units=hidden_units)

# Manually create initial states wrapped in LSTMStateTuple
initial_state = LSTMStateTuple(
    c=tf.random.normal(shape=(batch_size, hidden_units)),
    h=tf.random.normal(shape=(batch_size, hidden_units))
)

# Run the RNN with custom initial state
outputs, final_state = tf.nn.dynamic_rnn(lstm_cell, inputs, initial_state=initial_state)

内容的提问来源于stack exchange，提问作者mrk