You are given two short ML coding problems from a machine-learning engineer screen. Both are implementation-focused but probe whether you understand the underlying numerics and linear-algebra structure, not just the API calls. Use NumPy-style vectorization throughout; avoid Python loops over individual examples. ### Constraints & Assumptions - Use float64 (or the host dtype) NumPy; assume `numpy as np` is available. - Part 1: a single block can be held in memory, but the concatenation of all blocks cannot. Logits are finite reals; $n \ge 1$ over all blocks. Aim for one pass over the stream and $O(1)$ extra state. - Part 2: distances are exact (no approximate-NN structures required for correctness, though you may mention them for scale). Ties may be broken arbitrarily. The query set $X_{\text{query}}$ is given as a batch; the $m \times n$ distance matrix is assumed to fit (discuss batching if it doesn't). - "No Python loops over examples" forbids per-row Python iteration; vectorized NumPy reductions and one matrix multiply are expected. ### Clarifying Questions to Ask - Part 1: is the entropy needed in nats (natural log) or bits (log base 2)? Are degenerate inputs (single class, or one logit dominating to the point that $p$ is one-hot) in scope? - Part 1 streaming: do blocks arrive with known sizes, and is a single final answer required, or a running estimate after each block? - Part 2: are the labels arbitrary class IDs (so the network selects a *prototype index*, then maps to a label) or already one-hot? - Part 2: how large are $n$, $d$, $m$ — does the $m \times n$ distance matrix fit in memory, or is batching expected? - Part 2: for the "neural network" framing, is an exact equivalence required, or is a soft/temperature-scaled approximation acceptable? ### Part 1 — Numerically stable categorical entropy from logits Given a vector of real-valued logits $z = [z_1, z_2, \ldots, z_n]$, let $p = \mathrm{softmax}(z)$. Compute the categorical (Shannon) entropy $$H(p) = -\sum_i p_i \log p_i.$$ Your implementation must be **numerically stable** (it should not overflow when some $z_i$ are large) and should rely on the **log-sum-exp** trick rather than naively forming `exp(z)` and dividing. Then extend the method to a **streaming** setting: the logits arrive as a sequence of blocks (e.g. chunks of a very long vector), and the full vector never fits in memory at once. Describe and implement an **online** algorithm that produces the same entropy after consuming all blocks, using $O(1)$ state per block (independent of $n$). ```hint Rewrite the objective Substitute $\log p_i = z_i - \log Z$ where $Z = \sum_j e^{z_j}$. The entropy collapses to a closed form in terms of $\log Z$ and the softmax-weighted average of the logits $\mathbb{E}_p[z]$ — so you never have to materialize $p$ or take its log. ``` ```hint Stability Factor out the max: $\log\sum_j e^{z_j} = m + \log\sum_j e^{z_j - m}$ with $m = \max_j z_j$. Every exponent is then $\le 0$. ``` ```hint Streaming state Keep a running maximum $M$ plus two accumulators *relative to* $M$ — one for $\sum e^{z_i - M}$ and one for $\sum z_i\, e^{z_i - M}$ — and those reconstruct $\log Z$ and $\mathbb{E}_p[z]$. When a new block raises the maximum, rescale the old accumulators to the new $M$ — the same trick used to merge two log-sum-exp partials. ``` #### What This Part Should Cover - A reformulation that avoids forming $p$ or $\log p_i$ directly (those are the catastrophic operations as $p_i \to 0$), reducing entropy to a small number of stable scalar quantities. - A correct stabilization scheme (max-subtraction) and a clear statement of *why* it guarantees no overflow. - A genuinely online streaming algorithm: constant per-block state, a single pass, and correct handling of the case where a later block raises the running maximum. - An explicit claim — and justification — that the streaming result equals the batch computation up to floating-point roundoff, not just approximately. ### Part 2 — Vectorized 1-NN and its neural-network interpretation Given training examples `X_train` with shape $(n, d)$, integer/categorical labels `y_train` with shape $(n,)$, and query examples `X_query` with shape $(m, d)$: 1. Implement **1-nearest-neighbor** prediction under **L2** distance, fully vectorized — no Python loop over training or query examples. 2. Express L2 1-NN *inference* as a single **neural-network forward pass**: a linear layer followed by a softmax. Show precisely what the layer's weights and biases are and why `argmax` of the output recovers the L2 nearest neighbor. 3. Explain how the construction must change if the metric is **L1** instead of L2. ```hint Vectorized distances You do not need the full distance — `argmin` over $\lVert q - x_i\rVert_2^2$ is enough. Expand the square: $\lVert q - x_i\rVert_2^2 = \lVert q\rVert^2 + \lVert x_i\rVert^2 - 2\,q\cdot x_i$. The cross term is a single matrix product, the norms are broadcasts. ``` ```hint Which terms survive the argmax When you turn "nearest" into "highest score" $\big(-\lVert q-x_i\rVert_2^2\big)$, one of the three expanded terms is identical for every $i$ and can be dropped. What's left is affine in $q$ — that's your linear layer. ``` ```hint L1 is piecewise-linear, not linear A bare affine layer can't produce $\sum_j |q_j - x_{i,j}|$. What two-piece identity lets a one-hidden-layer ReLU network compute an absolute value exactly? ``` #### What This Part Should Cover - Vectorized squared-L2 distances via the $\lVert q\rVert^2 + \lVert x_i\rVert^2 - 2\,q\cdot x_i$ expansion, with correct broadcasting shapes, one matrix multiply, and an `argmin` — no per-example Python loop. - A precise weight/bias mapping for the linear layer, an explicit argument for which expanded term is constant across prototypes and therefore drops out, and the observation that softmax is monotonic so `argmax` is preserved. - For L1: recognition that an affine layer is fundamentally insufficient (the kink), and an exact one-hidden-layer ReLU construction built from a two-piece absolute-value identity. - Awareness of cost and scale: the $O(mn)$ distance matrix and $O(mnd)$ time, batching when the matrix does not fit, and the fact that softmax is only needed for the "as-a-network" framing — exact NN needs only the argmin/argmax. ### What a Strong Answer Covers These dimensions span both parts. - Consistently vectorized NumPy: reductions, broadcasts, and at most one matrix multiply per stage; no Python iteration over individual examples or coordinates. - Numerical-stability awareness throughout — the log-sum-exp max trick in Part 1 and the catastrophic-cancellation risk in the squared-distance identity for near-coincident points in Part 2. - Clear separation of *what is exact* from *what is a convenience*: the algebraic reductions are exact, while softmax (Part 2) and base-2 vs nats (Part 1) are read-out choices that do not change the argmin/argmax. ### Follow-up Questions - The softmax in the 1-NN framing is monotonic, so it doesn't change the `argmax`. What *is* it buying you — and how would a temperature on the scores turn this into soft (weighted) k-NN? - In the streaming entropy algorithm, can you also maintain a running *variance* of the logits under $p$ with $O(1)$ state? What would that give you? - For the L1 network, how does the hidden-unit count scale with $n$ and $d$, and could you share structure across training points to shrink it? - If $X_{\text{train}}$ is fixed but $X_{\text{query}}$ streams in, how would you precompute to make each query cheaper, and where does numerical error creep into the squared-distance identity for nearly-coincident points?