File size: 3,008 Bytes
b3112c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Projection adapter: encoder output → LFM2.5 hidden dim.
This is the load-bearing module that lifts the per-transaction encoder's
output into the pretrained LFM2.5 backbone's continuous-token space. The
architecture mirrors LFM2-VL's `Lfm2VlMultiModalProjector` exactly:
LayerNorm → Linear(d_encoder → hidden) → GELU → Linear(hidden → d_lfm)
Why exactly this shape:
- **LayerNorm at the front.** The encoder is freshly initialized; its
outputs have arbitrary norm. LayerNorm gives the projector a unit-norm
input distribution and stabilizes early training. LFM2-VL's config
defaults to `projector_use_layernorm=True`.
- **2 layers, not 1.** LLaVA-1.5 ([arXiv 2310.03744](https://arxiv.org/abs/2310.03744))
switched from a single Linear (LLaVA-1.0) to a 2-layer MLP and got
materially better benchmarks. LFM2-VL adopted this. The non-linear
bridge handles the encoder-domain ↔ text-pretrained-domain gap that a
single Linear cannot.
- **GELU activation.** LFM2-VL uses GELU (`projector_hidden_act="gelu"`).
The encoder uses SiLU internally (matching LFM2's SwiGLU MLPs), but the
projector sits outside the LFM backbone and follows the VL projector
convention. Cross-modality consistency matters more here than
intra-encoder consistency.
- **hidden = 2 * d_lfm by default.** LFM2-VL uses `projector_hidden_size=2560`
for `text_hidden_size=2048` (ratio 1.25). We use 2x for the 350M (hidden
2048 for d_lfm=1024) — slightly more capacity at minor parameter cost.
Defensible default; revisit if the projector becomes a bottleneck.
Shape contract:
(B, T, d_encoder) → (B, T, d_lfm)
"""
from __future__ import annotations
import torch
import torch.nn as nn
class ProjectionAdapter(nn.Module):
"""LFM2-VL-shaped 2-layer MLP projector with input LayerNorm.
Args:
d_encoder: input feature dim from the per-transaction encoder.
d_lfm: output feature dim — must match the LFM2.5 backbone's hidden
size (1024 for LFM2.5-350M, 2048 for LFM2.5-1.2B).
hidden: intermediate projector hidden size. Defaults to `2 * d_lfm`
following LFM2-VL's pattern of `projector_hidden_size > d_lfm`.
use_layernorm: include LayerNorm at the input (LFM2-VL default: True).
"""
def __init__(
self,
d_encoder: int = 256,
d_lfm: int = 1024,
hidden: int | None = None,
use_layernorm: bool = True,
) -> None:
super().__init__()
if hidden is None:
hidden = 2 * d_lfm
self.input_norm: nn.Module = (
nn.LayerNorm(d_encoder) if use_layernorm else nn.Identity()
)
self.up = nn.Linear(d_encoder, hidden)
self.act = nn.GELU()
self.down = nn.Linear(hidden, d_lfm)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (B, T, d_encoder)
x = self.input_norm(x)
x = self.up(x)
x = self.act(x)
return self.down(x)
# → (B, T, d_lfm)
|