Chapter 10 · Content revision 3

Separate input encoding from the transformer

Compose text lookup, the existing causal backbone and a tied text head. Follow vector shapes, live parameter checks and the two gradient paths into one shared table.

The problem: text IDs hide the vector boundary

The text model currently accepts token IDs and returns vocabulary logits. That convenient call contains three different transformations: look up vectors, transform those vectors with the causal decoder, then score the text vocabulary. The attention blocks operate on residual features. They need the vectors and their causal layout; token IDs and next-token targets serve other parts of the computation.

Making this boundary explicit lets us use the existing transformer as a component without making its input contract depend on a text vocabulary. The complete text model must still produce the same values and send both uses of its tied table back to one learned parameter.

Follow one text sequence through the three transformations

The existing BPE tokenizer with no merges encodes ab as document IDs [0,99,100,1]. The first and last IDs mark the document boundaries. Use [0,99,100] as the input and [99,100,1] as the next-token targets.

The fixture has one batch row, three supplied positions and four residual features per position. Its text table has $258$ rows, one for each vocabulary entry, and four features per row. Looking up the three input IDs produces $X$ with shape $[1,3,4]$ . The existing two decoder blocks and final RMSNorm produce hidden features $H$ with that same shape. The text head multiplies $H$ by the transpose of the lookup table, producing logits $Z$ with shape $[1,3,258]$ : one score for every vocabulary entry at every supplied position.

With initialization seed 39, the mean text loss is $5.668398842431571$ nats. Its f64 bits are 4016ac70bf03b620. The explicit composition and the ordinary text wrapper give identical logit bits in this fixture. The table remains one live node, and the complete model still has $1372$ learned scalars in $20$ named tensors, counting that table once.

Tokenize the fixed text, reuse its learned lookup, run the existing backbone, then apply the head that borrows the same table. rust/demos/practical-ch10-input-encoding-boundary/src/lib.rs#compose-text-path

/// All learned tensors belong to the existing model. The views borrow them.
pub fn run() -> Result<Value, Box<dyn std::error::Error>> {
    let tokenizer = BpeTokenizer::from_merge_pairs(&[])?;
    let document = tokenizer.encode_utf8_document("ab");
    let ids = &document[..3];
    let targets = document[1..]
        .iter()
        .map(|&id| id as usize)
        .collect::<Vec<_>>();
    let model = DecoderModel::new(
        DecoderModelConfig::new(258, 4, 1, 8, 2, 8, 10000.0, 1e-6),
        &mut SplitMix64::from_seed(39),
    )?;
    let context = AutogradContext::recording();
    let backbone = model.backbone();
    let encoded = TextInputEncoder::tokens(&model, ids, &[1, 3]).encode(context, &backbone)?;
    let embedding_shape = encoded.embeddings().shape().to_vec();
    let embedding_bits = bits(encoded.embeddings());
    let hidden = backbone.forward(encoded)?;
    let head = model.text_output_head();
    let logits = head.forward(context, hidden.hidden())?;
    let loss = logits.indexed_mean_nll(2, &targets)?;
    loss.backward()?;
    let total = model
        .tied_embedding()
        .tensor()
        .gradient_snapshot()
        .ok_or("missing tied gradient")?;
    let wrapper = model.forward(ids, &[1, 3])?;
    let same_text_logits = bits(&logits) == bits(wrapper.logits());
    let tied_alias = head
        .table()
        .is_same_node(model.embedding().table().tensor());

Separate lookup, hidden features and vocabulary scores

Let $B$ count batch rows, $T$ supplied positions per row, $D$ residual features per position and $V$ text vocabulary entries. The learned table $E$ has shape $[V,D]$ . Token ID $i_{b,t}$ selects its row:

X_{b,t,d}=E_{i_{b,t},d}.

Here $b$ selects a batch row, $t$ a position and $d$ a feature. The ID is in the range $0\le i_{b,t}<V$ . Repeated IDs select the same learned row; lookup backward adds the contributions from all those occurrences.

Write the existing backbone as $F_\theta$ , where $\theta$ contains its block and final-normalization parameters. Let $\Lambda$ describe active row lengths, causal range starts and RoPE positions. Then

H=F_\theta(X;\Lambda),\qquad X,H\in\mathbb R^{B\times T\times D}.

The backbone returns features of width $D$ . It does not select a vocabulary. The text head supplies that last operation:

Z_{b,t,v}=\sum_{d=0}^{D-1}H_{b,t,d}E_{v,d},\qquad Z=H E^{\mathsf T}.

Vocabulary index $v$ selects one output score. The multiplication uses the transpose of the same $[V,D]$ table, so $Z$ has shape $[B,T,V]$ . The text loss consumes these logits, target IDs and any valid-target mask. Targets and loss eligibility therefore stay outside the backbone.

Add the two derivatives of the shared table

An untied model could use separate lookup and head matrices. Tying constrains their values to one shared parameter. At fixed numerical weights, input and loss, changing that parameter affects the loss through both roles. The forward composition and the chain rule give the complete relation

\begin{aligned} X_{b,t,d}&=E_{i_{b,t},d},\\ H&=F_\theta(X;\Lambda),\\ Z_{b,t,v}&=\sum_{d=0}^{D-1}H_{b,t,d}E_{v,d},\\ \nabla_E\mathcal L&=G_{\mathrm{lookup}}+G_{\mathrm{head}}. \end{aligned}

$\mathcal L$ is the scalar text loss. $G_{\mathrm{lookup}}$ is the table derivative through lookup, the backbone and the head’s hidden input. $G_{\mathrm{head}}$ is the derivative through the head’s direct use of the table. Both have shape $[V,D]$ and accumulate into the same learned node.

The experiment fixes a detached copy of the output table’s values while leaving lookup connected to measure $G_{\mathrm{lookup}}$ . It then fixes detached text embeddings while leaving the head’s table connected to measure $G_{\mathrm{head}}$ . Neither isolation changes the numerical weights or input. The production text path keeps both connections.

For token row $99$ , feature $0$ , lookup contributes approximately $0.02265892536$ , the head contributes approximately $-0.00149137079$ , and the total tied-table derivative is approximately $0.02116755456$ . These displayed values are rounded; the demo prints the full f64 values. Across all $1032$ table coordinates, the maximum measured absolute difference between the total and the sum is zero. The comparison declares absolute tolerance $10^{-12}$ because floating-point additions can differ in order. This example does not determine the relative sizes or signs of the two contributions in other models or inputs.

Fix only the output table or only the text embeddings to isolate the two derivatives, then compare their sum with the full tied gradient. rust/demos/practical-ch10-input-encoding-boundary/src/lib.rs#tied-gradient-paths

    // Fix only the output table to isolate the lookup contribution.
    // detach is an experiment here; the production head uses the live table.
    clear_gradients(&model)?;
    let encoded = TextInputEncoder::tokens(&model, ids, &[1, 3]).encode(context, &backbone)?;
    let input_hidden = backbone.forward(encoded)?;
    let fixed_table = head.table().detach();
    TextOutputHead::new(&fixed_table)
        .forward(context, input_hidden.hidden())?
        .indexed_mean_nll(2, &targets)?
        .backward()?;
    let lookup = model
        .tied_embedding()
        .tensor()
        .gradient_snapshot()
        .ok_or("missing lookup gradient")?;

    // Fix only the text lookup result to isolate the head contribution.
    clear_gradients(&model)?;
    let encoded = TextInputEncoder::tokens(&model, ids, &[1, 3]).encode(context, &backbone)?;
    let fixed_input = PreparedDecoderInput::new(
        &backbone,
        context,
        encoded.embeddings().detach(),
        encoded.layout().clone(),
        "fixed text lookup for derivative partition",
        &[],
    )?;
    let head_hidden = backbone.forward(fixed_input)?;
    head.forward(context, head_hidden.hidden())?
        .indexed_mean_nll(2, &targets)?
        .backward()?;
    let output = model
        .tied_embedding()
        .tensor()
        .gradient_snapshot()
        .ok_or("missing head gradient")?;
    let max_gradient_error = total
        .as_slice()
        .iter()
        .zip(lookup.as_slice())
        .zip(output.as_slice())
        .map(|((&all, &a), &b)| (all - (a + b)).abs())
        .fold(0.0_f64, f64::max);
    // Show the first feature of the existing 'a' token row (ID99).
    let coordinate = 99 * 4;
    let gradient_sample = json!({
        "row":99, "feature":0, "total":total.as_slice()[coordinate],
        "lookup_path":lookup.as_slice()[coordinate], "head_path":output.as_slice()[coordinate],
    });

Keep the source axes separate from the vocabulary axis

$B,T,D$ describe vector sequences: rows, positions and residual features. $V$ describes the text vocabulary. Lookup maps text IDs to $X[B,T,D]$ ; the backbone maps $X$ to $H[B,T,D]$ ; the text head maps $H$ to $Z[B,T,V]$ . A generic vector input does not acquire a universal token vocabulary merely by sharing the residual width $D$ .

The fixture uses $V=258,D=4$ , two blocks, one attention head, feed-forward width $8$ and position capacity $8$ . Its RoPE base is $10000$ and its RMSNorm epsilon is $10^{-6}$ . These are the existing course decoder’s settings for this comparison.

Give the backbone checked vectors

InputEncoder::encode receives the explicit AutogradContext and the destination DecoderBackbone. It returns PreparedDecoderInput or the encoder’s typed error. Source-specific validation belongs in the encoder; the backbone receives the prepared vector sequence.

Require a source encoder to produce checked vectors for a particular backbone under an explicit graph-recording context. rust/crates/llm-from-scratch-practical/src/models/prepared_decoder_input.rs#input-encoder

/// A producer of differentiable vectors. Source-specific validation belongs in
/// the implementation; the transformer receives only the checked result.
pub trait InputEncoder {
    type Error: std::error::Error;
    fn encode<'execution>(
        &self,
        context: AutogradContext<'execution>,
        backbone: &DecoderBackbone<'_>,
    ) -> Result<PreparedDecoderInput<'execution>, Self::Error>;
}

SequenceLayout gives the batch and position shape, each row’s active prefix length, each cell’s causal range start and its explicit RoPE position. An active query may attend from its range start through itself, within that row’s active prefix. Ordinary text uses a range start of zero and positions starting at zero; packed text can retain its occurrence-local ranges and positions. This layout does not contain target IDs or loss masks.

The prepared carrier retains the original TensorValue handle, so recorded lookup operations remain connected to their learned table. Its checks cover residual width, layout, position capacity and graph policy. It also binds the destination’s actual parameter nodes and value revisions, plus every learned producer parameter declared by the encoder. Consumption checks those bindings again and holds read guards while the backbone runs.

Keep the original embedding handle and check vector shape, graph context and exact live parameter bindings when preparing and consuming it. rust/crates/llm-from-scratch-practical/src/models/prepared_decoder_input.rs#prepared-decoder-input

/// A checked [B,T,D] vector sequence. No token IDs, vocabulary, tokenizer,
/// targets or output table are required by this carrier.
///
/// Learned encoders must list every live parameter used to produce the vectors.
/// Handles are retained without copying learned values. A changed parameter
/// invalidates a prepared result before the transformer performs any work.
pub struct PreparedDecoderInput<'execution> {
    embedding: TensorValue,
    embedding_binding: TensorValueBinding,
    context: AutogradContext<'execution>,
    layout: SequenceLayout,
    provenance: String,
    config: DecoderBlockConfig,
    layers: usize,
    identity: Option<CoreIdentity>,
    parameters: Vec<TensorValueBinding>,
    producer_parameters: Vec<TensorValueBinding>,
}
impl<'execution> PreparedDecoderInput<'execution> {
    pub fn new(
        backbone: &DecoderBackbone<'_>,
        context: AutogradContext<'execution>,
        embedding: TensorValue,
        layout: SequenceLayout,
        provenance: impl Into<String>,
        producer_parameters: &[NamedParameter],
    ) -> Result<Self, PreparedInputError> {
        let input = Self {
            embedding_binding: TensorValueBinding::capture(&embedding),
            embedding,
            context,
            layout,
            provenance: provenance.into(),
            config: backbone.config(),
            layers: backbone.blocks().len(),
            identity: backbone.identity().cloned(),
            parameters: backbone
                .parameters()
                .map(|p| TensorValueBinding::capture(p.tensor()))
                .collect(),
            producer_parameters: producer_parameters
                .iter()
                .map(|p| TensorValueBinding::capture(p.tensor()))
                .collect(),
        };
        input.validate(backbone)?;
        Ok(input)
    }
    /// Recheck exact live nodes/revisions at consumption, not just preparation.
    pub fn validate(&self, backbone: &DecoderBackbone<'_>) -> Result<(), PreparedInputError> {
        let [b, t] = self.layout.shape();
        if self.embedding.shape() != [b, t, backbone.config().model_width()]
            || t > backbone.config().max_positions()
            || self.identity.as_ref().is_some_and(|id| b > id.max_rows)
        {
            return Err(PreparedInputError("B,T,D or admitted row count"));
        }
        if self
            .layout
            .positions()
            .iter()
            .any(|&p| p >= backbone.config().max_positions())
        {
            return Err(PreparedInputError("position capacity"));
        }
        if self.config != backbone.config()
            || self.layers != backbone.blocks().len()
            || self.identity.as_ref() != backbone.identity()
        {
            return Err(PreparedInputError(
                "backbone configuration or execution identity",
            ));
        }
        if self.provenance.is_empty() {
            return Err(PreparedInputError("empty producer provenance"));
        }
        if !self.embedding_binding.node_matches(&self.embedding)
            || !self.embedding_binding.revision_matches(&self.embedding)
        {
            return Err(PreparedInputError(
                "prepared hidden tensor identity or revision",
            ));
        }
        if !self.context.records_graph() && self.embedding.tracks_gradient() {
            return Err(PreparedInputError("graph-recording mode"));
        }
        let current: Vec<_> = backbone.parameters().collect();
        if self.parameters.len() != current.len()
            || self.parameters.iter().zip(current).any(|(binding, p)| {
                !binding.node_matches(p.tensor()) || !binding.revision_matches(p.tensor())
            })
        {
            return Err(PreparedInputError("backbone parameter node or revision"));
        }
        if self
            .producer_parameters
            .iter()
            .any(|binding| !binding.revision_matches(binding.value_handle()))
        {
            return Err(PreparedInputError("producer parameter revision"));
        }
        Ok(())
    }
    pub fn embeddings(&self) -> &TensorValue {
        &self.embedding
    }
    pub fn context(&self) -> AutogradContext<'execution> {
        self.context
    }
    pub fn layout(&self) -> &SequenceLayout {
        &self.layout
    }
    /// Descriptive provenance is never used in place of a live node binding.
    pub fn provenance(&self) -> &str {
        &self.provenance
    }
    pub(crate) fn producer_parameters(&self) -> &[TensorValueBinding] {
        &self.producer_parameters
    }
}

These checks prevent a prepared result from silently using a different model or stale learned vectors. Equal-shaped, equal-valued parameters in another model still have different nodes. Updating the lookup table after encoding invalidates the prepared vectors even though that table is outside the backbone. Descriptive provenance identifies the producer; changing a label cannot authorize different vectors or a stale parameter binding.

The real refusal cases include an incorrect residual width, malformed layout or out-of-capacity position, a graph-tracked lookup supplied under no-grad, different or changed backbone parameters, and a changed learned producer parameter. A no-grad call must prepare its text lookup under no-grad from the start; relabeling an already tracked tensor does not detach it.

TextInputEncoder retains text-specific checks, including the existing padded or packed batch binding and validity rules. The generic carrier requires no tokenizer, token IDs, vocabulary, text split identity, targets or tied output table. Only the text encoder is implemented here. A future source still needs compatible encoding and training, together with any additional attention or objective semantics it requires.

Borrow the existing transformer and table

DecoderBackbone borrows the existing blocks and final RMSNorm. It can also be assembled directly from DecoderBlockConfig, blocks and final norm, without constructing a text model or supplying a tokenizer or vocabulary. Its public full forward returns hidden features. Construction does not copy learned tensor values, and full and cached execution retain one shared block loop.

Build a borrowing view from existing decoder components and expose a full forward that returns hidden features. rust/crates/llm-from-scratch-practical/src/models/causal_decoder_core.rs#decoder-backbone

/// A borrowing transformer view. It owns neither a source encoder nor an output
/// table, and constructing it does not allocate or clone learned tensor values.
pub struct DecoderBackbone<'a> {
    config: DecoderBlockConfig,
    blocks: &'a [DecoderBlock],
    final_norm: &'a RmsNorm,
    identity: Option<CoreIdentity>,
}
impl<'a> DecoderBackbone<'a> {
    /// Reuses the component checks used by complete text-model assembly.
    pub fn new(
        config: DecoderBlockConfig,
        blocks: &'a [DecoderBlock],
        final_norm: &'a RmsNorm,
    ) -> Result<Self, DecoderModelError> {
        validate_backbone_config(config)?;
        for (layer, block) in blocks.iter().enumerate() {
            if block.model_width() != config.model_width() {
                return Err(DecoderModelError::BlockModelWidthMismatch {
                    layer,
                    expected: config.model_width(),
                    actual: block.model_width(),
                });
            }
            if block.attention().heads() != config.heads() {
                return Err(DecoderModelError::BlockHeadCountMismatch {
                    layer,
                    expected: config.heads(),
                    actual: block.attention().heads(),
                });
            }
            if block.attention().kv_heads() != config.kv_heads()
                || block.attention().context_window() != config.context_window()
            {
                return Err(DecoderModelError::PreparedInput {
                    reason: "block head sharing or context window differs from backbone",
                });
            }
            if block.feed_forward().hidden_width() != config.feed_forward_width() {
                return Err(DecoderModelError::BlockFeedForwardWidthMismatch {
                    layer,
                    expected: config.feed_forward_width(),
                    actual: block.feed_forward().hidden_width(),
                });
            }
            if block.attention().rope().max_positions() != config.max_positions() {
                return Err(DecoderModelError::BlockPositionCapacityMismatch {
                    layer,
                    expected: config.max_positions(),
                    actual: block.attention().rope().max_positions(),
                });
            }
            if block.attention().rope().base().to_bits() != config.rope_base().to_bits() {
                return Err(DecoderModelError::BlockRopeBaseMismatch {
                    layer,
                    expected: config.rope_base(),
                    actual: block.attention().rope().base(),
                });
            }
            for (norm, epsilon) in [
                (
                    DecoderModelNorm::Attention,
                    block.attention_norm().epsilon(),
                ),
                (
                    DecoderModelNorm::FeedForward,
                    block.feed_forward_norm().epsilon(),
                ),
            ] {
                if epsilon.to_bits() != config.rms_epsilon().to_bits() {
                    return Err(DecoderModelError::BlockRmsEpsilonMismatch {
                        layer,
                        norm,
                        expected: config.rms_epsilon(),
                        actual: epsilon,
                    });
                }
            }
        }
        if final_norm.feature_width() != config.model_width() {
            return Err(DecoderModelError::FinalNormWidthMismatch {
                expected: config.model_width(),
                actual: final_norm.feature_width(),
            });
        }
        if final_norm.epsilon().to_bits() != config.rms_epsilon().to_bits() {
            return Err(DecoderModelError::FinalNormEpsilonMismatch {
                expected: config.rms_epsilon(),
                actual: final_norm.epsilon(),
            });
        }

        Ok(Self {
            config,
            blocks,
            final_norm,
            identity: None,
        })
    }
    pub const fn config(&self) -> DecoderBlockConfig {
        self.config
    }
    pub fn blocks(&self) -> &[DecoderBlock] {
        self.blocks
    }
    pub const fn final_norm(&self) -> &RmsNorm {
        self.final_norm
    }
    pub fn parameters(&self) -> impl Iterator<Item = &NamedParameter> {
        self.blocks
            .iter()
            .flat_map(|block| block.parameters())
            .chain(self.final_norm.parameters())
    }
    pub(crate) fn identity(&self) -> Option<&CoreIdentity> {
        self.identity.as_ref()
    }
    pub(crate) fn with_identity(mut self, identity: Option<CoreIdentity>) -> Self {
        self.identity = identity;
        self
    }
    /// Run the checked sequence and return [B,T,D] hidden features.
    pub fn forward(
        &self,
        input: PreparedDecoderInput,
    ) -> Result<BackboneForward, DecoderModelError> {
        consume_with_dropout(self, input, &mut DropoutExecution::Identity, false)
            .map_err(full_error)
    }

    /// Same transformer and checked carrier, retaining only block boundaries.
    /// With graph recording enabled, the block-intermediate list is empty;
    /// final hidden features remain. No-grad execution needs no checkpoints.
    pub fn forward_recomputed(
        &self,
        input: PreparedDecoderInput,
    ) -> Result<BackboneForward, DecoderModelError> {
        consume_with_dropout(self, input, &mut DropoutExecution::Identity, true).map_err(full_error)
    }
    pub(super) fn forward_with_dropout(
        &self,
        input: PreparedDecoderInput,
        dropout: &mut DropoutExecution<'_>,
    ) -> Result<BackboneForward, DecoderModelError> {
        consume_with_dropout(self, input, dropout, false).map_err(full_error)
    }
    pub(super) fn forward_recomputed_with_dropout(
        &self,
        input: PreparedDecoderInput,
        dropout: &mut DropoutExecution<'_>,
    ) -> Result<BackboneForward, DecoderModelError> {
        consume_with_dropout(self, input, dropout, true).map_err(full_error)
    }
    pub(crate) fn forward_cached(
        &self,
        input: PreparedDecoderInput,
        layers: &[LayerKvCache],
        prefix: usize,
    ) -> Result<BackboneForward, DecoderKvCacheError> {
        consume(
            self,
            &input,
            Mode::Cached(layers),
            Some(prefix),
            &mut DropoutExecution::Identity,
            false,
        )
        .map_err(|e| match e {
            CoreError::Cached(e) => e,
            CoreError::Prepared(reason) => DecoderKvCacheError::PreparedInput { reason },
            CoreError::Full(_) => unreachable!("cached execution"),
        })
    }
}

TextOutputHead borrows the same live table used by lookup. It delegates to the existing transpose and matrix multiplication operations:

Borrow the live lookup table and reuse its transpose for the text vocabulary projection. rust/crates/llm-from-scratch-practical/src/models/text_input.rs#text-output-head

/// The existing tied projection, borrowing its live [V,D] lookup table.
/// A constant detached table is used only by the demo's gradient partition.
pub struct TextOutputHead<'a> {
    table: &'a TensorValue,
}
impl<'a> TextOutputHead<'a> {
    pub const fn new(table: &'a TensorValue) -> Self {
        Self { table }
    }
    pub const fn table(&self) -> &TensorValue {
        self.table
    }
    pub fn forward(
        &self,
        context: AutogradContext,
        hidden: &TensorValue,
    ) -> Result<TensorValue, DecoderModelError> {
        let tied = self
            .table
            .transpose_with_context(context, 0, 1)
            .map_err(|source| DecoderModelError::Autodiff {
                stage: DecoderModelStage::TiedWeightTranspose,
                source,
            })?;
        hidden
            .matmul_with_context(context, &tied)
            .map_err(|source| DecoderModelError::Autodiff {
                stage: DecoderModelStage::TiedVocabularyProjection,
                source,
            })
    }
}

The existing model calls now compose the encoder, backbone and head for ordinary, padded and packed text, optional dropout and replay, and text loss. The text wrapper holds parameter reads through the head, retains the original operation order, and keeps parameter names, order, counts and tying. The fixed before-and-after comparison preserves embedding, hidden and logit bits, loss, every named parameter value and gradient, and the tested cache results.

The existing text cache selects no-grad before encoding its next token, runs the same backbone in cached mode and applies the text head before publishing layer cache tickets. positioned_token chooses an input’s RoPE position; it does not choose cached execution. The generic cached entry remains internal to this text cache transaction.

Prefilling [0,99] and decoding 100 reaches a logical prefix of three tokens. Across two layers, the measured work is six cache appends and twelve attention-score values. The cached result is graph-free, and its newest logits have zero bit mismatches against the corresponding full result in this fixture. Those observations do not promise equality for every sequence or future backend. The existing cache and dropout commit rules retain their own failure boundaries; they do not undo an earlier optimizer update or roll back a whole job.

The current numerical execution remains CPU f64. This separation adds no GPU implementation and demonstrates no full-data training or quality result.

Separate the roles even when the weights are shared

A language model can learn different matrices for input lookup and output vocabulary projection. Tying instead constrains those two roles to one parameter. Press and Wolf’s 2017 paper studies sharing input and output embeddings and the update of the common parameter. Its reported language-model results are not measurements of this course fixture. The two derivative paths in our Rust experiment make that shared update visible.

Section 3.4 of the 2017 Transformer paper describes learned token embeddings and the final linear/softmax transform as separate operations while sharing selected weight matrices. The course’s existing RMSNorm, RoPE and SwiGLU operations are its own decoder architecture, with different normalization and position treatment from that paper.

Separating components does not require separate learned storage: both text operations borrow one live table and add their derivative contributions.

Here an input encoder means the source-to-vector adapter. It is not the separate bidirectional encoder stack in the original encoder-decoder Transformer. Our causal backbone is still the decoder we already built.

Follow the forward sequence and the shared gradient

The forward sequence changes the final axis from residual features to text vocabulary entries only at the head. The two derivative paths then meet at the table used in both text operations.

One text table, two roles and one gradient sum

Text IDs select rows of one shared table to produce vectors for one batch row and three positions, with four residual features at each position. The existing causal backbone preserves those axes in its hidden features. The text head uses the same table's transpose to produce 258 vocabulary scores at each position. The lookup-path and head-path derivatives join at that single table of 1032 learned scalars, counted once within the complete model's 1372 scalars and 20 named tensors.

Forward: text IDs become vectors, hidden features and text logits

1. Text encoder

Input token IDs: [0,99,100]

Look up rows in shared table $E$

$X\;[1,3,4]$
2. Existing causal backbone

Two blocks and final RMSNorm preserve one row, three positions and four residual features per position.

$X\;[1,3,4]$ $H\;[1,3,4]$
3. Tied text head

Project through the same table's transpose $E^{\mathsf T}$

$H E^{\mathsf T}=Z$

One row, three positions and 258 text vocabulary scores per position: $[1,3,258]$

Backward: both table roles contribute

Lookup-path derivative

The loss derivative travels through the hidden features, backbone and original text lookup to the selected table rows.

$G_{\mathrm{lookup}}$

Head-path derivative

The loss derivative reaches the table directly through the text vocabulary projection.

$G_{\mathrm{head}}$

One live shared text table

$E\;[258,4]$

The lookup and head borrow this same allocation: 258 vocabulary rows, four features per row and 1032 learned scalars.

The total table gradient $\nabla_E\mathcal L$

is the sum of both contributions: $G_{\mathrm{lookup}}+G_{\mathrm{head}}$

Count this table once within the complete fixture's 1372 learned scalars in 20 named tensors.

Optional practice: reproduce and explain the composition

Run from the repository root with Docker running. The course launcher prepares the workspace image; its initial build may need network access to provision the toolchain and dependencies. COURSE_CORPUS=false disables corpus acquisition, and this fixed CPU text example needs no corpus or model download. After setup, the command runs locked offline Cargo with container networking disabled.

COURSE_CORPUS=false ./course run cargo run --offline --locked -p practical-ch10-input-encoding-boundary

The no-argument demo prints one JSON object. It reports lookup_and_head_same_node: true, same_text_wrapper_logits: true, loss bits 4016ac70bf03b620 and within_tolerance: true. Its cached result reports newest_logit_bit_mismatches: 0 for this fixture.

From the same directory and Docker setup, run the focused boundary tests and the demo library test:

COURSE_CORPUS=false ./course run cargo test --offline --locked -p llm-from-scratch-practical --lib models::prepared_decoder_input::tests
COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch10-input-encoding-boundary --lib

The first command selects six tests covering component-only construction, composition and gradients, layout and graph policy, stale parameters and provenance. The second selects one test covering the demonstrated text composition, shared gradient and cached result.

Print the fixed CPU text composition and its derivative, cache and refusal observations as one JSON object. rust/demos/practical-ch10-input-encoding-boundary/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!(
        "{}",
        serde_json::to_string(&practical_ch10_input_encoding_boundary::run()?)?
    );
    Ok(())
}

Using the demo output above, explain why detaching only the table passed to the experimental head removes the head contribution while preserving a lookup contribution to the original table.
Trace the demo’s malformed text-derived vector with shape $[1,1,3]$ into a backbone of residual width $4$ . Identify the refusal and explain why the vector cannot enter the first block.
For the worked text input and targets, identify which data the backbone needs and which data stay with text lookup or text loss. Explain why a valid vector shape alone does not establish that the model understands another source.

Checked explanations for derivative isolation, width refusal and text-specific data

The detached head table keeps the same numerical values but has no gradient connection to the original learned node. The original lookup remains connected to that node, so the loss still differentiates through the head’s hidden input, the backbone and the lookup. This experiment isolates the lookup contribution. Fixing the text embeddings instead isolates the head contribution; the production path keeps both.
The layout describes one row and one position, but the vector has three features where the backbone requires four. PreparedDecoderInput rejects it with B,T,D or admitted row count before the backbone runs a block. The vector is not padded, truncated or reinterpreted as a token ID.
The backbone needs checked vectors, their active lengths, causal range starts, RoPE positions, graph context and current live bindings. Text token IDs select lookup rows; tokenizer and batch-source checks stay with text encoding. Next-token targets [99,100,1] and valid-target masks stay with the text loss. Another source needs suitable encoding, training and any additional attention or objective semantics; matching residual width supplies only a tensor compatibility condition.

Carry the composition into numerical execution

We can now name the complete text computation as encoding, backbone and head without changing its learned nodes or tied gradient. Chapter 11 will add GPU numerical execution behind a narrow execution boundary while preserving this composition and CPU inference. Device computation must preserve the meaning of the vectors, layout and model parameters.

Persistence will follow in Chapter 15. It will need to restore the actual component composition, required text assets and shared parameter aliases, including the single table used by lookup and the text head. Neither GPU execution nor that persistence capability is implemented in this chapter.