Skip to main content

Module lm_encoder

Module lm_encoder 

Source
Expand description

LM encoder adapter: projects ESMC-6B hidden states → single representation.

Four standard pre-norm transformer blocks (self-attention + SwiGLU FFN) running at d_model=2560, followed by a LayerNorm and a linear projection to d_single=384.

Weight layout (rooted at lm_encoder):

blocks.{0..3}.attn.layernorm_qkv.0.*  — pre-attn LayerNorm
blocks.{0..3}.attn.layernorm_qkv.1.*  — QKV linear (no bias)
blocks.{0..3}.attn.out_proj.*         — output projection (no bias)
blocks.{0..3}.attn.q_ln.*             — per-head Q LayerNorm (no bias)
blocks.{0..3}.attn.k_ln.*             — per-head K LayerNorm (no bias)
blocks.{0..3}.ffn.0.*                 — pre-FFN LayerNorm
blocks.{0..3}.ffn.1.*                 — gate+up projection (no bias)
blocks.{0..3}.ffn.3.*                 — down projection (no bias)
norm.*                                — final LayerNorm
proj.*                                — linear d_in → d_out (no bias)

Structs§

LMEncoder
4-layer transformer adapter that maps ESMC-6B hidden states ([B, L, d_in=2560]) to the ESMFold2 single representation ([B, L, d_out=384]).