Skip to content

← Flax step 58 of 100

Medium Primitives

Vision-Language Fusion

Jump to the task

Combine two feature vectors

Suppose an image encoder and a text encoder have already produced one vector each. Their widths may differ: an image vector might have six features while a text vector has four. This exercise combines them into one vector with a chosen width, hidden.

Use a separate Dense layer for each input so both projections have the same shape. Add the projections, apply tanh elementwise, then apply a third Dense layer. The module receives feature vectors, not raw pixels or text. It does not include encoders, tokenization, attention or a training objective, and is not a reproduction of a particular vision-language model.

The calculation

With row-vector inputs i and t, and hidden width H:

v = i @ Wv + bv             Wv: (D_v, H), bv: (H,)
u = t @ Wt + bt             Wt: (D_t, H), bt: (H,)
f = tanh(v + u)             f: (H,)
out = f @ Wo + bo           Wo: (H, H), bo: (H,)

The projections have separate weights even when the input widths are equal. A shared output width makes addition possible; it does not by itself mean the features have learned matching semantic meanings. This exercise initializes the parameters and returns an untrained output.

What tanh changes

Tanh keeps the sign of each summed feature and limits its magnitude. In exact arithmetic its output lies strictly between -1 and 1; floating point can round large inputs to either endpoint. Large magnitudes also produce small derivatives, so tanh can saturate.

A negative feature is not inherently disagreement between an image and text. Feature coordinates acquire meaning through training. Here tanh simply introduces a specified nonlinear operation. Without it, the three Dense layers could be combined into one affine map of the concatenated inputs. With it, an input's effect can depend on the other input through the summed value entering tanh.

Only the intermediate vector f is bounded by tanh. The final Dense layer can produce values outside [-1, 1]. It does not apply another tanh.

A small example

Use chosen weights to make the steps visible:

import jax.numpy as jnp

image = jnp.array([1., 2.])
text = jnp.array([3.])
v = image @ jnp.eye(2)  # [1., 2.], zero bias
u = text @ jnp.array([[1., -1.]])  # [3., -3.], zero bias
fused = jnp.tanh(v + u)  # tanh([4., -1.])
out = fused @ (2 * jnp.eye(2)) + jnp.array([0.5, -0.5])
# approximately [2.498659, -2.023188]

The example demonstrates the operation and its unbounded final output. The graded function instead uses the default Dense initializers and the provided seed. With default zero biases, two all-zero inputs yield an all-zero output at initialization. That property need not hold after training.

Problem

Implement vision_language_fusion_forward(seed, image_features, text_features, hidden) using Flax Linen. Inputs are floating-point vectors (D_v,) and (D_t,); both widths and hidden are positive. Return a one-dimensional array of length hidden. There is no batch axis or classification softmax.

Define a compact VLFusion(nn.Module) with a hidden field. Call three separate unnamed nn.Dense(hidden) instances in this order:

  1. Dense_0 projects image_features.
  2. Dense_1 projects text_features.
  3. Dense_2 projects jnp.tanh(v + u).

Keep default kernels and biases and this module nesting. Initialize with jax.random.PRNGKey(int(seed)), passing both input vectors to model.init. Apply the resulting variables to the same two inputs and flatten the result. Treat seed and hidden as fixed Python integers when compiling over inputs.

The parameter count, including all three biases, is D_v * H + D_t * H + H * H + 3 * H. The two input projections are independent in the parameter tree, although their biases enter the sum as bv + bt.