TensorFlow
`RNN`
neural networks
machine learning
deep learning

How to feed back `RNN` output to input in tensorflow

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Feeding an RNN output back into its next input is a standard sequence-generation pattern rather than a special TensorFlow trick. The core idea is simple: run one step, take the predicted token or value, transform it into the next-step input, and carry the recurrent state forward at the same time.

Understand What Is Being Fed Back

An RNN step usually produces two related things:

  • an output for the current time step
  • a hidden state that carries memory to the next time step

The hidden state is already fed forward internally by the recurrent cell. What people usually mean by "feed the output back to the input" is that the model's emitted prediction becomes the external input for the next time step.

That pattern is common in:

  • character generation
  • sequence-to-sequence decoding
  • autoregressive time-series forecasting
  • sampling from language models

So the design has two feedback paths:

  • internal state transition handled by the cell
  • external next-input construction handled by your loop or decoder logic

Build the Step Function Explicitly

In TensorFlow, the clearest way to express this is with an RNN cell and an explicit decoding loop.

python
1import tensorflow as tf
2
3vocab_size = 5
4hidden_size = 8
5embedding = tf.keras.layers.Embedding(vocab_size, 4)
6cell = tf.keras.layers.GRUCell(hidden_size)
7projection = tf.keras.layers.Dense(vocab_size)
8
9batch_size = 1
10state = cell.get_initial_state(batch_size=batch_size, dtype=tf.float32)
11current_token = tf.constant([0])
12
13for _ in range(4):
14    x = embedding(current_token)
15    output, state = cell(x, states=state)
16    logits = projection(output)
17    current_token = tf.argmax(logits, axis=-1, output_type=tf.int32)
18    print(current_token.numpy())

Here the predicted token from one iteration becomes the input token for the next iteration. That is the feedback loop.

Use Teacher Forcing During Training

During training, you often do not want the model to consume its own predictions at every step. A common training strategy is teacher forcing, where the true next token is provided as input instead.

That means the model learns from the correct historical context even before it can generate well on its own.

A simplified example looks like this:

python
1import tensorflow as tf
2
3inputs = tf.constant([[0, 1, 2, 3]])
4embedding = tf.keras.layers.Embedding(10, 4)
5rnn = tf.keras.layers.GRU(8, return_sequences=True)
6dense = tf.keras.layers.Dense(10)
7
8x = embedding(inputs)
9sequence_outputs = rnn(x)
10logits = dense(sequence_outputs)
11print(logits.shape)

This is not autoregressive feedback yet. It is the training-side structure where the full input sequence is known.

Separate Training and Inference Logic

This distinction matters a lot:

  • training often uses teacher forcing
  • inference often uses the model's previous prediction as the next input

If you blur those two modes together, debugging becomes difficult. In practice, many models use one graph or function for training and a separate decoder loop for generation.

That is not a flaw in TensorFlow. It is the normal shape of sequence models.

Use the State, Not Only the Output Token

A common misunderstanding is to think that feeding the output token back is enough by itself. It is not. The recurrent state must also be passed from one step to the next.

That is why manual decoding loops usually keep both:

  • the current predicted token or generated value
  • the updated recurrent state returned by the cell

If you drop the state and only recycle the token, the model stops behaving like a recurrent sequence generator.

Common Pitfalls

The most common mistake is confusing the hidden state with the predicted output. TensorFlow recurrent cells already manage state flow if you pass the returned state forward. The extra work is building the next input from the emitted prediction.

Another issue is training the model with teacher forcing and then being surprised when free-running inference behaves differently. That gap is normal and should be accounted for in evaluation.

Developers also often use argmax in generation loops without thinking about whether greedy decoding is actually the intended sampling strategy. Greedy feedback is simple, but not always the best generation policy.

Summary

  • Recurrent state is already passed forward by the RNN cell; output feedback is the extra external loop.
  • Feeding output back means using the model's prediction as the next-step input.
  • Training often uses teacher forcing, while inference usually uses autoregressive decoding.
  • A manual loop with an RNN cell makes the feedback pattern easiest to see.
  • Keep both the recurrent state and the generated input token flowing through the decoding loop.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.